内存要求

内存使用估算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)[source]

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

如果您有实际的模型对象，请使用此函数，所有信息将自动导出。

如果这是一个假设模型，请使用estimate_zero2_model_states_mem_needs_all_cold，您必须明确传递total_params。

参数

  • model (-) – 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)[source]

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

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

如果您有实际的模型对象，请使用estimate_zero2_model_states_mem_needs_all_live，所有信息将自动导出。

参数

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

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

  • num_nodes (-) – 有多少个节点 (默认为 1)，

  • additional_buffer_factor (-) – 估算因子 (默认为 1.5)

示例

让我们尝试一个只有 1 个节点、8 个 GPU 的 30亿参数模型，使用实际模型

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)[source]

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

如果您有实际的模型对象，请使用此函数，所有信息将自动导出。

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

参数

  • model (-) – 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)[source]

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

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

如果您有实际的模型对象，请使用estimate_zero3_model_states_mem_needs_all_live，所有信息将自动导出。

参数

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

  • largest_layer_params (-) – 最大层参数

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

  • num_nodes (-) – 有多少个节点 (默认为 1)，

  • additional_buffer_factor (-) – 估算因子 (默认为 1.5)

示例

让我们尝试一个只有 1 个节点、8 个 GPU 的 30亿参数模型，使用实际模型

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 的 pytorch_model.bin大小为 42.1GB，所以如果我们将其除以 4，我们可以立即判断这是一个 110亿参数的模型。

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

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

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

GPU 内存

主要问题是您现有的硬件能容纳多大的模型？或者说，您需要多大的 GPU 内存来容纳目标模型。

  • ZeRO-2

    • "offload_optimizer": {"device": "cpu"}: 2 * params

    示例：一个 40GB 的 GPU 可以容纳大约 110亿参数的模型（无论使用多少个 GPU）。这里模型以fp16加载，因此模型权重本身大约占用 22GB，剩余的 18GB 用于其他组件。在这种情况下，您几乎无法适应非常小的批量大小。

    • "offload_optimizer": {"device": "none"}: 4 * params + 16 * params/ (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"} - largest_layer_memory + 18 * params / 所有节点上的 GPU 总数

  • 情况 2: "offload_param": {"device": "cpu"}, "offload_optimizer": {"device": "cpu"}- largest_layer_memory。这里的主要限制是通用内存 (RAM)。

  • 情况 3: "offload_param": {"device": "none"}, "offload_optimizer": {"device": "cpu"}- largest_layer_memory + 2 * params / 所有节点上的 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 内存所需的成本。而且，通常甚至不可能购买具有大量内存的 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"}:

      params * max(4 * n_gpus, 16) * additional_buffer_factor

    示例: xxx

  • ZeRO-3

    gpus_factor = n_gpus / total_gpus

    • 情况 1: "offload_param": {"device": "none"}, "offload_optimizer": {"device": "none"}

      不使用zero.Init

      params * 4 * n_gpus * additional_buffer_factor

      这是仅在开始时在 CPU 内存中初始化模型所需的内存。一旦模型传输到 GPU，此内存即被释放。

      使用zero.Init

      largest_layer_params * 4 * n_gpus * additional_buffer_factor

      假设一旦张量通过 ZeRO.Init 移动到 GPU，Pytorch 就会释放内存

    • 情况 2: "offload_param": {"device": "cpu"}, "offload_optimizer": {"device": "cpu"}

      不使用zero.Init

      params * max(4 * n_gpus, 18 * gpus_factor) * additional_buffer_factor

      使用zero.Init

      params * 18 * gpus_factor * additional_buffer_factor

    • 情况 3: "offload_param": {"device": "none"}, "offload_optimizer": {"device": "cpu"}

      不使用zero.Init

      params * max(4 * n_gpus, 16 * gpus_factor) * additional_buffer_factor

      使用zero.Init

      params * 16 * gpus_factor * additional_buffer_factor

以下是 16 和 18 乘数的细分（b = 字节）

4 （在4*n_gpus中）

  • 当 pytorch 创建模型时，默认情况下会以 fp32 格式创建（4 字节）

16:

  • fp32 为 16 字节：每个参数 4 字节参数、4 字节梯度、4 字节动量和 4 字节方差

18:

  • fp32 为 16 字节：每个参数 4 字节参数、4 字节梯度、4 字节动量和 4 字节方差

  • +2 字节用于 fp16 参数

关于梯度的注意事项：虽然梯度以 fp16（2 字节）存储，但在权重更新期间，所有梯度在进行权重更新之前都会转换为 fp32，因为在 DeepSpeed 的 FusedAdam 优化器中，权重更新几乎是在整个模型粒度（param_group 粒度）上完成的。因此，在转换之后，我们几乎需要为所有权重提供每个梯度 4 字节的存储空间。

锁页内存

锁页通用内存 (RAM) 包含在正常的通用内存 (RAM) 分配中（即，这并非额外的内存分配，而是简单地显示有多少通用内存 (RAM) 被锁页）。

  • ZeRO-2: 无法控制

  • ZeRO-3

要启用，请添加: "cpu_offload_use_pin_memory" : true

现在有两种子情况

  1. "cpu_offload_params": true:

    • 6 * params (2 字节用于 fp16 参数 + 4 字节用于 fp32 梯度)

    • 如果gradient_accumulation_steps > 1，则额外 2 字节的 fp16 梯度被锁页

  2. "cpu_offload_params": false:

    • 4 字节用于 fp32 梯度

激活内存

XXX：对于 Transformer 模型，大约是 (2* seq * attn_heads + 16 * hidden_size) * sequence * batch/gpu

这部分需要补充。