Diving in distributed training in PyTorch

def data_parallel(
	module,
	input,
	device_ids,
	output_device=None
):
    if not device_ids:
        return module(input)

    if output_device is None:
        output_device = device_ids[0]

    # 复制模型
    replicas = nn.parallel.replicate(module, device_ids)
    # 拆分数据
    inputs = nn.parallel.scatter(input, device_ids)
    replicas = replicas[:len(inputs)]
    # 各自前向传播
    outputs = nn.parallel.parallel_apply(replicas, inputs)
    # 汇聚输出
    return nn.parallel.gather(outputs, output_device)

class Model(nn.Module):
    # Our model
    def __init__(self, input_size, output_size):
        super(Model, self).__init__()
        # for convenience
        self.fc = nn.Linear(input_size, output_size)

    def forward(self, input):
        output = self.fc(input)
        print("\tIn Model: input size", input.size(),
              "output size", output.size())
        return output

bs, input_size, output_size = 6, 8, 10
# define inputs
inputs = torch.randn((bs, input_size)).cuda()
model = Model(input_size, output_size)
if torch.cuda.device_count() > 1:
  print("Let's use", torch.cuda.device_count(), "GPUs!")
  # dim = 0 [6, xxx] -> [2, ...], [2, ...], [2, ...] on 3 GPUs
  model = nn.DataParallel(model)
# 先 DataParallel，再 cuda
model = model.cuda()
outputs = model(inputs)
print("Outside: input size", inputs.size(),
	  "output_size", outputs.size())
# assume 2 GPUS are available
# Let's use 2 GPUs!
#    In Model: input size torch.Size([3, 8]) output size torch.Size([3, 10])
#    In Model: input size torch.Size([3, 8]) output size torch.Size([3, 10])
# Outside: input size torch.Size([6, 8]) output_size torch.Size([6, 10])

# save the model
torch.save(model.module.state_dict(), PATH)
# load again
model.module.load_state_dict(torch.load(PATH))
# do anything you want

def scatter(self, inputs, kwargs, device_ids):
    # 不同于源码，获取 batch size 更加灵活
    # 支持只有 kwargs 的情况，如 model(**inputs)
    if len(inputs) > 0:
        bsz = inputs[0].size(self.dim)
    elif kwargs:
        bsz = list(kwargs.values())[0].size(self.dim)
    else:
        raise ValueError("You must pass inputs to the model!")

    num_dev = len(self.device_ids)
    gpu0_bsz = self.gpu0_bsz
    # 除第一块之外每块GPU的bsz
    bsz_unit = (bsz - gpu0_bsz) // (num_dev - 1)
    if gpu0_bsz < bsz_unit:
        # adapt the chunk sizes
        chunk_sizes = [gpu0_bsz] + [bsz_unit] * (num_dev - 1)
        delta = bsz - sum(chunk_sizes)
        # 补足偏移量
        # 会有显存溢出的风险，因而最好给定的 bsz 是可以整除的
        # e.g., 总的 =52 => bsz_0=16, bsz_1=bsz_2=18
        # 总的 =53 => bsz_0=16, bsz_1=19, bsz_2=18
        for i in range(delta):
            chunk_sizes[i + 1] += 1
        if gpu0_bsz == 0:
            chunk_sizes = chunk_sizes[1:]
    else:
        return super().scatter(inputs, kwargs, device_ids)

    return scatter_kwargs(inputs, kwargs, device_ids, chunk_sizes, dim=self.dim)

# 判断数据集长度是否可以整除 GPU 数
# 如果不能，选择舍弃还是补全，进而决定总数
# If the dataset length is evenly divisible by # of replicas
# then there is no need to drop any data, since the dataset
# will be split equally.
if (self.drop_last and len(self.dataset) % self.num_replicas != 0):
    # num_replicas = num_gpus
    self.num_samples = math.ceil((len(self.dataset) -
        self.num_replicas) /self.num_replicas)
else:
    self.num_samples = math.ceil(len(self.dataset) /
        self.num_replicas)
self.total_size = self.num_samples * self.num_replicas

# 根据是否 shuffle 来创建 indices
if self.shuffle:
    # deterministically shuffle based on epoch and seed
    g = torch.Generator()
    g.manual_seed(self.seed + self.epoch)
    indices = torch.randperm(len(self.dataset), generator=g).tolist()
else:
    indices = list(range(len(self.dataset)))
if not self.drop_last:
    # add extra samples to make it evenly divisible
    padding_size = self.total_size - len(indices)
    if padding_size <= len(indices):
        # 不够就按 indices 顺序加
        # e.g., indices 为 [0, 1, 2, 3 ...]，而 padding_size 为4
        # 加好之后的 indices[..., 0, 1, 2, 3]
        indices += indices[:padding_size]
    else:
        indices += (indices * math.ceil(padding_size / len(indices)))[:padding_size]
else:
    # remove tail of data to make it evenly divisible.
    indices = indices[:self.total_size]
assert len(indices) == self.total_size
# subsample
# rank 代表进程 id
indices = indices[self.rank:self.total_size:self.num_replicas]
return iter(indices)

def dist_setup_launch(args):
    # tell DDP available devices [NECESSARY]
    os.environ['CUDA_VISIBLE_DEVICES'] = args.devices
    args.rank = int(os.environ['RANK'])
    args.world_size = int(os.environ['WORLD_SIZE'])
    args.local_rank = int(os.environ['LOCAL_RANK'])

    dist.init_process_group(args.backend,
                            args.init_method,
                            rank=args.rank,
                            world_size=args.world_size)
    # this is optional, otherwise you may need to specify the
    # device when you move something e.g., model.cuda(1)
    # or model.to(args.rank)
    # Setting device makes things easy: model.cuda()
    torch.cuda.set_device(args.rank)
    print('The Current Rank is %d | The Total Ranks are %d'
          %(args.rank, args.world_size))

from torch.utils.data import DataLoader, DistributedSampler

train_sampler = DistributedSampler(train_dataset, seed=args.seed)
train_dataloader = DataLoader(train_dataset,
                              pin_memory=True,
                              shuffle=(train_sampler is None),
                              batch_size=args.per_gpu_train_bs,
                              num_workers=args.num_workers,
                              sampler=train_sampler)

eval_sampler = DistributedSampler(eval_dataset, seed=args.seed)
eval_dataloader = DataLoader(eval_dataset,
                             pin_memory=True,
                             batch_size=args.per_gpu_eval_bs,
                             num_workers=args.num_workers,
                             sampler=eval_sampler)