Pytorch 分布式训练

区分概念：Machine vs. Device. 多机(Machine) 是指在不同主机上训练。每个主机有自己的 ip 。通过向 pytorch 提供 Master 的 ip:port 同步各个机器。多卡 (Device) 是指在同一台机器上有多张显卡。

分布式训练

采用 .to(cuda:%d) 方法指定GPU。eg.

cuda1 = torch.device("cuda:1")
data = data.to(cuda1)

Overview

三种分布式训练方法：

• DDP(Distributed Data Parallel，使用最多)

Single-Program; Multi-data training。会在每一个进程中拷贝一份模型，并且会用不同的数据拟合不同进程中的模型。DDP 会同步不同进程中模型的参数。

• RPC

适用更通用的训练框架，此时 DDP 无法满足要求。RPC 将自动微分器(autograd) 拓展到其它远程机器上。

• Collective Communication

进程间通讯。用于收集不同进程（卡）上的数据做整体处理。不同后端支持的操作不同。

DP or DDP

DP

• torch.nn.DataParallel

Single Machine, Multi GPU。 改动代码少。但是由于其本质是 Single Process，Multi-Thread，受到 Python GIL(全局解释锁)影响，运行速度受限。

eg.

# 先预先选定一个 GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 

# ... 构建模型
model = nn.DataParallel(model)
model.to(device)

# Train Loop
for data in data_loader:
	input = data.to(device)
	output = model(input)

总结：DP 将数据自动分成等大小的若干份（与 GPU 数量有关）。在不同 GPU 上传入模型（多线程）。每个模型运行过程中的 Batch 是 batch_size / num_gpu。等到每个 GPU 上模型结束后，DP 再收集所有的结果，拼接起来，所以上面 output 的 Batch 仍然是 batch_size

DDP

使用 DDP 时，需要开启多个线程，并在每个线程中创建一个 DDP 对象。由于 DDP 是多进程，所以相较于 DP ，需要先初始化进程组(init_process_group)，这就需要用到 Pytorch 多进程模块。以下是初始化进程的例子：

import torch.multiprocessing as mp

# setup environment
def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    torch.manual_seed(42)
    
# clean up environment
def cleanup():
    dist.destroy_process_group()
    
# 具体的训练过程， world_size 是总共的线程数
def train(rank, world_size):
    setup(rank, world_size)
    device = torch.device("cuda:%d"%(rank)) # 一般每个进程占用一块 GPU
    
    # build a simple model
    model = nn.Linear(10, 10).to(device)
    # construct DDP model
    ddp_model = DDP(model, device_ids=[rank])

	# ... Forward / Backward / Update Op