内存要求
内存使用估算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
现在有两种子情况
"cpu_offload_params": true
:6 * params (2 字节用于 fp16 参数 + 4 字节用于 fp32 梯度)
如果
gradient_accumulation_steps > 1
,则额外 2 字节的 fp16 梯度被锁页
"cpu_offload_params": false
:4 字节用于 fp32 梯度
激活内存
XXX:对于 Transformer 模型,大约是 (2* seq * attn_heads + 16 * hidden_size) * sequence * batch/gpu
这部分需要补充。