模型检查点
DeepSpeed 提供了在训练期间对模型状态进行检查点的例程。
加载训练检查点
保存训练检查点
ZeRO 检查点 fp32 权重恢复
DeepSpeed 提供了从保存的 ZeRO 检查点优化器状态中提取 fp32 权重的例程。
- deepspeed.utils.zero_to_fp32.get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag=None, exclude_frozen_parameters=False, lazy_mode=False)[source]
将 ZeRO 2 或 3 检查点转换为单个 fp32 合并的
state_dict
,该 `state_dict` 可以通过load_state_dict()
加载,并用于不使用 DeepSpeed 进行训练,或者通过模型中心等方式与他人共享。- 参数
checkpoint_dir (-) – 目标检查点文件夹的路径
tag (-) – 用作检查点唯一标识符的检查点标签。如果未提供,将尝试从“latest”文件中加载标签。例如,
global_step14
exclude_frozen_parameters (-) – 排除冻结参数
lazy_mode (-) – 以惰性模式获取 `state_dict`。它返回一个伪张量(pesduo tensor)字典,而不是 PyTorch 张量,这样更节省内存。通过
.contiguous()
将伪张量转换为 PyTorch 张量。
- 返回值
PyTorch
state_dict
典型用法可能如下
from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint # do the training and checkpoint saving state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu model = model.cpu() # move to cpu model.load_state_dict(state_dict) # submit to model hub or save the model to share with others
在此示例中,
model
将无法在同一应用程序的 DeepSpeed 上下文中使用。即,您需要重新初始化 DeepSpeed 引擎,因为model.load_state_dict(state_dict)
将移除其所有的 DeepSpeed 特性。如果您希望自动完成所有操作,请改用
load_state_dict_from_zero_checkpoint
。注意:如果您的应用程序没有足够的空闲 CPU 内存,上述用法可能无法奏效。您可能需要使用离线方法,即使用随检查点保存的
zero_to_fp32.py
脚本。或者您可以以惰性模式加载 `state_dict`。from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu for name, lazy_tensor in state_dict.item(): tensor = lazy_tensor.contiguous() # to cpu print(name, tensor) # del tensor to release memory if it no longer in use
- deepspeed.utils.zero_to_fp32.load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None)[source]
将提供的模型放到 CPU 上
将 ZeRO 2 或 3 检查点转换为单个 fp32 合并的
state_dict
将其加载到提供的模型中
- 参数
model (-) – 要更新的模型对象
checkpoint_dir (-) – 目标检查点文件夹的路径。(包含标签文件夹的那个,例如
global_step14
)tag (-) – 用作检查点唯一标识符的检查点标签。如果未提供,将尝试从检查点文件夹中名为
latest
的文件中加载标签,例如global_step14
- 返回值
修改后的模型
- 返回类型
``model`
在调用此函数之前,请确保有足够的 CPU 内存可用。如果内存不足,请使用
zero_to_fp32.py
工具进行转换。您会发现它已方便地放置在检查点文件夹中。典型用法可能如下
from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir) # submit to model hub or save the model to share with others
请注意,一旦运行此操作,
model
将无法在同一应用程序的 DeepSpeed 上下文中使用。即,您需要重新初始化 DeepSpeed 引擎,因为model.load_state_dict(state_dict)
将移除其所有的 DeepSpeed 特性。
- deepspeed.utils.zero_to_fp32.convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, output_dir, max_shard_size='5GB', safe_serialization=False, tag=None, exclude_frozen_parameters=False)[source]
将 ZeRO 2 或 3 检查点转换为单个 fp32 合并的
state_dict
文件,该文件可以通过torch.load(file)
+load_state_dict()
加载,并用于不使用 DeepSpeed 进行训练。- 参数
checkpoint_dir (-) – 目标检查点文件夹的路径。(包含标签文件夹的那个,例如
global_step14
)output_dir (-) – PyTorch fp32 `state_dict` 输出文件的目录
max_shard_size (-) – 检查点在分片之前的最大大小,默认值为 5GB
safe_serialization (-) – 是否使用 safetensors 或传统的 PyTorch 方式(使用 pickle)保存模型。
tag (-) – 用作检查点唯一标识符的检查点标签。如果未提供,将尝试从检查点文件夹中名为
latest
的文件中加载标签,例如global_step14
exclude_frozen_parameters (-) – 排除冻结参数
避免 ZeRO 检查点膨胀
使用 torch.save()
创建的 ZeRO 阶段 1 和 2 检查点有时会比预期更大。这种膨胀是由于 ZeRO 的张量展平与 PyTorch 的张量 存储管理 之间的相互作用造成的。您可以通过使用 DeepSpeed 的 clone_tensors_for_torch_save
工具来避免此问题,如下所示。
- deepspeed.checkpoint.utils.clone_tensors_for_torch_save(item, device=device(type='cpu'))[source]
返回
item
的副本,其中所有包含的张量都替换为指定设备上的克隆。适用于单个张量,以及包含/嵌套在列表、元组和字典中的张量。- 参数
item (-) – 要克隆的张量或(可能嵌套的)要克隆的张量容器。
device (-) – 目标设备(默认为“cpu”)
- 返回值
item
的副本,其中包含目标设备上的克隆张量
以下代码片段演示了创建 HuggingFace 模型检查点的此功能
ds_config = {
...
}
model = AutoModelForCausalLM.from_pretrained("facebook/opt-13b", torch_dtype=torch.float16)
ds_engine, _, _, _ = deepspeed.initialize(model=model, config_params=ds_config)
lean_state_dict = deepspeed.checkpoint.utils.clone_tensors_for_torch_save(ds_engine.module.state_dict())
ds_engine.module.save_pretrained("lean_after", state_dict=lean_state_dict)
通用检查点(开发中)
ZeRO 数据并行 (DP)、张量并行 (TP)、流水线并行 (PP) 等并行技术会将模型和/或优化器状态进行分片,这使得难以使用在不同数量 GPU 上创建的检查点恢复训练。DeepSpeed 提供了通用检查点机制来解决此问题。通用检查点为用户提供了在使用 3D(TP、PP 和 DP)并行训练时更改 GPU 数量的灵活性,并能更有效地利用弹性训练硬件。开始使用通用检查点的最简单方法是查阅 Megatron-DeepSpeed 和 BLOOM 示例。