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 参数，以防止进程进入他人的设备。如果map_location 缺失，torch.load 将首先将模块加载到 CPU，然后将每个参数复制到保存位置，这将导致同一台机器上的所有进程使用同一组设备.

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

为避免该问题，您应该为map_location 或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")