分布式检查点（DCP）入门

创建于：2023年10月02日 | 最后更新：2024年10月30日 | 最后验证：2024年11月05日

作者: Iris Zhang, Rodrigo Kumpera, Chien-Chin Huang, Lucas Pasqualin

注意

在github上查看和编辑本教程。

先决条件：

• FullyShardedDataParallel API 文档

• torch.load API 文档

在分布式训练期间检查点AI模型可能具有挑战性，因为参数和梯度分布在训练器之间，并且在恢复训练时可用训练器的数量可能会发生变化。 Pytorch分布式检查点（DCP）可以帮助简化这一过程。

在本教程中，我们展示了如何使用DCP API与一个简单的FSDP封装模型。

DCP 的工作原理

torch.distributed.checkpoint() 允许并行地从多个等级保存和加载模型。您可以使用此模块在任意数量的等级上并行保存，然后在加载时重新分片到不同的集群拓扑中。

此外，通过使用torch.distributed.checkpoint.state_dict()中的模块， DCP提供了对在分布式设置中优雅处理state_dict生成和加载的支持。 这包括管理模型和优化器之间的完全限定名称（FQN）映射，并为PyTorch提供的并行性设置默认参数。

DCP 与 torch.save() 和 torch.load() 在一些重要方面有所不同：

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

• 它在原地操作，意味着模型应该首先分配其数据，DCP使用该存储空间。

• DCP 提供了对 Stateful 对象（在 torch.distributed.checkpoint.stateful 中正式定义）的特殊处理，如果定义了 state_dict 和 load_state_dict 方法，则会自动调用它们。

注意

本教程中的代码在8-GPU服务器上运行，但可以轻松推广到其他环境。

如何使用DCP

这里我们使用一个用FSDP包装的玩具模型进行演示。同样地，这些API和逻辑可以应用于更大的模型以进行检查点保存。

保存

现在，让我们创建一个玩具模块，用FSDP包装它，给它提供一些虚拟输入数据，并保存它。

import os

import torch
import torch.distributed as dist
import torch.distributed.checkpoint as dcp
import torch.multiprocessing as mp
import torch.nn as nn

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
from torch.distributed.checkpoint.stateful import Stateful
from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

CHECKPOINT_DIR = "checkpoint"


class AppState(Stateful):
    """This is a useful wrapper for checkpointing the Application State. Since this object is compliant
    with the Stateful protocol, DCP will automatically call state_dict/load_stat_dict as needed in the
    dcp.save/load APIs.

    Note: We take advantage of this wrapper to hande calling distributed state dict methods on the model
    and optimizer.
    """

    def __init__(self, model, optimizer=None):
        self.model = model
        self.optimizer = optimizer

    def state_dict(self):
        # this line automatically manages FSDP FQN's, as well as sets the default state dict type to FSDP.SHARDED_STATE_DICT
        model_state_dict, optimizer_state_dict = get_state_dict(self.model, self.optimizer)
        return {
            "model": model_state_dict,
            "optim": optimizer_state_dict
        }

    def load_state_dict(self, state_dict):
        # sets our state dicts on the model and optimizer, now that we've loaded
        set_state_dict(
            self.model,
            self.optimizer,
            model_state_dict=state_dict["model"],
            optim_state_dict=state_dict["optim"]
        )

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(16, 16)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(16, 8)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))


def setup(rank, world_size):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355 "

    # initialize the process group
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)


def cleanup():
    dist.destroy_process_group()


def run_fsdp_checkpoint_save_example(rank, world_size):
    print(f"Running basic FSDP checkpoint saving example on rank {rank}.")
    setup(rank, world_size)

    # create a model and move it to GPU with id rank
    model = ToyModel().to(rank)
    model = FSDP(model)

    loss_fn = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

    optimizer.zero_grad()
    model(torch.rand(8, 16, device="cuda")).sum().backward()
    optimizer.step()

    state_dict = { "app": AppState(model, optimizer) }
    dcp.save(state_dict, checkpoint_id=CHECKPOINT_DIR)

    cleanup()


if __name__ == "__main__":
    world_size = torch.cuda.device_count()
    print(f"Running fsdp checkpoint example on {world_size} devices.")
    mp.spawn(
        run_fsdp_checkpoint_save_example,
        args=(world_size,),
        nprocs=world_size,
        join=True,
    )

请继续检查checkpoint目录。您应该会看到8个检查点文件，如下所示。

加载

保存后，让我们创建相同的FSDP封装模型，并将保存的状态字典从存储加载到模型中。您可以在相同的世界大小或不同的世界大小中加载。

请注意，在加载之前，您需要调用model.state_dict()并将其传递给DCP的load_state_dict() API。 这与torch.load()有根本的不同，因为torch.load()只需要检查点的路径即可加载。 我们需要在加载之前获取state_dict的原因是：

• DCP 使用模型 state_dict 中预分配的存储从检查点目录加载。在加载过程中，传入的 state_dict 将会被就地更新。

• DCP 在加载之前需要从模型中获取分片信息以支持重新分片。

import os

import torch
import torch.distributed as dist
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.stateful import Stateful
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
import torch.multiprocessing as mp
import torch.nn as nn

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

CHECKPOINT_DIR = "checkpoint"


class AppState(Stateful):
    """This is a useful wrapper for checkpointing the Application State. Since this object is compliant
    with the Stateful protocol, DCP will automatically call state_dict/load_stat_dict as needed in the
    dcp.save/load APIs.

    Note: We take advantage of this wrapper to hande calling distributed state dict methods on the model
    and optimizer.
    """

    def __init__(self, model, optimizer=None):
        self.model = model
        self.optimizer = optimizer

    def state_dict(self):
        # this line automatically manages FSDP FQN's, as well as sets the default state dict type to FSDP.SHARDED_STATE_DICT
        model_state_dict, optimizer_state_dict = get_state_dict(self.model, self.optimizer)
        return {
            "model": model_state_dict,
            "optim": optimizer_state_dict
        }

    def load_state_dict(self, state_dict):
        # sets our state dicts on the model and optimizer, now that we've loaded
        set_state_dict(
            self.model,
            self.optimizer,
            model_state_dict=state_dict["model"],
            optim_state_dict=state_dict["optim"]
        )

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(16, 16)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(16, 8)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))


def setup(rank, world_size):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355 "

    # initialize the process group
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)


def cleanup():
    dist.destroy_process_group()


def run_fsdp_checkpoint_load_example(rank, world_size):
    print(f"Running basic FSDP checkpoint loading example on rank {rank}.")
    setup(rank, world_size)

    # create a model and move it to GPU with id rank
    model = ToyModel().to(rank)
    model = FSDP(model)

    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

    state_dict = { "app": AppState(model, optimizer)}
    dcp.load(
        state_dict=state_dict,
        checkpoint_id=CHECKPOINT_DIR,
    )

    cleanup()


if __name__ == "__main__":
    world_size = torch.cuda.device_count()
    print(f"Running fsdp checkpoint example on {world_size} devices.")
    mp.spawn(
        run_fsdp_checkpoint_load_example,
        args=(world_size,),
        nprocs=world_size,
        join=True,
    )

如果您希望将保存的检查点加载到非分布式设置中的非FSDP包装模型中，可能是为了推理，您也可以使用DCP来实现。 默认情况下，DCP以单程序多数据（SPMD）风格保存和加载分布式state_dict。然而，如果没有初始化进程组，DCP会推断意图是以“非分布式”风格保存或加载，这意味着完全在当前进程中完成。

注意

多程序多数据的分布式检查点支持仍在开发中。

import os

import torch
import torch.distributed.checkpoint as dcp
import torch.nn as nn


CHECKPOINT_DIR = "checkpoint"


class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(16, 16)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(16, 8)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))


def run_checkpoint_load_example():
    # create the non FSDP-wrapped toy model
    model = ToyModel()
    state_dict = {
        "model": model.state_dict(),
    }

    # since no progress group is initialized, DCP will disable any collectives.
    dcp.load(
        state_dict=state_dict,
        checkpoint_id=CHECKPOINT_DIR,
    )
    model.load_state_dict(state_dict["model"])

if __name__ == "__main__":
    print(f"Running basic DCP checkpoint loading example.")
    run_checkpoint_load_example()

格式

尚未提到的一个缺点是，DCP以一种与使用torch.save生成的格式本质上不同的格式保存检查点。由于当用户希望与习惯于torch.save格式的用户共享模型时，或者通常只是希望为他们的应用程序增加格式灵活性时，这可能会成为一个问题。对于这种情况，我们在torch.distributed.checkpoint.format_utils中提供了format_utils模块。

为了方便用户，提供了一个命令行实用程序，其格式如下：

python -m torch.distributed.checkpoint.format_utils <mode> <checkpoint location> <location to write formats to>

在上面的命令中，mode 是 torch_to_dcp 或 dcp_to_torch 之一。

另外，还为希望直接转换检查点的用户提供了方法。

import os

import torch
import torch.distributed.checkpoint as DCP
from torch.distributed.checkpoint.format_utils import dcp_to_torch_save, torch_save_to_dcp

CHECKPOINT_DIR = "checkpoint"
TORCH_SAVE_CHECKPOINT_DIR = "torch_save_checkpoint.pth"

# convert dcp model to torch.save (assumes checkpoint was generated as above)
dcp_to_torch_save(CHECKPOINT_DIR, TORCH_SAVE_CHECKPOINT_DIR)

# converts the torch.save model back to DCP
dcp_to_torch_save(TORCH_SAVE_CHECKPOINT_DIR, f"{CHECKPOINT_DIR}_new")

结论

总之，我们已经学会了如何使用DCP的save()和load() API，以及它们与torch.save()和torch.load()的不同之处。 此外，我们还学会了如何使用get_state_dict()和set_state_dict()在状态字典生成和加载过程中自动管理并行性特定的FQN和默认值。

欲了解更多信息，请参阅以下内容：

• Saving and loading models tutorial

• FullyShardedDataParallel 教程入门