模型检查点

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)[来源]

将ZeRO 2或3检查点转换为单个fp32合并的state_dict，可以使用load_state_dict()加载，并用于无需DeepSpeed的训练或与他人共享，例如通过模型中心。

Parameters

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

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

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

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

Returns

  • 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)[来源]

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

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

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

Parameters

  • model (-) – 要更新的模型对象

  • checkpoint_dir (-) – 指向所需检查点文件夹的路径。（包含标签文件夹的路径，例如 global_step14）

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

Returns

修改后的模型

Return type

  • ``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)[来源]

将ZeRO 2或3检查点转换为单个fp32合并的state_dict文件，该文件可以使用torch.load(file) + load_state_dict()加载，并用于无需DeepSpeed的训练。

Parameters

  • 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'))[来源]

返回item的副本，其中所有包含的张量都被替换为指定设备上的克隆。 适用于单个张量，以及包含/嵌套在列表、元组和字典中的张量。

Parameters

  • item (-) – 要克隆的张量或（可能是嵌套的）张量容器。

  • device (-) – 目标设备（默认为‘cpu’）

Returns

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