【分布式】大模型分布式训练入门与实践 - 04

相关系列

【分布式】NCCL部署与测试 - 01
【分布式】入门级NCCL多机并行实践 - 02
【分布式】小白看Ring算法 - 03
【分布式】大模型分布式训练入门与实践 - 04

数据并行-Distributed Data Parallel

1.1 背景

数据并行（Distributed Data Parallel）是一种用于加快深度学习模型训练速度的技术。在过去，训练大型模型往往受限于单卡训练的瓶颈，尤其是当处理大规模数据集时。数据并行通过在多个处理单元上同时训练模型，并通过增加BatchSize来提高并行度，有效地减少了训练时间。这种技术在加快深度学习模型训练速度的同时，还提高了模型处理大规模数据集的能力，为解决现实世界中的复杂问题提供了强有力的支持。

1.2 PyTorch DDP

1） DDP训练流程

这里我们通过一段代码过一下DDP的训练流

import torch
import torchvision
import argparse

# Step 1: import distributed
import torch.distributed as dist
parser = argparse.ArgumentParser()

# Step 2: torch.distributed.launchlocal_rank
parser.add_argument("--local_rank", default=-1)
FLAGS = parser.parse_args()
local_rank = FLAGS.local_rank
torch.cuda.set_device(local_rank)

# Step 3: DDP backend
dist.init_process_group(backend='nccl') # ncclgloompi
train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True)

# Step 4: DistributedSamplerDataLoaderbatch_sizebatch_sizebatch_sizebatch_size)
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(my_trainset, batch_size=batch_size, sampler=train_sampler)
device = torch.device("cuda", local_rank)
model = nn.Linear(batch_size*32*32, 10).to(device)

# Step 5: DDP wrapper model
model = DDP(model, device_ids=[local_rank], output_device=local_rank)
loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)

for epoch in range(num_epochs):
 # Step 6: sampler epoch and shuffle
 train_loader.sampler.set_epoch(epoch)
 for step, (input, label) in