分布式检查点

PyTorch/XLA SPMD 通过专用的 Planner 实例兼容 torch.distributed.checkpoint 库。用户可以通过此通用接口同步保存和加载检查点。

SPMDSavePlanner 和 SPMDLoadPlanner（源码）类使 save 和 load 函数可以直接在 XLAShardedTensor 的分片上操作，从而在 SPMD 训练中实现分布式检查点的所有优势。

以下是同步分布式检查点 API 的演示

import torch.distributed.checkpoint as dist_cp
import torch_xla.experimental.distributed_checkpoint as xc

# Saving a state_dict
state_dict = {
    "model": model.state_dict(),
    "optim": optim.state_dict(),
}

dist_cp.save(
    state_dict=state_dict,
    storage_writer=dist_cp.FileSystemWriter(CHECKPOINT_DIR),
    planner=xc.SPMDSavePlanner(),
)
...

# Loading the model's state_dict from the checkpoint. The model should
# already be on the XLA device and have the desired sharding applied.
state_dict = {
    "model": model.state_dict(),
}

dist_cp.load(
    state_dict=state_dict,
    storage_reader=dist_cp.FileSystemReader(CHECKPOINT_DIR),
    planner=xc.SPMDLoadPlanner(),
)
model.load_state_dict(state_dict["model"])

CheckpointManager

实验性的 CheckpointManager 接口提供了 torch.distributed.checkpoint 函数之上的更高级 API，以实现一些关键功能：

• 托管检查点：CheckpointManager 每次保存的检查点都由保存时所在的步骤标识。所有跟踪的步骤都可以通过 CheckpointManager.all_steps 方法访问，并且可以使用 CheckpointManager.restore 恢复任何跟踪的步骤。

• 异步检查点：通过 CheckpointManager.save_async API 保存的检查点会异步写入持久化存储，以避免检查点期间阻塞训练。输入的分片 state_dict 会首先移动到 CPU，然后检查点被分发到后台线程。

• 发生抢占时自动检查点：在 Cloud TPU 上，可以检测到抢占，并在进程终止前保存检查点。要使用此功能，请确保您的 TPU 通过启用了 Autocheckpointing 的 QueuedResource 进行配置，并确保在构造 CheckpointManager 时设置 chkpt_on_preemption 参数（此选项默认启用）。

• FSSpec 支持：CheckpointManager 使用 fsspec 存储后端，可以直接将检查点保存到任何 fsspec 兼容的文件系统，包括 GCS。

CheckpointManager 的示例用法如下：

from torch_xla.experimental.distributed_checkpoint import CheckpointManager, prime_optimizer

# Create a CheckpointManager to checkpoint every 10 steps into GCS.
chkpt_mgr = CheckpointManager('gs://my-bucket/my-experiment', 10)

# Select a checkpoint to restore from, and restore if applicable
tracked_steps = chkpt_mgr.all_steps()
if tracked_steps:
    # Choose the highest step
    best_step = max(tracked_steps)
    # Before restoring the checkpoint, the optimizer state must be primed
    # to allow state to be loaded into it.
    prime_optimizer(optim)
    state_dict = {'model': model.state_dict(), 'optim': optim.state_dict()}
    chkpt_mgr.restore(best_step, state_dict)
    model.load_state_dict(state_dict['model'])
    optim.load_state_dict(state_dict['optim'])

# Call `save` or `save_async` every step within the train loop. These methods
# return True when a checkpoint is taken.
for step, data in enumerate(dataloader):
    ...
    state_dict = {'model': model.state_dict(), 'optim': optim.state_dict()}
    if chkpt_mgr.save_async(step, state_dict):
        print(f'Checkpoint taken at step {step}')

恢复优化器状态

在分布式检查点中，state_dicts 会就地加载，并且仅加载检查点的必需分片。由于优化器状态是惰性创建的，直到第一次调用 optimizer.step 时才存在该状态，因此尝试加载未初始化的优化器将会失败。

为此提供了实用方法 prime_optimizer：它通过将所有梯度设置为零并调用 optimizer.step 来运行一个模拟的训练步骤。这是一个破坏性方法，会修改模型参数和优化器状态，因此应仅在恢复之前调用。

进程组

要使用 torch.distributed API（例如分布式检查点），需要一个进程组。在 SPMD 模式下，不支持 xla 后端，因为编译器负责所有集合通信。

相反，必须使用 CPU 进程组，例如 gloo。在 TPU 上，仍然支持 xla:// init_method 来发现主 IP、全局世界大小和主机秩。以下是一个示例初始化：

import torch.distributed as dist
# Import to register the `xla://` init_method
import torch_xla.distributed.xla_backend
import torch_xla.runtime as xr

xr.use_spmd()

# The `xla://` init_method will automatically discover master worker IP, rank,
# and global world size without requiring environment configuration on TPUs.
dist.init_process_group('gloo', init_method='xla://')