Shortcuts

hfai.nn.parallel

DistributedDataParallel

分布式数据并行工具,封装了 hfreduce

class hfai.nn.parallel.DistributedDataParallel(module, device_ids=None, broadcast_buffers=True, process_group=None, find_unused_parameters=False)[source]

分布式数据并行工具,封装了 hfreduce

使用幻方 AI 自研的 hfreduce 多卡通信工具,可以替换 PyTorch DDP 加速训练。 使用方法与 torch.nn.parallel.DistributedDataParallel 相同。

Parameters
  • module (torch.nn.Module) – PyTorch 模型

  • device_ids (list) – 模型所在的 GPU id,如果是 None 则会用 torch.cuda.current_device() 的返回值,默认是 None

  • broadcast_buffers (bool) – 是否在 forward 之前把 rank-0 上的 buffer 广播到其他 rank 上,默认是 True

  • process_group (ProcessGroup) – ProcessGroup 对象,如果是 None 会用默认的分组

  • find_unused_parameters (bool) – 是否遍历计算图,找到不参与 backward 的参数; 在少数情况下(比如训练 GAN)需要设置成 True;默认是 False

Note

进程退出时可能会报错,可通过退出之前调用 model.reducer.stop() 来解决

Examples:

from hfai.nn.parallel import DistributedDataParallel
model = DistributedDataParallel(model, device_ids=[local_rank])
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# training ...
for epoch in range(epochs):
    for step, (x, y) in enumerate(dataloader):
        # training
        optimizer.zero_grad()
        output = model(x)
        loss_fn(y, output).backward()
        optimizer.step()