• 文档 >
  • 多进程最佳实践
快捷方式

多进程最佳实践

torch.multiprocessing 是 Python 标准库 multiprocessing 模块的直接替代品。它支持完全相同的操作,并对其进行了扩展,使得所有通过 multiprocessing.Queue 发送的 tensor 的数据都会被移动到共享内存中,并仅将一个句柄发送给另一个进程。

注意

当一个 Tensor 被发送到另一个进程时,该 Tensor 的数据是共享的。如果 torch.Tensor.grad 不是 None,它也会被共享。一个没有 torch.Tensor.grad 字段的 Tensor 被发送到其他进程后,会在接收进程中创建一个标准的、进程专属的 .grad Tensor,与 Tensor 数据已共享的方式不同,这个 .grad Tensor 不会自动在所有进程中共享。

这使得实现各种训练方法成为可能,例如 Hogwild、A3C 或任何其他需要异步操作的方法。

多进程中的 CUDA

CUDA 运行时不支持 fork 启动方法;在子进程中使用 CUDA 需要使用 spawnforkserver 启动方法。

注意

启动方法可以通过创建上下文 (`multiprocessing.get_context(...)`) 或直接使用 multiprocessing.set_start_method(...) 来设置。

与 CPU tensor 不同,只要接收进程保留了该 tensor 的副本,发送进程就需要保留原始 tensor。这一点在底层已经实现,但要求用户遵循最佳实践以确保程序正确运行。例如,只要消费者进程持有 tensor 的引用,发送进程就必须保持存活;如果消费者进程因致命信号异常退出,引用计数机制无法拯救你。参见 本节

另请参阅:使用 nn.parallel.DistributedDataParallel 而非 multiprocessing 或 nn.DataParallel

最佳实践和技巧

避免和解决死锁

当创建新进程时,有很多事情可能会出错,其中最常见的死锁原因是后台线程。如果存在任何持有锁或导入了模块的线程,并且调用了 fork,那么子进程很可能处于损坏状态并导致死锁或以其他方式失败。请注意,即使你没有这样做,Python 的内置库也会——只需看看 multiprocessing 就知道了。 multiprocessing.Queue 实际上是一个非常复杂的类,它会生成多个用于序列化、发送和接收对象的线程,这些线程也可能导致上述问题。如果你发现自己处于这种情况,请尝试使用不使用任何额外线程的 SimpleQueue

我们正尽力让你轻松使用并确保这些死锁不会发生,但有些事情超出了我们的控制范围。如果你遇到暂时无法解决的问题,请尝试在论坛上求助,我们会看看是否是我们能够修复的问题。

复用通过 Queue 传递的缓冲区

请记住,每次将 Tensor 放入 multiprocessing.Queue 时,都需要将其移动到共享内存。如果它已经是共享的,则无需操作;否则会产生额外的内存复制,这可能会减慢整个进程。即使你有一个进程池向单个进程发送数据,也要让该进程将缓冲区发回——这几乎是免费的,并且可以在发送下一个批次时避免一次复制。

异步多进程训练(例如 Hogwild)

使用 torch.multiprocessing,可以异步训练模型,参数要么始终共享,要么定期同步。在前一种情况下,我们建议传递整个模型对象;在后一种情况下,我们建议只发送 state_dict()

我们建议使用 multiprocessing.Queue 在进程之间传递各种 PyTorch 对象。例如,在使用 fork 启动方法时,可以继承共享内存中已有的 tensor 和 storage,但这非常容易出错,应谨慎使用,且仅限高级用户。Queue,尽管有时不是那么优雅的解决方案,但在所有情况下都能正常工作。

警告

你应该小心那些没有使用 if __name__ == '__main__' 保护的全局语句。如果使用了 fork 以外的启动方法,它们会在所有子进程中执行。

Hogwild

可以在 examples 仓库中找到一个具体的 Hogwild 实现,但为了展示代码的总体结构,下面也有一个最小示例

import torch.multiprocessing as mp
from model import MyModel

def train(model):
    # Construct data_loader, optimizer, etc.
    for data, labels in data_loader:
        optimizer.zero_grad()
        loss_fn(model(data), labels).backward()
        optimizer.step()  # This will update the shared parameters

if __name__ == '__main__':
    num_processes = 4
    model = MyModel()
    # NOTE: this is required for the ``fork`` method to work
    model.share_memory()
    processes = []
    for rank in range(num_processes):
        p = mp.Process(target=train, args=(model,))
        p.start()
        processes.append(p)
    for p in processes:
        p.join()

多进程中的 CPU

不恰当的多进程使用可能导致 CPU 过度占用(oversubscription),使得不同进程竞争 CPU 资源,从而导致效率低下。

本教程将解释 CPU 过度占用是什么以及如何避免它。

CPU 过度占用

CPU 过度占用是一个技术术语,指系统中分配的 vCPU 总数超过硬件上可用的 vCPU 总数的情况。

这会导致对 CPU 资源的严重争夺。在这种情况下,进程之间会频繁切换,这增加了进程切换开销并降低了整体系统效率。

参考 example 仓库中 Hogwild 实现的代码示例来了解 CPU 过度占用。

在 CPU 上使用 4 个进程运行以下训练示例命令时

python main.py --num-processes 4

假设机器上有 N 个 vCPU 可用,执行上述命令将生成 4 个子进程。每个子进程会为自己分配 N 个 vCPU,总共需要 4*N 个 vCPU。然而,机器上只有 N 个 vCPU 可用。因此,不同进程会竞争资源,导致频繁的进程切换。

以下观察结果表明存在 CPU 过度占用

  1. 高 CPU 利用率:通过使用 htop 命令,你可以观察到 CPU 利用率持续很高,经常达到或超过其最大容量。这表明对 CPU 资源的需求超过了可用的物理核心数,导致进程之间争夺 CPU 时间。

  2. 频繁的上下文切换和低系统效率:在 CPU 过度占用的情况下,进程争夺 CPU 时间,操作系统需要快速在不同进程之间切换以公平分配资源。这种频繁的上下文切换增加了开销,降低了整体系统效率。

避免 CPU 过度占用

避免 CPU 过度占用的一个好方法是进行适当的资源分配。确保同时运行的进程或线程数量不超过可用的 CPU 资源。

在这种情况下,一个解决方案是在子进程中指定适当的线程数。这可以通过在子进程中使用 torch.set_num_threads(int) 函数来设置每个进程的线程数实现。

假设机器上有 N 个 vCPU,将生成 M 个进程,则每个进程使用的最大 num_threads 值为 floor(N/M)。为了避免 mnist_hogwild 示例中的 CPU 过度占用,需要对 example 仓库中的 train.py 文件进行以下更改。

def train(rank, args, model, device, dataset, dataloader_kwargs):
    torch.manual_seed(args.seed + rank)

    #### define the num threads used in current sub-processes
    torch.set_num_threads(floor(N/M))

    train_loader = torch.utils.data.DataLoader(dataset, **dataloader_kwargs)

    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)
    for epoch in range(1, args.epochs + 1):
        train_epoch(epoch, args, model, device, train_loader, optimizer)

使用 torch.set_num_threads(floor(N/M)) 为每个进程设置 num_thread。其中 N 替换为可用的 vCPU 数量,M 替换为选择的进程数量。合适的 num_thread 值会根据具体任务而有所不同。然而,作为一般准则,num_thread 的最大值应为 floor(N/M) 以避免 CPU 过度占用。在 mnist_hogwild 训练示例中,避免 CPU 过度占用后,性能可以提升 30 倍。

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深度教程

查看教程

资源

查找开发资源并获得解答

查看资源