在 PyTorch 中如何将一个很大的模型分布到多个 GPU 上训练？

[daa233]

在PyTorch 中，如果我们需要训练一个很大的模型，即使 batch size 很小，放在单个 GPU 上也会显存溢出，而我们又没有大显存的 GPU，这时应该怎么办呢？我们能否将一个大的模型分成几个小的部分，分别在不同的模型上训练呢？

[daa233]

可以的，参考下面的例子：

class MyModel(nn.Module):
    def __init__(self, split_gpus):
        self.large_submodule1 = ...
        self.large_submodule2 = ...

        self.split_gpus = split_gpus
        if split_gpus:
            self.large_submodule1.cuda(0)
            self.large_submodule2.cuda(1)

    def forward(self, x):
        x = self.large_submodule1(x)
        if split_gpus:
            x = x.cuda(1)  # P2P GPU transfer
        return self.large_submodule2(x)

参考：Model parallelism in pytorch for large(r than 1 GPU) models?

[daa233]

PyTorch 新版本给了官方教程：Model Parallel Best Practices 和 Combine DDP with Model Parallelism