模型检查点

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，可以使用 load_state_dict() 加载并用于无需 DeepSpeed 的训练，或者与其他人共享，例如通过模型中心。

参数

  • checkpoint_dir (-) – 目标检查点文件夹的路径

  • tag (-) – 用作检查点唯一标识符的检查点标签。如果未提供，将尝试加载“latest”文件中的标签。例如，global_step14

  • exclude_frozen_parameters (-) – 排除冻结的参数

  • lazy_mode (-) – 以惰性模式获取 state_dict。它返回一个伪张量的字典而不是 torch 张量，这更节省内存。通过 .contiguous() 将伪张量转换为 torch 张量

返回值

  • 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 的张量扁平化和 torch 的张量 存储管理 之间的交互引起的。您可以通过使用 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)、流水线并行 (TP) 等并行技术，会对模型和/或优化器状态进行分片，这使得难以使用在不同数量的 GPU 上创建的检查点恢复训练。DeepSpeed 提供了通用检查点机制来解决此问题。通用检查点使用户能够在使用 3D（TP、PP 和 DP）并行训练时更改 GPU 数量，并能够更有效地利用弹性训练硬件。开始使用通用检查点的最简单方法是参考 Megatron-DeepSpeed 和 BLOOM 示例。