内存需求

用于估算内存用量的 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）

示例

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

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）

示例

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

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（11B 参数）、gpt-neo-1.3B（1.3B 参数）等。

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

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

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

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

GPU RAM

最大的问题是，您可以在拥有的硬件上容纳多大的模型？或者更确切地说，您需要多大的 GPU RAM 来容纳所需的模型。

ZeRO-2
- "offload_optimizer": {"device": "cpu"}: 2 * params
示例：40GB GPU 可以容纳约 11B 参数模型（无论使用多少 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 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"}:
  
  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
  
  假设 Pytorch 在 ZeRO.Init 将张量移动到 GPU 后会释放内存。
- 情况 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 占用 16b：每个参数占用 4b 参数、4b 梯度、4b 动量和 4b 方差。

18:

fp32 占用 16b：每个参数占用 4b 参数、4b 梯度、4b 动量和 4b 方差。
+2b 用于 fp16 参数。

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

固定内存

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

ZeRO-2：无法控制。
ZeRO-3

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

现在有两个子情况

"cpu_offload_params": true:
- 6 * params（fp16 参数占用 2b + fp32 梯度占用 4b）
- 如果 gradient_accumulation_steps > 1，则会额外固定 2b 用于 fp16 梯度。
"cpu_offload_params": false:
- fp32 梯度占用 4b。

激活内存

XXX：对于 Transformer，可能约为 (2* seq * attn_heads + 16 * hidden_size) * sequence * batch/gpu

这需要完善。