Torch 1.9.1 DDP 并行优化与多模块调用问题

DDP 基础实现

由于 DataParallel (DP) 采取的是多线程并行，出于其特性，会造成通信瓶颈 （GIL 限制），因此更高效的方式是使用 DistributedDataParallel 实现更高效的 GPU 使用。DDP 相关基础实现参考此处文章，亲测可以使用。

注：目前使用 1.7+ pytorch nccl 初始化 DDP 会报错，亲测也有该问题，因此建议使用 ‘gloo’，虽然相对速度可能较慢。

多模块调用问题

假设考虑如下问题：

import torch

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

import torch.optim as optim
import torch.nn as nn

import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=-1, type=int)
FLAGS = parser.parse_args()
local_rank = FLAGS.local_rank

torch.cuda.set_device(local_rank)
dist.init_process_group(backend='gloo')

device = torch.device("cuda", local_rank)


class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.device = device

        self.nn1 = nn.Linear(10, 20)
        self.nn2 = nn.Linear(20, 10)

    def forward(self, x):
        mid = self.nn1(x)
        model = self.nn2(mid)

        return model

    def middle(self, x):
        return self.nn1(x)


model = Net().to(device)
model = DDP(model, device_ids=[local_rank], output_device=local_rank)

x = torch.randn(20, 10).to(local_rank)
outputs = model(x)
mid = model.middle(x)

labels = torch.randn(20, 10).to(local_rank)
loss_fn = nn.MSELoss()
loss_fn(outputs, labels).backward()

optimizer = optim.SGD(model.parameters(), lr=0.001)
optimizer.step()

如果没有 model_ddp = DDP(model, device_ids=[local_rank], output_device=local_rank) 的操作，上述程序是个简单例子。注意到此处 Net 模型有一个 intrinsic forwar() 函数 和 其他函数 middle()，在调用 mid = model.middle(x) 会报错，因为 model 转换为 DDP 形式之后只会执行 forward()。

考虑如下解决方式：

DDP.module.func()

查看 DDP 实现源码可以发现，DDP 可以通过模块引用从而调用 其他函数 middle()。将

mid = model.middle(x)

改为

mid = model.module.middle(x)

可以找到 middle() 函数，注意此处 mid 是非 DDP 形式的。

并行非并行分离

类似 StyleGAN2, 我们可以考虑创建并行 model_ddp() 模型和 model() 非并行模型，由于其他函数一般不参与模型学习，而是中间状态输出，因此这种情况下我们可以使用双模型。如下：

model = Net().to(device)
model_ddp = DDP(model, device_ids=[local_rank], output_device=local_rank)

forward() 条件引入 middle

由于问题是 DDP 只会执行 forward(), 因此我们可以考虑将 middle 进行条件引入：

def forward(self, x, run_middle=False):
    if run_middle:
        return self.middle(x)
    mid = self.nn1(x)
    model = self.nn2(mid)

    return model

此为一种最为通用的策略。

上述三种策略根据需求可以结合使用。