分布式检查点 - torch.distributed.checkpoint

分布式检查点 (DCP) 支持从多个等级并行加载和保存模型。它处理加载时重新分片，这使您能够在一个集群拓扑中保存并在另一个集群拓扑中加载。

DCP 在以下几个重要方面不同于 torch.save 和 torch.load

• 它为每个检查点生成多个文件，每个等级至少有一个文件。

• 它在原地操作，这意味着模型应首先分配其数据，然后 DCP 使用该存储而不是重新分配。

加载和保存检查点的入口点如下所示

torch.distributed.checkpoint.state_dict_saver.save(state_dict, *, checkpoint_id=None, storage_writer=None, planner=None, process_group=None)

以 SPMD 样式保存分布式模型。

此函数不同于 torch.save()，因为它处理 ShardedTensor 和 DTensor，每个等级只保存其本地分片。

对于每个 Stateful 对象（同时具有 state_dict 和 load_state_dict），保存将在序列化之前调用 state_dict。

警告

对于保存的 state_dict，在 PyTorch 版本之间没有向后兼容性的保证。

警告

如果使用 process_group 参数，请确保只有其等级调用 save_state_dict 并且 state_dict 中的所有数据都属于它。

注意

保存 FSDP 的 ShardingStrategy.HYBRID_SHARD 的检查点时，只有分片组中的一个应该调用 save_state_dict，并且需要传入相应的进程组。

注意

如果没有进程组可用，此函数将假设意图是在本地进程中保存

state_dict。

参数

• state_dict (Dict[str, Any]) – 要保存的 state_dict。

• checkpoint_id (Union[str, os.PathLike, None]) – 此检查点实例的 ID。checkpoint_id 的含义取决于存储。它可以是文件夹或文件的路径。如果存储是键值存储，它也可以是键。（默认值：None)

• storage_writer (Optional[StorageWriter]) – 用于执行写入的 StorageWriter 实例。如果未指定，DCP 将根据 checkpoint_id 自动推断写入器。如果 checkpoint_id 也为 None，则会引发异常。（默认值：None)

• planner (Optional[SavePlanner]) – SavePlanner 实例。如果未指定，将使用默认计划程序。（默认值：None)

• process_group (Optional[ProcessGroup]) – 用于跨等级同步的 ProcessGroup。（默认值：None)

返回值

保存的检查点的元数据对象。

返回类型

Metadata

示例

>>> my_model = MyModule()

>>> state_dict = {"model": my_model}

>>> fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter("/checkpoint/1")
>>> torch.distributed.checkpoint.save(
>>>     state_dict=state_dict,
>>>     storage_writer=fs_storage_writer,
>>> )

注意

save_state_dict 使用集体来协调跨等级的写入。对于基于 NCCL 的进程组，在通信发生之前，必须将对象的内部张量表示移动到 GPU 设备。在这种情况下，使用的设备由 torch.cuda.current_device() 给出，用户有责任确保通过 torch.cuda.set_device() 设置此设备，以便每个等级都有一个单独的 GPU。

torch.distributed.checkpoint.state_dict_saver.async_save(state_dict, *, checkpoint_id=None, storage_writer=None, planner=None, process_group=None)

save 的异步版本。此代码首先在 CPU 上取消分阶段 state_dict，然后在单独的线程中调用 save。

警告

此功能处于实验阶段，可能会发生变化。

参数

• state_dict (Dict[str, Any]) – 要保存的 state_dict。

• checkpoint_id (Union[str, os.PathLike, None]) – 此检查点实例的 ID。checkpoint_id 的含义取决于存储。它可以是文件夹或文件的路径。如果存储是键值存储，它也可以是键。（默认值：None)

• storage_writer (Optional[StorageWriter]) – 用于执行写入的 StorageWriter 实例。如果未指定，DCP 将根据 checkpoint_id 自动推断写入器。如果 checkpoint_id 也为 None，则会引发异常。（默认值：None)

• planner (Optional[SavePlanner]) – SavePlanner 实例。如果未指定，将使用默认计划程序。（默认值：None)

• process_group (Optional[ProcessGroup]) – 用于跨等级同步的 ProcessGroup。（默认值：None)

返回值

保存save方法返回的元数据对象的未来。

返回类型

未来

示例

>>> my_model = MyModule()

>>> state_dict = {"model": my_model}

>>> fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter("/checkpoint/1")
>>> checkpoint_future = torch.distributed.checkpoint.async_save(
>>>     state_dict=state_dict,
>>>     storage_writer=fs_storage_writer,
>>> )
>>>
>>> # ... do some work ...
>>>
>>> checkpoint_future.result()

torch.distributed.checkpoint.state_dict_saver.save_state_dict(state_dict, storage_writer, process_group=None, coordinator_rank=0, no_dist=False, planner=None)

此方法已弃用。请切换到 'save' 方法。

返回类型

Metadata

torch.distributed.checkpoint.state_dict_loader.load(state_dict, *, checkpoint_id=None, storage_reader=None, planner=None, process_group=None)

在 SPMD 样式中加载分布式 state_dict。

每个进程将尝试读取最少的数据量来满足请求的 state_dict。在加载 ShardedTensor 或 DTensor 实例时，每个进程只读取其本地分片的相关数据。

对于每个 Stateful 对象（同时具有 state_dict 和 load_state_dict），load 方法会首先调用 state_dict 来进行反序列化，之后在反序列化完成后调用 load_state_dict。

警告

state_dict 中的所有张量必须在调用此函数之前，在目标设备上分配。

所有非张量数据都使用 torch.load() 进行加载，并在 state_dict 上就地修改。

警告

用户必须在根模块上调用 load_state_dict 以确保加载后处理和非张量数据能够正确传播。

参数

• state_dict (Dict[str, Any]) – 要保存的 state_dict。

• checkpoint_id (Union[str, os.PathLike, None]) – 此检查点实例的 ID。checkpoint_id 的含义取决于存储。它可以是文件夹或文件的路径。如果存储是键值存储，它也可以是键。（默认值：None)

• storage_reader (Optional[StorageReader]) – 用于执行读取操作的 StorageWriter 实例。如果未指定此参数，DCP 会根据 checkpoint_id 自动推断读取器。如果 checkpoint_id 也为 None，则会抛出异常。（默认： None)

• planner (Optional[LoadPlanner]) – LoadPlanner 实例。如果未指定此参数，将使用默认的规划器。（默认： None)

• process_group (Optional[ProcessGroup]) – 用于跨等级同步的 ProcessGroup。（默认值：None)

返回值

无。

返回类型

无

示例

>>> my_model = MyModule()
>>> optimizer = Adagrad(my_model.parameters())
>>> model_state_dict = my_model.state_dict()
>>> fs_storage_reader = torch.distributed.checkpoint.FileSystemReader("/checkpoint/1")

>>> torch.distributed.checkpoint.load_state_dict(
>>>     state_dict=model_state_dict,
>>>     storage_reader=fs_storage_reader,
>>> )

>>> # module.load_state_dict() function might have customized steps
>>> # to flush the state_dict, must call it to
>>> # ensure correct behavior.
>>> my_model.load_state_dict(model_state_dict)

注意

load_state_dict 方法使用集体操作来协调跨进程的读取。对于基于 NCCL 的进程组，对象的内部张量表示必须在进行通信之前移动到 GPU 设备上。在这种情况下，使用的设备由 torch.cuda.current_device() 指定，用户负责确保每个进程都拥有单独的 GPU，方法是使用 torch.cuda.set_device()。

torch.distributed.checkpoint.state_dict_loader.load_state_dict(state_dict, storage_reader, process_group=None, coordinator_rank=0, no_dist=False, planner=None)

此方法已弃用。请切换到 'load' 方法。

以下模块也适用于对异步检查点机制进行额外定制（torch.distributed.checkpoint.async_save）

class torch.distributed.checkpoint.staging.AsyncStager(*args, **kwargs)

此协议旨在为 dcp.async_save 提供定制和扩展性，允许用户自定义在并行执行常规 dcp.save 路径之前如何对数据进行暂存。预期的操作顺序（在 torch.distributed.state_dict_saver.async_save 中具体定义）如下：

1. AsyncStager.stage_data(state_dict)

   此调用使 AsyncStager 有机会对 state_dict 进行“暂存”。在此上下文中，暂存的预期和目的是创建 state_dict 的“训练安全”表示形式，这意味着在暂存完成后对模块数据的任何更新都应该不会反映在从该方法返回的 state_dict 中。例如，在默认情况下，会在 CPU RAM 上创建整个 state_dict 的副本并在此返回，允许用户继续训练而不会冒更改正在序列化的数据的风险。

2. 并行调用 dcp.save 对从暂存方法返回的 state_dict 进行操作。此调用负责

   对 state_dict 进行序列化并将它写入存储。

3. 如果 AsyncStager.should_synchronize_after_execute 为 True，则此方法将在

   序列化线程启动后，在从 dcp.async_save 返回之前立即调用。如果此参数设置为 False，则假设用户已经为进一步优化保存延迟在训练循环中定义了自定义同步点（例如，通过将暂存与前向/反向传递重叠），并且用户有责任在适当的时间调用 AsyncStager.synchronize_staging。

property should_synchronize_after_execute: bool

是否在执行暂存后进行同步。

stage(state_dict)

返回 state_dict 的“暂存”副本。暂存副本的预期是，它不受在暂存调用完成后发生的任何更新的影响。

返回类型

Dict[str, Union[StatefulT, Any]]

synchronize_staging()

在 stage 异步进行的情况下，应该调用此方法以确保暂存完成，并且可以安全地开始修改原始 state_dict

class torch.distributed.checkpoint.staging.BlockingAsyncStager(cache_staged_state_dict=False, type_check=False)

AsyncStager 的一种实现，它在 CPU RAM 上对 state_dict 进行暂存，并阻塞直到复制完成。此实现还提供了一个选项，可以使用固定内存来优化暂存延迟。

注意：在这种情况下，synchronize_staging 是一个空操作。

stage(state_dict)

返回 state_dict 在 CPU 上的副本。

返回类型

Dict[str, Union[StatefulT, Any]]

synchronize_staging()

空操作函数，因为暂存是阻塞的。

除了上述入口点之外，Stateful 对象（如下所述）在保存/加载期间提供额外的定制。 .. automodule:: torch.distributed.checkpoint.stateful

class torch.distributed.checkpoint.stateful.Stateful(*args, **kwargs)

用于可进行检查点和恢复的对象的 Stateful 协议。

load_state_dict(state_dict)

从提供的 state_dict 中恢复对象的状态。

参数

state_dict (Dict[str, Any]) – 要从中恢复的 state dict

state_dict()

对象应将其 state_dict 表示形式作为字典返回。该函数的输出将被检查点，并在以后在 load_state_dict() 中恢复。

警告

由于恢复检查点的就地性质，此函数在 torch.distributed.checkpoint.load 期间也会被调用。

返回值

对象的 state dict

返回类型

Dict

此 示例 展示了如何使用 Pytorch Distributed Checkpoint 保存 FSDP 模型。

以下类型定义了检查点期间使用的 IO 接口

class torch.distributed.checkpoint.StorageReader

load_state_dict 用于从存储读取数据的接口。

一个 StorageReader 实例在分布式检查点中同时充当协调器和跟随者。作为初始化的一部分，每个实例都会被告知其角色。

子类应该预期 load_state_dict 的以下调用顺序

0. (所有排名) 如果用户传递有效的 checkpoint_id，则设置 checkpoint_id。

1. (所有排名) read_metadata()

2. (所有排名) set_up_storage_reader()

3. (所有排名) prepare_local_plan()

4. (协调器) prepare_global_plan()

5. (所有排名) read_data()

abstract prepare_global_plan(plans)

执行存储加载的集中式规划。

此方法仅在协调器实例上调用。

虽然此方法可以生成完全不同的计划，但首选方法是在 LoadPlan::storage_data 中存储特定于存储的数据。

参数

plans (List[LoadPlan]) – LoadPlan 实例的列表，每个排名一个。

返回值

存储全局规划后，经过转换的 LoadPlan 的列表

返回类型

List[LoadPlan]

abstract prepare_local_plan(plan)

执行特定于存储的本地规划。

虽然此方法可以生成完全不同的计划，但推荐的方法是在 LoadPlan::storage_data 中存储特定于存储的数据。

参数

plan (LoadPlan) – 使用中的 LoadPlan 中的本地计划。

返回值

存储本地规划后，经过转换的 LoadPlan

返回类型

LoadPlan

abstract read_data(plan, planner)

使用 planner 从 plan 中读取所有项目以解析数据。

子类应调用 LoadPlanner::load_bytes 将 BytesIO 对象反序列化到正确的位置。

子类应调用 LoadPlanner::resolve_tensor 以获取对应加载数据的张量的访问权限。

StorageLayer 负责正确调度所需的任何跨设备复制操作。

参数

• plan (LoadPlan) – 要执行的本地计划

• planner (LoadPlanner) – 用于解析项目的规划器对象。

返回值

所有读取完成后完成的 future。

返回类型

Future[None]

abstract read_metadata()

读取检查点元数据。

返回值

与正在加载的检查点关联的元数据对象。

返回类型

Metadata

abstract reset(checkpoint_id=None)

调用以指示将要进行全新的检查点读取。如果用户为该检查点读取设置了 checkpoint_id，则可能会存在 checkpoint_id。checkpoint_id 的含义取决于存储。它可以是文件夹/文件的路径，也可以是键值存储的键。(默认值: None)

参数

checkpoint_id (Union[str, os.PathLike, None]) – 该检查点实例的 ID。checkpoint_id 的含义取决于存储。它可以是文件夹的路径或文件的路径。如果存储更像键值存储，它也可以是键。

abstract set_up_storage_reader(metadata, is_coordinator)

初始化此实例。

参数

• metadata (Metadata) – 要使用的元数据模式。

• is_coordinator (bool) – 此实例是否负责协调检查点。

abstract classmethod validate_checkpoint_id(checkpoint_id)

检查给定的 checkpoint_id 是否受 stroage 支持。这使我们能够启用自动存储选择。

返回类型

bool

class torch.distributed.checkpoint.StorageWriter

save_state_dict 用于写入存储的接口。

一个 StorageWriter 实例在分布式检查点中同时充当协调器和跟随者。作为初始化的一部分，每个实例都会被告知其角色。

子类应该预期以下调用顺序。

0. (所有排名) 如果用户传递有效的 checkpoint_id，则设置 checkpoint_id。

1. (所有排名) set_up_storage_writer()

2. (所有排名) prepare_local_plan()

3. (协调器) prepare_global_plan()

4. (所有排名) write_data()

5. (协调器) finish()

abstract finish(metadata, results)

写入元数据并将当前检查点标记为成功。

用于序列化 metadata 的实际格式/模式是实现细节。唯一的要求是它可以恢复到相同的对象图。

参数

• metadata (Metadata) – 新检查点的元数据

• results (List[List[WriteResult]]) – 所有排名提供的 WriteResults 列表。

返回值

无

返回类型

无

abstract prepare_global_plan(plans)

执行存储的集中式规划。

此方法仅在协调器实例上调用。

虽然此方法可以生成完全不同的计划，但首选方法是在 SavePlan::storage_data 中存储特定于存储的数据。

参数

plans (List[SavePlan]) – 每个进程的 SavePlan 实例列表。

返回值

存储全局规划后，经过转换的 SavePlan 列表。

返回类型

List[SavePlan]

abstract prepare_local_plan(plan)

执行特定于存储的本地规划。

虽然此方法可以生成完全不同的计划，但推荐的方法是在 SavePlan::storage_data 中存储存储特定数据。

参数

plan (SavePlan) – 当前使用的 SavePlanner 中的本地计划。

返回值

存储本地规划后，经过转换的 SavePlan。

返回类型

SavePlan

abstract reset(checkpoint_id=None)

调用此方法表示即将发生新的检查点写入。如果用户为此检查点写入设置了 checkpoint_id，则可能存在 checkpoint_id。checkpoint_id 的含义取决于存储。它可以是文件夹/文件的路径或键值存储的键。

参数

checkpoint_id (Union[str, os.PathLike, None]) – 此检查点实例的 ID。checkpoint_id 的含义取决于存储。它可以是文件夹或文件的路径。如果存储是键值存储，它也可以是键。（默认值：None)

abstract set_up_storage_writer(is_coordinator)

初始化此实例。

参数

is_coordinator (bool) – 此实例是否负责协调检查点。

storage_meta()

返回存储特定的元数据。这用于在检查点中存储其他信息，这些信息可能有助于提供请求级可观察性。StorageMeta 在保存调用期间传递给 SavePlanner。默认情况下返回 None。

TODO：提供示例

返回类型

Optional[StorageMeta]

abstract classmethod validate_checkpoint_id(checkpoint_id)

检查给定的 checkpoint_id 是否受 stroage 支持。这使我们能够启用自动存储选择。

返回类型

bool

abstract write_data(plan, planner)

使用 planner 写入 plan 中的所有项目以解析数据。

子类应该对计划中的每个项目调用 SavePlanner::resolve_data 以访问要写入的底层对象。

子类应该延迟调用 resolve_data，因为它可能会分配内存。在张量的情况下，做出以下假设

• 它们可能位于任何设备上，包括与 WriteItem::tensor_data 上的设备不匹配的设备。

• 它们可能是视图，也可能不是连续的。只需要保存投影。

参数

• plan (SavePlan) – 要执行的保存计划。

• planner (SavePlanner) – 用于解析项目到数据的计划程序对象。

返回值

完成到 WriteResult 列表的未来。

返回类型

Future[List[WriteResult]]

以下类型定义了检查点期间使用的计划程序接口

class torch.distributed.checkpoint.LoadPlanner

定义了 load_state_dict 用于规划加载过程的协议的抽象类。

LoadPlanner 是有状态的对象，可用于自定义整个加载过程。

LoadPlanner 充当 state_dict 的访问代理，因此对它的任何转换都将对整个进程可见。

在 load_state_dict 期间，计划程序子类可以预期以下调用顺序

1. set_up_planner - 在所有进程上调用。

   表示加载检查点开始。

2. create_local_plan - 在所有进程上调用。

   处理 state_dict 并生成一个将发送进行全局规划的 LoadPlan。

3. create_global_plan - 仅在协调器进程上调用。

   接收来自所有进程的 LoadPlan 并做出任何全局决策。

4. load_bytes - 在每个进程上调用多次

   这将对 state_dict 中的每个非张量值调用一次。

5. resolve_tensor 和 commit_tensor - 在每个进程上调用多次

   它们将对 state_dict 中的每个张量值成对调用。

建议用户扩展 DefaultLoadPlanner 而不是直接扩展此接口，因为大多数更改可以通过单个方法的更改来表达。

有两种常见的扩展模式

重写 state_dict。这是扩展加载过程的最简单方法，因为它不需要了解 LoadPlan 如何工作的复杂性。我们需要保留对原始 state_dict 的引用，因为加载是在原地进行的，因此我们需要能够在原地执行它

>>> class RenamePlanner(DefaultLoadPlanner):
>>>     def set_up_planner(
>>>         self,
>>>         state_dict: STATE_DICT_TYPE,
>>>         metadata: Metadata,
>>>         is_coordinator: bool,
>>>     ) -> None:
>>>         self.original_state_dict = state_dict
>>>         state_dict = {"foo_" + k: v for k, v in state_dict.items()}
>>>
>>>         if self.flatten_sharded_tensors:
>>>             state_dict = _flatten_sharded_tensors(state_dict)
>>>
>>>         if self.flatten_state_dict:
>>>             state_dict, self.mappings = flatten_state_dict(state_dict)
>>>
>>>         self.state_dict = state_dict
>>>         self.metadata = metadata
>>>         self.is_coordinator = is_coordinator
>>>
>>>     def load_bytes(self, read_item, value):
>>>         # Remove the "foo_" prefix
>>>         self.original_state_dict[read_item.dest_index.fqn[4:]] = torch.load(value)

修改 resolve_tensor 和 commit_tensor 以处理加载时转换。

>>> class MetaModelMaterialize(DefaultSavePlanner):
>>>     def resolve_tensor(self, read_item):
>>>         tensor = super().resolve_tensor(read_item)
>>>         return torch.empty_like(tensor, device="cpu")
>>>
>>>     def commit_tensor(self, read_item, tensor):
>>>         self.state_dict[read_item.dest_index.fqn] = tensor

abstract commit_tensor(read_item, tensor)

在 StorageReader 完成将数据加载到 tensor 中后调用一次。

提供的张量与调用 resolve_tensor 返回的张量相同。仅当此 LoadPlanner 需要在将 tensor 复制回 state_dict 中的张量之前对其进行后处理时，才需要此方法。

张量的内容将遵循其设备同步模型。

abstract create_global_plan(global_plan)

计算全局加载计划并返回每个进程的计划。

. 注意：这仅在协调器进程上调用

返回类型

List[LoadPlan]

abstract create_local_plan()

根据 state_dict 和 set_up_planner 提供的元数据创建 LoadPlan。

. 注意：这在每个进程上调用。

返回类型

LoadPlan

abstract finish_plan(central_plan)

接受来自协调器的计划并返回最终的 LoadPlan。

返回类型

LoadPlan

abstract load_bytes(read_item, value)

加载由 read_item``和 ``value 描述的项目。

此方法预计将修改底层 state_dict 的位置。

value 的内容由用于生成正在加载的检查点的 SavePlanner 定义。

resolve_bytes(read_item)

返回 StorageReader 用于加载 read_item 的 BytesIO。

BytesIO 应该与底层 state_dict 中的 BytesIO 具有别名，因为 StorageReader 将替换其内容。

返回类型

BytesIO

abstract resolve_tensor(read_item)

返回由 read_item 描述的张量，以便 StorageReader 用于加载 read_item。

张量应该与基础 state_dict 中的一个张量建立别名，因为 StorageReader 将替换其内容。 如果由于任何原因，这不可行，规划器可以使用 commit_tensor 方法将数据复制回 state_dict 中的张量。

返回类型

张量

abstract set_up_planner(state_dict, metadata=None, is_coordinator=False)

初始化此实例以将数据加载到 state_dict 中。

. 注意：这在每个进程上调用。

class torch.distributed.checkpoint.LoadPlan(items: List[torch.distributed.checkpoint.planner.ReadItem], storage_data: Any = None, planner_data: Any = None)

class torch.distributed.checkpoint.ReadItem(type: torch.distributed.checkpoint.planner.LoadItemType, dest_index: torch.distributed.checkpoint.metadata.MetadataIndex, dest_offsets: torch.Size, storage_index: torch.distributed.checkpoint.metadata.MetadataIndex, storage_offsets: torch.Size, lengths: torch.Size)

class torch.distributed.checkpoint.SavePlanner

抽象类，定义 save_state_dict 用于规划保存过程的协议。

SavePlanner 是有状态的对象，可用于自定义整个保存过程。

SavePlanner 充当 state_dict 的访问代理，因此对它的任何转换都将对整个过程可见。

规划器子类可以预期在 save_state_dict 期间以下列顺序调用：

1. set_up_planner - 在所有进程上调用。

   发出检查点保存开始的信号。

2. create_local_plan - 在所有进程上调用。

   处理 state_dict 并生成一个 SavePlan，该计划将被发送用于全局规划。

3. create_global_plan - 仅在协调器进程上调用。

   获取所有等级的 SavePlan，并做出任何全局决策。

4. finish_plan - 在所有等级上调用。

   这使每个等级有机会适应全局规划决策。

5. resolve_data - 在每个等级上调用多次。

   在 state_dict 上查找存储层要写入的值。

建议用户扩展 DefaultSavePlanner 而不是直接扩展此接口，因为大多数更改可以通过单个方法中的更改来表达。

扩展通常有 3 种模式：

重写 state_dict。 这是扩展保存过程的最简单方法，因为它不需要了解 SavePlan 的工作原理。

>>> class RenamePlanner(DefaultSavePlanner):
>>>     def set_up_planner(
>>>         self,
>>>         state_dict: STATE_DICT_TYPE,
>>>         storage_meta: Optional[StorageMeta],
>>>         is_coordinator: bool,
>>>     ) -> None:
>>>         # prefix all keys with `foo_``
>>>         super().set_up_planner({"foo_" + k: v for k, v in state_dict.items()}, storage_meta, is_coordinator)

同时修改本地计划和查找。 这在需要精确控制如何持久化数据时很有用。

>>> class FP16Planner(DefaultSavePlanner):
>>>     def create_local_plan(self):
>>>         plan = super().create_local_plan()
>>>         for p in plan:
>>>             if p.tensor_data is not None:
>>>                 p.tensor_data.properties.dtype = torch.float16
>>>         return plan
>>>
>>>     def resolve_data(self, write_item):
>>>         item = super().resolve_data(write_item)
>>>         return item if write_item.type == WriteItemType.BYTE_IO else item.to(torch.float16)

使用全局规划步骤做出不能由每个等级单独做出的集中决策。

>>> from itertools import islice
>>> from dataclasses import replace
>>> class DDPLoadBalancingPlanner(DefaultSavePlanner):
>>>     # This uses the default local plan behavior of having all non-sharded writes in rank 0
>>>     # This sample doesn't handle ShardedTensors
>>>     def create_global_plan(self, all_plans):
>>>         def chunk(it, size):
>>>             it = iter(it)
>>>         return list(iter(lambda: tuple(islice(it, size)), ()))
>>>         all_plans = [
>>>             replace(plan, items=items) for plan, items in
>>>                 zip(all_plans, chunk(all_plans[0].items, len(all_plans)))
>>>         ]
>>>         return super().create_global_plan(all_plans)

最后，一些规划器需要在检查点中保存额外的元数据，这是通过让每个等级在其本地计划中贡献其数据项并让全局规划器聚合它们来完成的。

>>> class SaveExtraDataPlanner(DefaultSavePlanner):
>>>     def create_local_plan(self) -> SavePlan:
>>>         plan = super().create_local_plan()
>>>         return replace(plan, planner_data="per-rank-data")
>>>
>>>     def create_global_plan(self, all_plans: List[SavePlan]) -> Tuple[List[SavePlan], Metadata]:
>>>         global_plan, metadata = super().create_global_plan(all_plans)
>>>         merged_data = [p.planner_data for p in global_plan]
>>>         metadata = replace(metadata, planner_data=merged_data)
>>>         return global_plan, metadata

abstract create_global_plan(all_plans)

计算全局检查点计划并返回每个等级的本地计划。

仅在协调器等级上调用。

返回类型

Tuple[List[SavePlan], Metadata]

abstract create_local_plan()

计算当前等级的保存计划。

这将被聚合并传递给 create_global_plan。 规划器特定数据可以通过 SavePlan::planner_data 传递。

在所有等级上调用。

返回类型

SavePlan

abstract finish_plan(new_plan)

合并由 create_local_plan 创建的计划和 create_global_plan 的结果。

在所有等级上调用。

返回类型

SavePlan

abstract resolve_data(write_item)

转换和准备 write_item（来自 state_dict）以供存储，确保幂等性和线程安全性。

在 state_dict 中查找与 write_item 关联的对象，并在存储层使用它之前应用任何转换（如序列化）。

在每个等级上调用多次，每个等级至少调用一次（在最终 SavePlan 中的每个 WriteItem 上调用一次）。

此方法应该是幂等且线程安全的。 StorageWriter 实现可以根据需要尽可能频繁地调用它。

任何分配内存的转换都应该在调用此方法时延迟完成，以减少检查点所需的峰值内存。

返回张量时，它们可以位于任何设备或格式上，它们也可以是视图。 存储层负责弄清楚如何保存它们。

返回类型

Union[Tensor, BytesIO]

abstract set_up_planner(state_dict, storage_meta=None, is_coordinator=False)

初始化此规划器以保存 state_dict。

实现应该保存这些值，因为它们在保存过程中不会被提供。

在所有等级上调用。

class torch.distributed.checkpoint.SavePlan(items: List[torch.distributed.checkpoint.planner.WriteItem], storage_data: Any = None, planner_data: Any = None)

class torch.distributed.checkpoint.planner.WriteItem(index, type, tensor_data=None)

用于保存存储信息的 Dataclass。

tensor_storage_size()

计算底层张量的存储大小，如果这不是张量写入，则返回 None。

返回值

可选 [int] 存储大小，以字节为单位，如果存在底层张量。

返回类型

Optional[int]

我们提供基于文件系统的存储层

class torch.distributed.checkpoint.FileSystemReader(path)

property checkpoint_id: Union[str, PathLike]

返回用于保存检查点的 checkpoint_id。

class torch.distributed.checkpoint.FileSystemWriter(path, single_file_per_rank=True, sync_files=True, thread_count=1, per_thread_copy_ahead=10000000, cache_staged_state_dict=False, overwrite=True)

使用文件 IO 的 StorageWriter 的基本实现。

此实现做出以下假设和简化

• 检查点路径为空或不存在的目录。

• 文件创建是原子的

检查点包含每个写入请求一个文件，以及一个带有序列化元数据的 .metadata 文件。

stage(state_dict)

AsyncStager.stage 的重写

返回类型

Dict[str, Union[StatefulT, Any]]

我们提供 LoadPlanner 和 SavePlanner 的默认实现，可以处理所有 torch.distributed 结构，例如 FSDP、DDP、ShardedTensor 和 DistributedTensor。

class torch.distributed.checkpoint.DefaultSavePlanner(flatten_state_dict=True, flatten_sharded_tensors=True, dedup_replicated_tensors=None, dedup_save_to_lowest_rank=False)

lookup_object(index)

从 planner 接口扩展，使其易于扩展默认 planner。

返回类型

任何

transform_object(write_item, object)

从 planner 接口扩展，使其易于扩展默认 planner。

class torch.distributed.checkpoint.DefaultLoadPlanner(flatten_state_dict=True, flatten_sharded_tensors=True, allow_partial_load=False)

在 LoadPlanner 之上添加多个功能的 DefaultLoadPlanner。

特别是它添加了以下内容

flatten_state_dict: 处理具有嵌套字典的 state_dict flatten_sharded_tensors: 对于 2D 并行模式下的 FSDP allow_partial_load: 如果为 False，则如果 state_dict 中存在键，但在检查点中不存在，则会引发运行时错误。

lookup_tensor(index)

从 planner 接口扩展，使其易于扩展默认 planner。

返回类型

张量

transform_tensor(read_item, tensor)

从 planner 接口扩展，使其易于扩展默认 planner。

由于遗留的设计决策，FSDP 和 DDP 的状态字典可能具有不同的键或完全限定名称（例如，layer1.weight），即使原始的非并行模型相同。此外，FSDP 提供了各种类型的模型状态字典，例如完整状态字典和分片状态字典。此外，优化器状态字典使用参数 ID 而不是完全限定名称来标识参数，这可能在使用并行性（例如，管道并行性）时会导致问题。

为了解决这些挑战，我们为用户提供了一组 API，以便轻松管理 state_dict。get_model_state_dict 返回一个模型状态字典，其键与非并行模型状态字典返回的键一致。类似地，get_optimizer_state_dict 提供优化器状态字典，其键在应用的所有并行性中保持一致。为了实现这种一致性，get_optimizer_state_dict 将参数 ID 转换为与非并行模型状态字典中找到的完全限定名称相同的名称。

请注意，这些 API 返回的结果可以直接与 torch.distributed.checkpoint.save() 和 torch.distributed.checkpoint.load() 方法一起使用，无需任何其他转换。

请注意，此功能是实验性的，API 签名将来可能会更改。

torch.distributed.checkpoint.state_dict.get_state_dict(model, optimizers, *, submodules=None, options=None)

返回模型 state_dict 和优化器 state_dict。

get_state_dict 可以处理由 PyTorch FSDP/fully_shard、DDP/replicate、tensor_parallel/parallelize_module 以及这些并行性的任何组合并行的任何模块。 get_state_dict 的主要功能是：1.) 返回可以用不同数量的训练器和/或不同并行性重新分片的模型和优化器 state_dict。 2.) 隐藏特定于并行性的 state_dict API。用户无需调用这些 API。 3.) 对结果 state_dict 进行健全性检查。

结果状态字典的键是规范的 FQN（完全限定名称）。规范的 FQN 指的是基于参数在 nn.Module 层次结构中的位置的 FQN。更具体地说，参数的规范 FQN 是当模块未由任何并行性分布时由 module.named_parameters() 或 module.named_buffers() 返回的 FQN。

由于优化器在内部使用参数 ID 来表示参数，因此在调用此 API 时，将从参数 ID 转换为规范的 FQN。

示例

>>> import torch
>>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
>>> from torch.nn.parallel import DistributedDataParallel as DDP
>>> from torch.distributed.checkpoint.state_dict import get_state_dict

>>> fsdp_model = FSDP(copy.deepcopy(model))
>>> fsdp_optim = torch.optim.Adam(model.parameters(), lr=1e-3)
>>> ddp_model = DDP(copy.deepcopy(model))
>>> ddp_optim = torch.optim.Adam(model.parameters(), lr=1e-3)

>>> ddp_state_dict, ddp_optim_state_dict = get_state_dict(ddp_model, ddp_optim)
>>> fsdp_state_dict, fsdp_optim_state_dict = get_state_dict(fsdp_model, fsdp_optim)

>>> # if we simply call ddp_model.state_dict() and fsdp_model.state_dict(),
>>> # the asserts will fail.
>>> assert ddp_state_dict == fsdp_state_dict
>>> assert ddp_optim_state == fsdp_optim_state_dict

参数

• get_state_dict 也可以处理未并行的模块。在这种情况下，get_state_dict 只执行一个功能——将优化器参数 ID 转换为规范的 FQN。

• 优化器 (Union[None, 优化器, Iterable[优化器]]) – 用于优化 model 的优化器。

• 子模块 (已弃用) – Optional[Set[nn.Module]]: 仅返回属于子模块的模型参数。

• 选项 (StateDictOptions) – 用于控制模型状态字典和优化器状态字典返回值的选项。有关详细信息，请参阅 StateDictOptions。

返回值

Tuple，其中包含模型状态字典和优化器状态字典。

返回类型

Tuple[Dict[str, ValueType], OptimizerStateType]

torch.distributed.checkpoint.state_dict.get_model_state_dict(model, *, submodules=None, options=None)

返回 model 的模型状态字典。

有关详细信息，请参阅 get_state_dict。

参数

• get_state_dict 也可以处理未并行的模块。在这种情况下，get_state_dict 只执行一个功能——将优化器参数 ID 转换为规范的 FQN。

• 子模块 (已弃用) – Optional[Set[nn.Module]]: 仅返回属于子模块的模型参数。

• 选项 (StateDictOptions) – 用于控制模型状态字典和优化器状态字典返回值的选项。有关详细信息，请参阅 StateDictOptions。

返回值

model 的状态字典。

返回类型

Dict[str, ValueType]

torch.distributed.checkpoint.state_dict.get_optimizer_state_dict(model, optimizers, *, submodules=None, options=None)

返回优化器的组合状态字典。

有关详细信息，请参阅 get_state_dict。

参数

• get_state_dict 也可以处理未并行的模块。在这种情况下，get_state_dict 只执行一个功能——将优化器参数 ID 转换为规范的 FQN。

• 优化器 (Union[None, 优化器, Iterable[优化器]]) – 用于优化 model 的优化器。

• 子模块 (已弃用) – Optional[Set[nn.Module]]: 仅返回属于子模块的模型参数。

• 选项 (StateDictOptions) – 用于控制模型状态字典和优化器状态字典返回值的选项。有关详细信息，请参阅 StateDictOptions。

返回值

optimizers 的状态字典。

返回类型

OptimizerStateType

torch.distributed.checkpoint.state_dict.set_state_dict(model, optimizers, *, model_state_dict, optim_state_dict, options=None)

加载模型状态字典和优化器状态字典。

get_state_dict 的对应函数，用于将状态字典设置到模型和优化器。给定的 model_state_dict 和 optim_state_dict 不必由 get_state_dict 返回，但必须满足以下要求：1) 所有 FQN 都是 get_state_dict 中定义的规范 FQN，2) 如果张量被分片，则它必须是 ShardedTensor 或 DTensor，3) 优化器状态字典不能包含参数 ID；键应该是规范的 FQN。

参数

• get_state_dict 也可以处理未并行的模块。在这种情况下，get_state_dict 只执行一个功能——将优化器参数 ID 转换为规范的 FQN。

• 优化器 (Union[优化器, Iterable[优化器]]) – 用于优化 model 的优化器。

• model_state_dict (Dict[str, ValueType]) – (Union[Dict[nn.Module, Dict[str, ValueType]], Dict[str, ValueType]]): 要加载的模型状态字典。如果 model_state_dict 的键是 nn.Module，则键是 model 的子模块，而值应该是子模块的状态字典。加载状态字典时，将子模块的前缀附加到状态字典。

• optim_state_dict (OptimizerStateType) – OptimizerStateType: 要加载的优化器状态字典。

• 选项 (StateDictOptions) – 用于控制模型状态字典和优化器状态字典加载方式的选项。有关详细信息，请参阅 StateDictOptions。

返回值

• missing_keys 是一个包含模型状态字典中缺失键的 str 列表。

• unexpected_keys 是一个包含模型状态字典中意外键的 str 列表。

返回类型

NamedTuple，其中包含 missing_keys 和 unexpected_keys 字段

torch.distributed.checkpoint.state_dict.set_model_state_dict(model, model_state_dict, *, options=None)

加载模型状态字典。

get_model_state_dict 的对应函数，用于将状态字典设置到模型。有关详细信息，请参阅 set_state_dict。

参数

• get_state_dict 也可以处理未并行的模块。在这种情况下，get_state_dict 只执行一个功能——将优化器参数 ID 转换为规范的 FQN。

• model_state_dict (Dict[str, ValueType]) – (Dict[str, ValueType]): 要加载的模型状态字典。如果 model_state_dict 的键是 nn.Module，则键是 model 的子模块，而值应该是子模块的状态字典。加载状态字典时，将子模块的前缀附加到状态字典。

• 选项 (StateDictOptions) – 用于控制模型状态字典和优化器状态字典加载方式的选项。有关详细信息，请参阅 StateDictOptions。

返回值

• missing_keys 是一个包含缺失键的 str 列表

• unexpected_keys 是一个包含意外键的 str 列表

返回类型

NamedTuple，其中包含 missing_keys 和 unexpected_keys 字段

torch.distributed.checkpoint.state_dict.set_optimizer_state_dict(model, optimizers, optim_state_dict, *, options=None)

加载优化器状态字典。

get_optimizer_state_dict 的对应函数，用于将状态字典设置到优化器。有关详细信息，请参阅 set_state_dict。

参数

• get_state_dict 也可以处理未并行的模块。在这种情况下，get_state_dict 只执行一个功能——将优化器参数 ID 转换为规范的 FQN。

• 优化器 (Union[优化器, Iterable[优化器]]) – 用于优化 model 的优化器。

• optim_state_dict (OptimizerStateType) – OptimizerStateType: 要加载的优化器状态字典。

• 选项 (StateDictOptions) – 用于控制模型状态字典和优化器状态字典加载方式的选项。有关详细信息，请参阅 StateDictOptions。

返回值

无

返回类型

无

class torch.distributed.checkpoint.state_dict.StateDictOptions(full_state_dict=False, cpu_offload=False, ignore_frozen_params=False, keep_submodule_prefixes=True, strict=True, broadcast_from_rank0=False, flatten_optimizer_state_dict=False)

此数据类指定 get_state_dict/set_state_dict 的工作方式。

• full_state_dict: 如果将其设置为 True，则返回的状态字典中的所有张量将被收集。返回的状态字典中不会包含任何 ShardedTensor 和 DTensor。

• cpu_offload: 将所有张量卸载到 CPU。为了防止 CPU 内存不足，如果 full_state_dict 也为 True，则只有 rank0 将获得状态字典，而所有其他 rank 将获得空状态字典。

• ignore_frozen_params: 如果值为 True，则返回的状态字典将不包含任何冻结参数 - requires_grad 为 False。默认值为 False。

• keep_submodule_prefixes (已弃用): 当 submodules 不为 None 时，此选项指示是否保留 state_dict 键中的子模块前缀。例如，如果子模块为 module.pretrain 且参数的完整 FQN 为 pretrain.layer1.weight。当此选项为 True 时，参数在返回的 state_dict 中的键将为 pretrain.layer1.weight。如果选项为 False，则键将为 layer1.weight。请注意，如果 keep_submodule_prefixes 为 False，则可能存在冲突的 FQN，因此 submodules 中应只有一个子模块。

• strict: 当 set_state_dict 调用 model.load_state_dict() 时的 strict 选项。

• broadcast_from_rank0: 当此选项为 True 时，rank0 应接收一个

  完整的 state_dict 并将 state_dict/optim_state_dict 中的张量逐个广播到其他 rank。其他 rank 将接收这些张量并根据模型和优化器中的本地分片进行分片。full_state_dict 在使用此选项时必须设置为 True。此选项目前仅支持 DTensor，不支持旧的 ShardedTensor。

对于习惯于使用和共享 torch.save 格式的模型的用户，提供了以下方法，这些方法提供了在格式之间转换的离线实用程序。

torch.distributed.checkpoint.format_utils.dcp_to_torch_save(dcp_checkpoint_dir, torch_save_path)

给定包含 DCP 检查点的目录，此函数将其转换为 Torch 保存文件。

参数

• dcp_checkpoint_dir (Union[str, PathLike]) – 包含 DCP 检查点的目录。

• torch_save_path (Union[str, PathLike]) – 用于存储转换后的 Torch 保存文件的文件名。

警告

为了避免 OOM，建议仅在一个 rank 上运行此函数。

torch.distributed.checkpoint.format_utils.torch_save_to_dcp(torch_save_path, dcp_checkpoint_dir)

给定 torch 保存文件的路径，将其转换为 DCP 检查点。

参数

• torch_save_path (Union[str, PathLike]) – Torch 保存文件的名称。

• dcp_checkpoint_dir (Union[str, PathLike]) – 用于存储 DCP 检查点的目录。

警告

为了避免 OOM，建议仅在一个 rank 上运行此函数。

以下类也可以用于从 torch.save 格式在线加载和重新分片模型。

class torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader(checkpoint_id=None, coordinator_rank=0)

用于读取 Torch Save 文件的 StorageReader。此 reader 将在协调器 rank 上读取整个检查点，然后将每个张量广播并分片到所有 rank。

. 注意，旨在与 DynamicMetaLoadPlanner 一起使用

警告

当前实现仅支持加载张量。

>>> sd = {"mode": model}
>>> dcp.load(
>>>    sd,
>>>    storage_reader=BroadcastingTorchSaveReader(),
>>>    planner=DynamicMetaLoadPlanner(),
>>>    checkpoint_id="path_to_model.pt"
>>> )

prepare_global_plan(global_plan)

StorageReader 方法的实现

返回类型

List[LoadPlan]

prepare_local_plan(plan)

StorageReader 方法的实现

返回类型

LoadPlan

read_data(plan, planner)

在协调器 rank 上读取 torch save 数据，并在之后广播，这会产生通信成本，但避免了在每个 rank 上加载整个检查点，有望防止 OOM 问题

返回类型

Future[None]

read_metadata()

扩展默认 StorageReader 以支持构建元数据文件

返回类型

Metadata

reset(checkpoint_id=None)

StorageReader 方法的实现

set_up_storage_reader(metadata, is_coordinator)

StorageReader 方法的实现

classmethod validate_checkpoint_id(checkpoint_id)

StorageReader 方法的实现

返回类型

bool

class torch.distributed.checkpoint.format_utils.DynamicMetaLoadPlanner(flatten_state_dict=True, flatten_sharded_tensors=True, allow_partial_load=False)

DefaultLoadPlanner 的扩展，它根据传入的 state dict 创建一个新的 Metadata 对象，避免了从磁盘读取元数据的需要。这在读取没有元数据文件的格式（如 Torch Save 文件）时非常有用。

. 注意，旨在与 BroadcastingTorchSaveReader 一起使用

警告

当前实现仅支持加载张量。

>>> sd = {"mode": model}
>>> dcp.load(
>>>    sd,
>>>    storage_reader=BroadcastingTorchSaveReader(),
>>>    planner=DynamicMetaLoadPlanner(),
>>>    checkpoint_id="path_to_model.pt"
>>> )

set_up_planner(state_dict, metadata=None, is_coordinator=False)

规划器的设置，通过从 state dict 创建 Metadata 对象来扩展默认行为

以下实验接口可用于在生产环境中提高可观察性