分布式检查点 - torch.distributed.checkpoint¶
分布式检查点 （DCP） 支持并行加载和保存来自多个 rank 的模型。 它处理加载时重新分片，从而支持在一个集群拓扑中保存并加载到另一个集群拓扑中。
DCP 与 torch.save 和 torch.load 在几个重要方面有所不同：
它为每个检查点生成多个文件，每个等级至少生成一个文件。
它在原地运行，这意味着模型应该首先分配其数据，然后 DCP 使用该存储。
加载和保存检查点的入口点如下：
- torch.distributed.checkpoint 中。load_state_dict（state_dict， storage_reader， process_group=无， coordinator_rank=0， no_dist=False， 规划者=无）[来源]¶
以 SPMD 样式加载 distributed。
state_dict每个等级将尝试读取所需的最少数据量 完成请求的state_dict。加载实例时，每个排名仅读取其本地分片的数据。
ShardedTensor警告
中的所有张量都必须在其 destination device 之前。
state_dict所有非张量数据都使用 torch.load（） 加载并就地修改 在 state_dict。
警告
用户必须在根模块上调用 load_state_dict 以确保负载 pos-processing 和非 Tensor 数据正确传播。
- 参数
state_dict （Dict[str， Any]） – 要加载state_dict。请注意，此 state dict 将就地更新。
storage_reader （StorageReader） – 用于从中加载数据的 StorageReader。
process_group （ProcessGroup） – 用于跨等级同步的 ProcessGroup。
coordinator_rank （int） – 用于协调检查点的排名。 默认情况下使用 rank0。
no_dist （bool） – 如果 ，则分布式检查点不会保存 在 SPMD 样式中。（默认：
TrueFalse)
- 结果
没有。
- 返回类型：
没有
- 例子
>>> my_model = MyModule() >>> optimizer = Adagrad(my_model.parameters()) >>> model_state_dict = my_model.state_dict() >>> fs_storage_loader = torch.distributed.checkpoint.FileSystemLoader("/checkpoint/1")
>>> torch.distributed.checkpoint.load_state_dict( >>> state_dict=model_state_dict, >>> storage_reader=fs_storage_loader, >>> )
>>> # 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 使用 collectives 来协调跨等级的读取。 对于基于 NCCL 的进程组， 在进行通信之前，必须将对象移动到 GPU 设备。 在这种情况下，使用的设备由 提供，用户有责任确保将其设置为每个 rank 具有单个 GPU，通过 .
torch.cuda.current_device()torch.cuda.set_device()
- torch.distributed.checkpoint 中。save_state_dict（state_dict， storage_writer， process_group=无， coordinator_rank=0， no_dist=False， 规划者=无）[来源]¶
以 SPMD 样式保存分布式模型。
此函数与它处理的不同之处在于，每个排名仅保存其本地分片。
torch.save()ShardedTensor警告
无法保证 PyTorch 版本之间的向后兼容性 用于节省state_dicts。
警告
如果使用 process_group 参数，请确保只有其 ranks 调用 save_state_dict，并且 state_dict 中的所有数据都属于它。
注意
此函数可用于保存具有初始进程的 state_dict group 通过传递 .这可用于生成 checkpoint 可以被 load_state_dict 消费的是 SPMD 时尚。
no_dist=True- 参数
state_dict （Dict[str， Any]） – 一个state_dict
storage_writer （StorageWriter） – StorageWrite 实例用于执行写入。
process_group （ProcessGroup） – 用于跨等级同步的 ProcessGroup。
coordinator_rank （int） – 用于协调检查点的排名。 默认情况下使用 rank0。
no_dist （bool） – 如果 ，则分布式检查点不会保存 在 SPMD 样式中。（默认：
TrueFalse)
- 结果
Metadata 对象。
- 返回类型：
元数据
例
>>> 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_stroage_writer, >>> )
注意
save_state_dict 使用 collectives 来协调跨等级的写入。 对于基于 NCCL 的进程组， 在进行通信之前，必须将对象移动到 GPU 设备。 在这种情况下，使用的设备由 提供，用户有责任确保将其设置为： 每个等级都有一个单独的 GPU，通过 .
torch.cuda.current_device()torch.cuda.set_device()
以下类型定义了 checkpoint 期间使用的 IO 接口：
- 类 torch.distributed.checkpoint 中。StorageReader[来源]¶
用于从存储中读取的接口。
load_state_dict一个 StorageReader 实例同时充当协调器和追随者 在分布式检查点中。作为初始化的一部分，每个实例 被告知其角色。
子类应期望以下调用序列：
load_state_dict（所有等级） read_metadata（）
（所有等级） set_up_storage_reader（）
（所有等级） prepare_local_plan（）
（协调） prepare_global_plan（）
（所有等级） read_data（）
- 摘要 prepare_global_plan（计划）[来源]¶
执行存储加载的集中规划。
此方法仅在协调器实例上调用。
虽然这种方法可以产生完全不同的计划，但首选的 方式是将存储特定数据存储在 LoadPlan：：storage_data 中。
- 摘要 prepare_local_plan（计划）[来源]¶
执行特定于存储的本地规划。
虽然此方法可以生成完全不同的计划，但建议的 方式是将存储特定数据存储在 LoadPlan：：storage_data 中。
- 参数
plan （LoadPlan） – 正在使用的本地计划。
LoadPlan- 结果
A 仓储后本地规划
LoadPlan- 返回类型：
- 摘要 read_data（Plan， Planner）[来源]¶
从 using 中读取所有项以解析数据。
planplanner应调用子类以反序列化 BytesIO object 移动到正确的位置。
LoadPlanner::load_bytes子类应该调用以获取对 应该将数据加载到的张量。
LoadPlanner::resolve_tensorStorageLayer 负责正确安排任何跨设备副本 必填。
- 参数
plan （LoadPlan） – 要执行的本地计划
planner （LoadPlanner） – 用于解析项目的 planner 对象。
- 结果
一个 future ，在完成所有读取后完成。
- 返回类型：
Future[无]
- 类 torch.distributed.checkpoint 中。StorageWriter[来源]¶
用于写入存储的接口。
save_state_dict一个 StorageWriter 实例同时充当协调器和追随者 在分布式检查点中。作为初始化的一部分，每个实例 被告知其角色。
子类应期望以下调用序列。
（所有等级） set_up_storage_writer（）
（所有等级） prepare_local_plan（）
（协调） prepare_global_plan（）
（所有等级） write_data（）
（协调器） finish（）
- 摘要 prepare_global_plan（计划）[来源]¶
执行集中存储规划。
此方法仅在协调器实例上调用。
虽然这种方法可以产生完全不同的计划，但首选的 方法是将存储特定数据存储在 SavePlan：：storage_data 中。
- 摘要 prepare_local_plan（计划）[来源]¶
执行特定于存储的本地规划。
虽然此方法可以生成完全不同的计划，但建议的 方法是将存储特定数据存储在 SavePlan：：storage_data 中。
- 摘要 write_data（Plan， Planner）[来源]¶
写入 using 中的所有项目以解析数据。
planplanner子类应该调用每个项目 从 plan 获取对要写入的基础对象的访问权。
SavePlanner::resolve_data子类应该延迟调用 resolve_data因为它可以分配内存。 如果是张量，请做出以下假设：
它们可能在任何设备上，包括不匹配的
WriteItem::tensor_data它们可能是视图，也可能不是连续的。只需保存投影。
- 参数
plan （SavePlan） – 要执行的保存计划。
planner （SavePlanner） – 用于将项目解析为数据的 Planner 对象。
- 结果
完成 WriteResult 列表的 future
- 返回类型：
以下类型定义了 checkpoint 期间使用的 planner 接口：
- 类 torch.distributed.checkpoint 中。LoadPlanner[来源]¶
定义 load_state_dict 用于规划加载过程的协议的抽象类。
LoadPlanner 是可用于自定义整个加载过程的有状态对象。
LoadPlanner 充当 state_dict 的访问代理，因此对它执行的任何转换 将对整个过程可见。
在 load_state_dict期间，Planner 子类可以预期以下 Sequences 调用：
- set_up_planner - 召集所有级别。
表示开始加载检查点。
- create_local_plan - 召集所有等级。
处理 state_dict 并生成 LoadPlan，该 LoadPlan 将发送以进行全局规划。
- create_global_plan - 仅根据协调者级别调用。
从所有级别获取 LoadPlan 并做出任何全局决策。
- load_bytes - 在每个等级上多次调用
在 state_dict 中，每个非张量值调用一次。
- resolve_tensor 和 commit_tensor - 在每个等级上多次调用
它们对 state_dict 中的每个 Tensor 值成对调用。
建议用户直接扩展 DefaultLoadPlanner 而不是这个接口 大多数更改可以通过单个方法中的更改来表示。
有两种常见的扩展模式：
重写state_dict。这是扩展 load 进程的最简单方法，因为它 并没有完全理解 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
- 摘要 commit_tensor（read_item， Tensor）[来源]¶
一旦 StorageReader 完成将数据加载到 中，就会调用此方法。
tensor提供的张量与调用 . 仅当此 LoadPlanner 需要在 将其复制回 state_dict 中的那个。
resolve_tensortensortensor 的内容将遵循其设备同步模型。
- 摘要 load_bytes（read_item， value）[来源]¶
加载 描述的项目。
read_item``and ``value此方法应就地修改基础state_dict。
的内容由用于生成 正在加载的 checkpoint。
value
- 类 torch.distributed.checkpoint 中。LoadPlan（items： List[torch.distributed.checkpoint.planner.ReadItem]， storage_data： Any = None， planner_data：任何 = 无）[来源]¶
- 类 torch.distributed.checkpoint 中。ReadItem（类型：torch.distributed.checkpoint.planner.LoadItemType，dest_index：torch.distributed.checkpoint.metadata.MetadataIndex，dest_offsets：torch.大小，storage_index：torch.distributed.checkpoint.metadata.MetadataIndex，storage_offsets：torch。尺寸、长度：Torch。尺寸）[来源]¶
- 类 torch.distributed.checkpoint 中。SavePlanner[来源]¶
定义 save_state_dict 用于规划保存过程的协议的抽象类。
SavePlanners 是可用于自定义整个保存过程的有状态对象。
SavePlanner 充当 state_dict 的访问代理，因此对它所做的任何转换 将对整个过程可见。
在 save_state_dict期间，Planner 子类可以预期以下 Sequences 调用：
- set_up_planner - 召集所有级别。
指示检查点保存开始。
- create_local_plan - 召集所有等级。
处理state_dict并生成 SavePlan，该 SavePlan 将发送用于全球规划。
- create_global_plan - 仅根据协调者级别调用。
从所有级别中获取 SavePlan 并做出任何全局决策。
- finish_plan - 召集所有军衔。
这使每个等级都有机会适应全局规划决策。
- resolve_data - 在每个等级上多次调用
在 state_dict 上查找要写入的存储层的值。
建议用户将 DefaultSavePlanner 而不是这个接口直接扩展为 大多数更改可以通过单个方法中的更改来表示。
有 3 种常见的扩展模式：
重写state_dict。这是扩展保存过程的最简单方法，因为它 并不完全理解 SavePlan 工作原理的内在含义：
>>> class RenamePlanner(DefaultSavePlanner): >>> def set_up_planner(self, state_dict, is_coordinator): >>> # prefix all keys with `foo_`` >>> super().set_up_planner(self, {"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 >>> >>> 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)
最后，一些 planner 需要在 checkpoint 中保存额外的元数据，这是 通过让每个 rank 在本地计划中贡献其数据项来完成，并且 全局规划器会聚合它们：
>>> 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
- 摘要 create_local_plan（）[来源]¶
计算当前排名的保存计划。 这将被聚合并传递给 create_global_plan。 Planner 特定的数据可以通过 SavePlan：:p lanner_data 传递。
这在所有等级上都调用。
- 返回类型：
- 类 torch.distributed.checkpoint 中。SavePlan（items： List[torch.distributed.checkpoint.planner.WriteItem]， storage_data： Any = None， planner_data：任何 = 无）[来源]¶
- 类 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）[来源]¶
我们提供基于文件系统的存储层：
- 类 torch.distributed.checkpoint 中。FileSystemWriter（路径， single_file_per_rank=True， sync_files=True， thread_count=1， per_thread_copy_ahead=10000000 元）[来源]¶
使用文件 IO 的 StorageWriter 的基本实现。
此实现进行了以下假设和简化：
检查点路径是一个空目录或不存在的目录。
文件创建是原子的
检查点由每个写入请求一个文件以及 包含序列化元数据的 .metadata 文件。
我们提供了 LoadPlanner 和 SavePlanner 的默认实现，这些 可以处理所有 torch.distributed 结构，例如 FSDP、DDP、ShardedTensor 和 DistributedTensor。
- 类 torch.distributed.checkpoint 中。DefaultSavePlanner（flatten_state_dict=True， flatten_sharded_tensors=True， dedup_replicated_tensors=True）[来源]¶