模型检查点

DeepSpeed 提供在训练期间检查模型状态的例程。

加载训练检查点

deepspeed.DeepSpeedEngine.load_checkpoint(self, load_dir, tag=None, load_module_strict=True, load_optimizer_states=True, load_lr_scheduler_states=True, load_module_only=False, custom_load_fn=None)

加载训练检查点

参数

  • load_dir – 必需。从该目录加载检查点。

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

  • load_module_strict – 可选。布尔值，用于严格强制执行模块和检查点状态字典中的键匹配。

  • load_optimizer_states – 可选。布尔值，用于从检查点加载训练优化器状态。例如 ADAM 的动量和方差。

  • load_lr_scheduler_states – 可选。布尔值，用于从检查点加载学习率调度器状态。

  • load_module_only – 可选。布尔值，用于仅从检查点加载模型权重。例如 预热。

  • custom_load_fn – 可选。自定义模型加载函数。

返回值

一个包含 load_path 和 client_state 的元组。*load_path: 加载的检查点的路径。如果加载检查点失败，则为 None。*client_state: 用于在客户端代码中加载所需训练状态的字典。

重要提示：在 ZeRO3 下，不能在 engine.save_checkpoint() 之后立即使用 engine.load_checkpoint() 加载检查点。因为 engine.module 已分区，而 load_checkpoint() 需要一个原始模型。如果坚持这样做，请在 load_checkpoint() 之前重新初始化引擎。

保存训练检查点

deepspeed.DeepSpeedEngine.save_checkpoint(self, save_dir, tag=None, client_state={}, save_latest=True, exclude_frozen_parameters=False)

保存训练检查点

参数

  • save_dir – 必需。保存检查点的目录。

  • tag – 可选。检查点标签，用作检查点的唯一标识符，如果未提供，则使用全局步数。标签名称在所有等级上必须相同。

  • client_state – 可选。用于在客户端代码中保存所需训练状态的字典。

  • save_latest – 可选。保存一个名为“最新”的文件，指向最新保存的检查点。

  • exclude_frozen_parameters – 可选。从检查点状态中排除冻结的参数。

重要提示：所有进程都必须调用此方法，而不仅仅是排名为 0 的进程。因为每个进程都需要保存其主权重和调度器 + 优化器状态。如果仅对排名为 0 的进程调用此方法，它将挂起，等待与其他进程同步。

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)[source]

将 ZeRO 2 或 3 检查点转换为单个 fp32 合并状态字典，该字典可以使用 load_state_dict() 加载，并在没有 DeepSpeed 的情况下用于训练，或者与其他人共享，例如通过模型中心。

参数

  • checkpoint_dir (-) – 指向所需检查点文件夹的路径。

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

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

返回值

  • pytorch state_dict

注意：如果您的应用程序没有足够的可用 CPU 内存，此方法可能无法正常工作，您可能需要使用 zero_to_fp32.py 脚本（与检查点一起保存）的离线方法。

典型的用法可能是

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。

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_file, 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_file (-) – 指向 pytorch fp32 状态字典输出文件的路径（例如 path/pytorch_model.bin）。

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