内存需求

用于估计内存使用量的 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 的 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)[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 的 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（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 方差
fp16 参数 +2b

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

固定内存

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

ZeRO-2：无法控制
ZeRO-3

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

现在有 2 种子情况

"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

这需要完成。