模型检查点

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]

  1. 将提供的模型放到 CPU 上

  2. 将 ZeRO 2 或 3 检查点转换为单个 fp32 合并的 state_dict

  3. 将其加载到提供的模型中

参数

  • 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 示例。