内存需求

API 用于估算内存使用量

ZeRO2:

deepspeed.runtime.zero.stage_1_and_2.estimate_zero2_model_states_mem_needs_all_live(model, num_gpus_per_node=1, num_nodes=1, additional_buffer_factor=1.5)

打印出给定model和硬件设置的ZeRO 2参数、优化状态和梯度的内存使用需求估计。

如果你有一个实际的模型对象，使用这个函数，所有内容都会自动派生。

如果这是一个假设模型，请使用estimate_zero2_model_states_mem_needs_all_cold，其中您必须显式传递total_params。

Parameters

• 模型 (-) – nn.Module 对象

• num_gpus_per_node (-) – 每个节点有多少个GPU（默认为1）

• num_nodes (-) – 节点数量（默认为1），

• additional_buffer_factor (-) – 估计因子（默认为1.5）：

deepspeed.runtime.zero.stage_1_and_2.estimate_zero2_model_states_mem_needs_all_cold(total_params, num_gpus_per_node=1, num_nodes=1, additional_buffer_factor=1.5)

打印出给定model和硬件设置下ZeRO 2参数、优化状态和梯度的内存使用需求估计。

如果这是一个假设模型，请使用此函数，您必须显式传递total_params和largest_layer_params。

如果你有一个实际的模型对象，使用 estimate_zero2_model_states_mem_needs_all_live，所有内容将自动派生。

Parameters

• total_params (-) – 模型参数总数

• num_gpus_per_node (-) – 每个节点有多少个GPU（默认为1）

• num_nodes (-) – 节点数量（默认为1），

• additional_buffer_factor (-) – 估计因子（默认为1.5）：

示例：

让我们尝试使用一个只有1个节点和8个GPU的3B模型，使用实时模型：

python -c 'from transformers import AutoModel; \
from deepspeed.runtime.zero.stage_1_and_2 import estimate_zero2_model_states_mem_needs_all_live; \
model = AutoModel.from_pretrained("t5-3b"); \
estimate_zero2_model_states_mem_needs_all_live(model, num_gpus_per_node=8, num_nodes=1)'

Estimated memory needed for params, optim states and gradients for a:
HW: Setup with 1 node, 8 GPUs per node.
SW: Model with 2851M total params.
  per CPU  |  per GPU |   Options
  127.48GB |   5.31GB | offload_optimizer=cpu
  127.48GB |  15.93GB | offload_optimizer=none

现在，没有实际的模型，这需要我们了解total_params和 largest_layer_params，但我们从上面的运行中得到了这些，所以未来的估计器现在更快了，因为我们不需要加载模型。

python -c 'from deepspeed.runtime.zero.stage_1_and_2 import estimate_zero2_model_states_mem_needs_all_cold; \
estimate_zero2_model_states_mem_needs_all_cold(total_params=2851e6, num_gpus_per_node=8, num_nodes=1)'

Estimated memory needed for params, optim states and gradients for a:
HW: Setup with 1 node, 8 GPUs per node.
SW: Model with 2851M total params.
  per CPU  |  per GPU |   Options
  127.45GB |   5.31GB | offload_optimizer=cpu
  127.45GB |  15.93GB | offload_optimizer=none

由于四舍五入的原因，存在轻微的差异 - 实际运行的模型有更多的参数

ZeRO3:

deepspeed.runtime.zero.stage3.estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=1, num_nodes=1, additional_buffer_factor=1.5)

打印出给定model和硬件设置下ZeRO 3参数、优化状态和梯度的内存使用需求估计。

如果你有一个实际的模型对象，使用这个函数，所有内容都会自动派生。

如果这是一个假设模型，请使用estimate_zero3_model_states_mem_needs_all_cold，其中您必须显式传递total_params和largest_layer_params。

Parameters

• 模型 (-) – nn.Module 对象

• num_gpus_per_node (-) – 每个节点有多少个GPU（默认为1）

• num_nodes (-) – 节点数量（默认为1），

• additional_buffer_factor (-) – 估计因子（默认为1.5）：

deepspeed.runtime.zero.stage3.estimate_zero3_model_states_mem_needs_all_cold(total_params, largest_layer_params, num_gpus_per_node=1, num_nodes=1, additional_buffer_factor=1.5)

打印出给定model和硬件设置下ZeRO 3参数、优化状态和梯度的内存使用需求估计。

如果这是一个假设模型，请使用此函数，您必须明确传递total_params和largest_layer_params。

如果你有一个实际的模型对象，使用 estimate_zero3_model_states_mem_needs_all_live，一切都会自动推导出来。

Parameters

• total_params (-) – 模型参数总数

• largest_layer_params (-) – 最大层的参数

• num_gpus_per_node (-) – 每个节点有多少个GPU（默认为1）

• num_nodes (-) – 节点数量（默认为1），

• additional_buffer_factor (-) – 估计因子（默认为1.5）：

示例：

让我们尝试使用一个只有1个节点和8个GPU的3B模型，使用实时模型：

python -c 'from transformers import AutoModel; \
from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live; \
model = AutoModel.from_pretrained("t5-3b"); \
estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=8, num_nodes=1)'

Estimated memory needed for params, optim states and gradients for a:
HW: Setup with 1 node, 8 GPUs per node.
SW: Model with 2851M total params, 32M largest layer params.
  per CPU  |  per GPU |   Options
   71.71GB |   0.12GB | offload_param=cpu , offload_optimizer=cpu , zero_init=1
  127.48GB |   0.12GB | offload_param=cpu , offload_optimizer=cpu , zero_init=0
   63.74GB |   0.79GB | offload_param=none, offload_optimizer=cpu , zero_init=1
  127.48GB |   0.79GB | offload_param=none, offload_optimizer=cpu , zero_init=0
    1.47GB |   6.10GB | offload_param=none, offload_optimizer=none, zero_init=1
  127.48GB |   6.10GB | offload_param=none, offload_optimizer=none, zero_init=0

现在，没有实际的模型，这需要我们了解total_params和 largest_layer_params，但我们从上面的运行中得到了这些，所以未来的估计器现在更快了，因为我们不需要加载模型。

python -c 'from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_cold; \
estimate_zero3_model_states_mem_needs_all_cold(total_params=2851e6, largest_layer_params=32e6, num_gpus_per_node=8, num_nodes=1)'

Estimated memory needed for params, optim states and gradients for a:
HW: Setup with 1 node, 8 GPUs per node.
SW: Model with 2851M total params, 32M largest layer params.
  per CPU  |  per GPU |   Options
   71.69GB |   0.12GB | offload_param=cpu , offload_optimizer=cpu , zero_init=1
  127.45GB |   0.12GB | offload_param=cpu , offload_optimizer=cpu , zero_init=0
   63.72GB |   0.78GB | offload_param=none, offload_optimizer=cpu , zero_init=1
  127.45GB |   0.78GB | offload_param=none, offload_optimizer=cpu , zero_init=0
    1.43GB |   6.09GB | offload_param=none, offload_optimizer=none, zero_init=1
  127.45GB |   6.09GB | offload_param=none, offload_optimizer=none, zero_init=0

由于四舍五入，存在轻微差异 - 实际实时模型有更多参数

讨论

让我们详细看看内存估算器API是如何计算这些数字的，并讨论一些API未涵盖的额外数字。

在以下讨论中：

• params - 模型参数的总数，可以通过以下公式计算：

print(sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values()))

一些模型已经在模型名称中包含了参数数量，例如 t5-11b（110亿参数），gpt-neo-1.3B（13亿参数）等。

此外，如果模型权重存储在fp32中，另一种快速计算模型大小的方法是将state_dict文件的大小除以4（fp32 == 4字节）。例如，你可以看到t5-11b’s pytorch_model.bin的大小是42.1GB，所以如果我们将其除以4，我们可以立即知道它是一个11B模型。

以下计算显示了模型参数、梯度和优化器状态所需的内存。除此之外，您还需要足够的内存来适应激活计算和任何用于中间计算的临时内存，这对于长序列来说可能非常重要（例如，可能需要与参数+梯度+优化器状态总和相同的内存）。

优化器状态假设使用了Adam，其中每个参数使用4字节用于动量，另外4字节用于方差（总共8字节）。

在fp32下的梯度占用4字节，参数在fp16下占用2字节，在fp32下占用4字节。

GPU 内存

大问题是你能在现有的硬件上安装多大的模型？或者更确切地说，你需要多大的GPU内存来安装所需的模型。

• ZeRO-2:

  • "offload_optimizer": {"device": "cpu"}: 2 * 参数

  示例：一个40GB的GPU可以容纳约110亿参数的模型（无论使用多少个GPU）。这里模型以fp16格式加载，因此仅模型权重就占用了约22GB，剩余的18GB被其他组件使用。在这种情况下，你几乎只能容纳一个非常小的批量大小。

  • "offload_optimizer": {"device": "none"}: 4 * 参数 + 16 * 参数/ (GPU总数)

• ZeRO-3:

largest_layer_memory = 4*largest_layer_params - 在单个GPU上收集最大层所需的GPU内存。收集2字节的fp16参数并计算2字节的fp16梯度（总共4倍）。优化器状态和fp32参数以分区形式更新，并以分区形式复制到fp16参数中。这发生在优化器步骤期间。之后，fp16参数就足够了。

• 案例1: "offload_param": {"device": "none"}, "offload_optimizer": {"device": "none"} - 最大层内存 + 18 * 参数 / 所有节点上的GPU总数

• 案例2: "offload_param": {"device": "cpu"}, "offload_optimizer": {"device": "cpu"}- 最大的层内存。这里的主要限制是通用RAM。

• 案例3: "offload_param": {"device": "none"}, "offload_optimizer": {"device": "cpu"}- 最大层内存 + 2 * 参数 / 所有节点上的GPU总数

  示例：

from transformers import AutoModel
model = AutoModel.from_pretrained("t5-large")

# shared params calculated only ones
total_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values())

largest_layer_params = 0
for m in model.modules():
    # assuming no shared params within a single layer
    layer_params = sum(p.numel() for p in m.parameters(recurse=False))
    largest_layer_params = max(largest_layer_params, layer_params)

largest_layer_memory = (4*largest_layer_params)

total_gpus = 4

case1 = largest_layer_memory + int(18*total_params/total_gpus)
case2 = largest_layer_memory
case3 = largest_layer_memory + int(2*total_params/total_gpus)

print(f"total params:         {total_params/1e6:6.2f}M")
print(f"largest layer params: {largest_layer_params/1e6:6.2f}M")
print(f"largest layer memory: {largest_layer_memory>>20:6}MB")
print(f"case1 gpu memory: {(case1)>>20:6}MB")
print(f"case2 gpu memory: {(case2)>>20:6}MB")
print(f"case3 gpu memory: {(case3)>>20:6}MB")

total params:         737.67M
largest layer params:  32.90M
largest layer memory:    125MB
case1 gpu memory:   3291MB
case2 gpu memory:    125MB
case3 gpu memory:    477MB

通用RAM:

ZeRO 的关键特性之一是其 CPU 卸载功能，它可以通过使用普通 RAM 显著扩展项目可访问的总内存池。用户可以轻松地将普通 RAM 扩展 10 倍，成本远低于获得相同 GPU RAM 的成本。而且，通常甚至无法购买具有大量 RAM 的 GPU（112GB GPU 有人要吗？），因为它们根本还不存在。

在以下计算中，我们将使用：

• additional_buffer_factor=1.5 作为一个额外的缓冲因子以保持保守

• n_gpus 单个节点（机器）上的GPU数量

• total_gpus 所有节点上的GPU总数

• params - 模型参数的总数（请参阅上文了解如何获取此数字）

• ZeRO-2:

  • "offload_optimizer": {"device": "none"}:

    params * 4 * n_gpus * additional_buffer_factor - 这是在开始时仅在CPU内存上初始化模型所需的内存

  • "offload_optimizer": {"device": "cpu"}:

    参数 * 最大值(4 * n_gpus, 16) * 额外缓冲因子

  示例：xxx

• ZeRO-3:

  gpus_factor = n_gpus / total_gpus

  • 案例1: "offload_param": {"device": "none"}, "offload_optimizer": {"device": "none"}:

    没有 zero.Init:

    参数 * 4 * GPU数量 * 额外缓冲因子

    这是仅在开始时在CPU内存上初始化模型所需的内存。一旦模型转移到GPU上，这部分内存就会被释放。

    使用 zero.Init：

    最大层参数 * 4 * GPU数量 * 额外缓冲因子

    假设Pytorch在张量通过ZeRO.Init移动到GPU后释放内存。

  • 案例2: "offload_param": {"device": "cpu"}, "offload_optimizer": {"device": "cpu"}:

    没有 zero.Init:

    参数 * 最大值(4 * n_gpus, 18 * gpus_factor) * 额外缓冲因子

    使用 zero.Init：

    参数 * 18 * GPU因子 * 额外缓冲因子

  • 案例 3: "offload_param": {"device": "none"}, "offload_optimizer": {"device": "cpu"}:

    没有 zero.Init:

    参数 * 最大值(4 * n_gpus, 16 * gpus_factor) * 额外缓冲因子

    使用 zero.Init：

    参数 * 16 * GPU因子 * 额外缓冲因子

以下是16和18乘数的详细说明（b = 字节）：

4 (在 4*n_gpus 中):

• 当 PyTorch 创建模型时，默认情况下会以 fp32（4 字节）创建它。

16:

• 16b 用于 fp32：每个参数 4b 参数，4b 梯度，4b 动量和 4b 方差

18:

• 16b 用于 fp32：每个参数 4b 参数，4b 梯度，4b 动量和 4b 方差

• +2b 用于 fp16 参数

关于梯度的说明：虽然梯度以fp16（2字节）存储，但在权重更新期间，所有梯度在权重更新之前都会转换为fp32，因为在DeepSpeed的FusedAdam优化器中，权重更新几乎是在整个模型粒度（参数组粒度）上进行的。因此，在转换后，我们几乎需要为整个权重集的每个梯度分配4字节。

固定内存

固定的通用RAM包含在正常的通用RAM分配中（即这不是额外的内存分配，而只是显示了有多少通用RAM被固定）

• ZeRO-2: 无法控制

• ZeRO-3

要启用添加："cpu_offload_use_pin_memory" : true

现在有2个子情况：

1. "cpu_offload_params": true:

   • 6 * 参数（2b 用于 fp16 参数 + 4b 用于 fp32 梯度）

   • 如果 gradient_accumulation_steps > 1，则会额外固定 2b 用于 fp16 梯度

2. "cpu_offload_params": false:

   • 4b 用于 fp32 梯度

激活记忆

XXX: 对于Transformers来说，大约为 (2* seq * attn_heads + 16 * hidden_size) * sequence * batch/gpu

这需要完成。