分布式检查点 - torch.distributed.checkpoint¶
分布式检查点（DCP）支持从多个排名平行加载和保存模型。 它处理加载时的重新分片，从而可以在一种集群拓扑中保存并在另一种集群拓扑中加载。
DCP与torch.save和torch.load在几个重要方面有所不同：
它为每个检查点生成多个文件，每个排名至少一个。
它以就地操作方式进行工作，这意味着模型应先分配其数据，然后DCP使用该存储空间。
加载和保存检查点的入口函数如下：
- torch.distributed.checkpoint.load_state_dict(state_dict, storage_reader, process_group=None, coordinator_rank=0, no_dist=False, planner=None)[source]¶
加载分布式
state_dict的 SPMD 风格。每个 rank 会尝试读取最少必要的数据 以满足所请求的 state_dict。当加载
ShardedTensor个实例时，每个 rank 只会读取其本地分片的数据。警告
所有在
state_dict中的张量必须在其目标设备上分配 之前 调用此函数。所有非张量数据都使用torch.load()加载，并在state_dict上就地修改。
警告
用户必须在根模块上调用load_state_dict以确保加载后处理和非张量数据正确传播。
- Parameters
state_dict (Dict[str, Any]) – 要加载的 state_dict。请注意，此 state_dict 将在原地更新。
storage_reader (StorageReader) – 用于从此处加载数据的 StorageReader。
process_group (ProcessGroup) – 用于跨等级同步的 ProcessGroup。
coordinator_rank (int) – 用于协调检查点的秩（rank）。 默认使用 rank0。
no_dist (bool) – 如果
True, 分布式检查点将不会以SPMD风格保存。 (默认值:False)
- Returns
None.
- Return type
请提供需要翻译的单词列表。
- Examples
>>> 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.save_state_dict(state_dict, storage_writer, process_group=None, coordinator_rank=0, no_dist=False, planner=None)[source]¶
以 SPMD 风格保存分布式模型。
此函数与
torch.save()不同，因为它通过让每个 rank 仅保存其本地分片来处理ShardedTensor。警告
在不同版本的 PyTorch 中，无法保证保存的 state_dicts 的向后兼容性。
警告
如果使用process_group参数，请确保只有它的排名调用save_state_dict，并且state_dict中的所有数据都属于它。
注意
在为FSDP的ShardingStrategy.HYBRID_SHARD保存检查点时，shard_group中应该只有一个调用save_state_dict，并且需要传入相应的进程组。
注意
此函数可以在不初始化进程组的情况下，通过传递
no_dist=True来保存 state_dict。- Parameters
state_dict (Dict[str, Any]) – 要保存的state_dict。
storage_writer (StorageWriter) – StorageWrite 的实例，用于执行写入操作。
process_group (ProcessGroup) – 用于跨等级同步的 ProcessGroup。
coordinator_rank (int) – 用于协调检查点的秩（rank）。 默认使用 rank0。
no_dist (bool) – 如果
True, 分布式检查点将不会以SPMD风格保存。 (默认值:False)
- Returns
保存检查点的元数据对象。
- Return type
元数据
示例
>>> my_model = MyModule()
>>> model_state_dict = my_model.state_dict()
>>> fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter("/checkpoint/1") >>> torch.distributed.checkpoint.save_state_dict( >>> state_dict=model_state_dict, >>> storage_writer=fs_storage_writer, >>> )
注意
save_state_dict 使用集体操作来协调不同排名之间的写入。 对于基于 NCCL 的进程组，在通信发生之前，对象的内部张量表示必须移动到 GPU 设备。 在这种情况下，使用的设备由
torch.cuda.current_device()指定， 并且用户有责任确保每个排名都有一个单独的 GPU，通过torch.cuda.set_device()设置。
这个 示例 展示了如何使用 Pytorch 分布式检查点保存 FSDP 模型。
以下类型定义了检查点过程中使用的 IO 接口：
- class torch.distributed.checkpoint.StorageReader[source]¶
由
load_state_dict使用的从存储中读取的接口。一个 StorageReader 实例在分布式检查点中同时充当协调者和跟随者。 作为初始化的一部分，每个实例都会被告知其角色。
子类应预期以下调用顺序由
load_state_dict：（所有排名）read_metadata()
（所有排名）设置存储阅读器()
（所有排名）prepare_local_plan()
(协调员) prepare_global_plan()
（所有排名）读取数据()
- abstract prepare_global_plan(plans)[source]¶
集中规划存储加载。
此方法仅在协调器实例上被调用。
虽然这种方法可以生成完全不同的计划，但更推荐的做法是将特定于存储的数据存储在 LoadPlan::storage_data 中。
- abstract prepare_local_plan(plan)[source]¶
执行存储特定的本地规划。
虽然这种方法可以生成完全不同的计划，但推荐的做法是将存储特定数据存储在 LoadPlan::storage_data 中。
- class torch.distributed.checkpoint.StorageWriter[source]¶
由
save_state_dict使用的接口，用于写入存储。一个 StorageWriter 实例在一个分布式检查点中同时充当协调者和跟随者。 在初始化过程中，每个实例都会被告知其角色。
一个子类应期望以下调用顺序。
（所有排名）设置存储写入器 ()
（所有排名）prepare_local_plan()
(协调员) prepare_global_plan()
所有排名 write_data()
(协调员) 结束()
- abstract finish(metadata, results)[source]¶
写入元数据并标记当前检查点为成功。
实际用于序列化的metadata格式/模式是一个实现细节。唯一的要求是它可以恢复到相同的对象图。
- abstract prepare_global_plan(plans)[source]¶
集中规划存储。
此方法仅在协调器实例上被调用。
虽然这种方法可以生成完全不同的计划，但更推荐的方式是将特定于存储的数据存储在 SavePlan::storage_data 中。
- abstract prepare_local_plan(plan)[source]¶
执行存储特定的本地规划。
虽然这种方法可以生成完全不同的计划，但推荐的做法是将存储特定数据保存在 SavePlan::storage_data 中。
- abstract set_up_storage_writer(is_coordinator)[source]¶
初始化此实例。
- Parameters
is_coordinator (bool) – 是否此实例负责协调检查点。
以下类型定义了检查点期间使用的计划器接口：
- class torch.distributed.checkpoint.LoadPlanner[source]¶
抽象类，定义了 load_state_dict 使用的协议，以规划加载过程。
LoadPlanner 是有状态的对象，可用于自定义整个加载过程。
LoadPlanner 作为状态字典的访问代理，因此对其所做的任何变换都将对整个过程可见。
在调用 load_state_dict 期间，计划器子类可以预期以下调用顺序：
- set_up_planner - called on all ranks.
表示开始加载检查点。
- create_local_plan - called on all ranks.
处理 state_dict 并生成一个LoadPlan，该值将用于全局规划。
- create_global_plan - called on the coordinator rank only.
从所有 ranks 获取 LoadPlan 并做出任何全局决策。
- load_bytes - called multiple times on each rank
这在状态字典中的每个非张量值上调用一次。
- resolve_tensor and commit_tensor - called multiple times on each rank
它们以成对的方式为 state_dict 中的每个张量值调用。
建议用户扩展 DefaultLoadPlanner 而不是直接扩展此接口，因为大多数更改都可以通过单个方法的更改来表达。
有两种常见的扩展模式：
重写 state_dict。这是扩展加载过程的最简单方式，因为它不需要理解 LoadPlan 的工作原理。在加载过程中需要保持对原始 state_dict 的引用，因此我们需要能够在原地进行操作。
>>> class RenamePlanner(DefaultLoadPlanner): >>> def set_up_planner(self, state_dict, metadata, is_coordinator): >>> self.original_state_dict = state_dict >>> super().set_up_planner(self, {"foo_" + k: v for k, v in state_dict.items()}, 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)[source]¶
此方法在StorageReader完成将数据加载到
tensor后调用一次。提供的张量与调用
resolve_tensor返回的张量相同。 此方法仅在该LoadPlanner需要在将其复制回state_dict中的张量之前对tensor进行后处理时才需要。张量的内容将遵循其设备同步模型。
- abstract create_local_plan()[source]¶
基于 set_up_planner 提供的 state_dict 和元数据创建一个 LoadPlan。
注意：这在每个排名上都会被调用。
- Return type
- abstract load_bytes(read_item, value)[source]¶
加载由
read_item``and ``value描述的项。此方法预计将就地修改底层 state_dict。
value的内容由用于生成正在加载的检查点的 SavePlanner 定义。
- class torch.distributed.checkpoint.LoadPlan(items: List[torch.distributed.checkpoint.planner.ReadItem], storage_data: Any = None, planner_data: Any = None)[source]¶
- 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)[source]¶
- class torch.distributed.checkpoint.SavePlanner[source]¶
抽象类，定义了 save_state_dict 使用的协议，以规划保存过程。
SavePlanner 是一种有状态的对象，可用于自定义整个保存过程。
SavePlanner 作为状态字典的访问代理，因此对其所做的任何变换都会在整个过程中可见。
在调用 save_state_dict 期间，计划子类可以预期以下调用顺序：
- set_up_planner - called on all ranks.
标志着检查点保存的开始。
- create_local_plan - called on all ranks.
处理 state_dict 并生成一个SavePlan，该值将用于全局规划。
- create_global_plan - called on the coordinator rank only.
从所有 ranks 中获取 SavePlan 并做出任何全局决策。
- finish_plan - called on all ranks.
这为每个排名有机会调整全局规划决策。
- resolve_data - called multiple times on each rank
在存储层写入时查找state_dict处的值。
建议用户扩展 DefaultSavePlanner 而不是直接扩展此接口，因为大多数更改都可以通过单个方法的修改来表达。
有三种常见的扩展模式：
重写 state_dict。这是扩展保存过程的最简单方式，因为它不需要理解 SavePlan 的工作原理：
>>> class RenamePlanner(DefaultSavePlanner): >>> def set_up_planner(self, state_dict, is_coordinator): >>> # prefix all keys with `foo_`` >>> super().set_up_planner({"foo_" + k: v for k, v in state_dict.items()}, 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_local_plan()[source]¶
计算当前 rank 的保存计划。 这将被聚合并传递给 create_global_plan。 特定于 Planner 的数据可以通过 SavePlan::planner_data 传递。
这在所有排名上都被调用。
- Return type
- abstract finish_plan(new_plan)[source]¶
合并由create_local_plan创建的计划和create_global_plan的结果。
这在所有排名上都被调用。
- Return type
- class torch.distributed.checkpoint.SavePlan(items: List[torch.distributed.checkpoint.planner.WriteItem], storage_data: Any = None, planner_data: Any = None)[source]¶
- class torch.distributed.checkpoint.WriteItem(index: torch.distributed.checkpoint.metadata.MetadataIndex, type: torch.distributed.checkpoint.planner.WriteItemType, tensor_data: Union[torch.distributed.checkpoint.planner.TensorWriteData, NoneType] = None)[source]¶
我们提供一种基于文件系统的存储层：
- class torch.distributed.checkpoint.FileSystemWriter(path, single_file_per_rank=True, sync_files=True, thread_count=1, per_thread_copy_ahead=10000000)[source]¶
使用文件 IO 实现的 StorageWriter 基本实现。
此实现做出了以下假设和简化：
检查点路径是一个空目录或不存在的目录。
文件创建是原子操作
检查点由每个写入请求对应的一个文件加上一个.metadata文件组成，该文件包含序列化的元数据。
我们提供了LoadPlanner和SavePlanner的默认实现， 可以处理所有torch.distributed构造，例如FSDP、DDP、ShardedTensor和DistributedTensor。
- class torch.distributed.checkpoint.DefaultSavePlanner(flatten_state_dict=True, flatten_sharded_tensors=True, dedup_replicated_tensors=True)[source]¶