Pytorch分布式数据并行混淆

Question

我只是看看DDP教程：

https://pytorch.org/tutorials/intermediate/ddp_tutorial.html

根据这一点：

在训练和从检查点恢复过程中通常使用torch.save和torch.load来检查检查点模块。有关详细信息，请参阅保存和加载模型。在使用DDP时，一种优化方法是只将模型保存在一个进程中，然后将其加载到所有进程，从而减少写入开销。这是正确的，因为所有进程都从相同的参数开始，梯度在反向传递中是同步的，因此优化器应该将参数设置为相同的值。如果使用此优化，请确保在保存完成之前，所有进程都不会开始加载。此外，在加载模块时，需要提供一个适当的map_location参数，以防止进程进入其他设备。如果缺少map_location，torch.load将首先将模块加载到CPU，然后将每个参数复制到保存参数的位置，这将导致使用相同的设备集在同一台计算机上的所有进程。关于更高级的故障恢复和弹性支持，请参考TorchElastic。

我不明白这意味着什么。不应该只有一个进程/第一个GPU保存模型吗？保存和加载是如何在进程/GPU之间共享权重的？

Answer 1

当您使用DistributedDataParallel时，您可以跨多个设备拥有相同的模型，这些设备被同步以具有完全相同的参数。

使用DDP时，一种优化方法是只将模型保存在一个进程中，然后将其加载到所有进程，从而减少写入开销。

因为它们是相同的，所以没有必要从所有进程中保存模型，因为它只需要多次写入相同的参数。例如，当您有4个进程/GPU时，您将编写相同的文件4次而不是一次。只有将其从主进程中保存起来，才能避免这种情况。

这是对拯救模型的一个优化。如果您在保存该模型后立即加载它，则需要更加小心。

如果使用此优化，请确保所有进程都不会在保存完成之前开始加载。

如果您只将其保存在一个进程中，则该进程将需要时间来编写该文件。同时，所有其他进程都在继续，它们可能会在文件完全写入磁盘之前加载该文件，这可能会导致各种意外行为或失败，无论该文件是否还不存在，您正在尝试读取一个不完整的文件，或者加载模型的旧版本(如果您覆盖相同的文件)。

此外，在加载模块时，需要提供一个适当的

参数，以防止进程进入其他设备。如果缺少map_location，torch.load将首先将模块加载到CPU，然后将每个参数复制到保存参数的位置，这将导致使用相同的设备集在同一台计算机上的所有进程。

当保存参数(或任何张量)时，PyTorch包括存储它的设备。假设您从使用GPU 0 (device = "cuda:0")的进程中保存它，该信息将被保存，当您加载它时，参数将自动放置到该设备上。但是，如果在使用GPU 1 (device = "cuda:1")的过程中加载它，就会错误地将它们加载到"cuda:0"中。现在，不再使用多个GPU，而是在一个GPU中多次使用相同的模型。最有可能的是，您将耗尽内存，但即使没有，您也不会再使用其他GPU了。

为了避免这个问题，您应该为map_location of torch.load设置适当的设备。

torch.load(PATH, map_location="cuda:1")

# Or load it on the CPU and later use .to(device) on the model
torch.load(PATH, map_location="cpu")