编辑|我不爱机器学习
1. 相关概念
- Process ：python的一个实例。可以使用进程控制GPU。简单理解：进程个数=使用的GPU的个数
- Node ：一个节点就像一台拥有所有资源的计算机。简单理解：多个节点=多台电脑/主机
- World-Size ：可用的GPU总数。它是节点总数 和每个节点的 GPU 总数的乘积。例如，如果有两台服务器，每台服务器有两个 GPU，则World-Size为 4。
- Rank ：它是在所有进程中标识进程的ID。例如，如果有两个节点服务器每个有四个 GPU，排名则为不同的0 - 7。Rank 0将识别进程0等等。
- Local Rank ：Rank是用来识别所有节点(主机个数*每个主机的GPU个数)，而Local Rank是用来识别本地服务器，比如节点2上的某个进程的rank可以为2和0，则可以理解为，在所有进程中排名为2(global rank)，在本地机器上中排名为0。
PyTorch 中的多进程处理
python
torch.multiprocessing.spawn(fn,
args=(),
nprocs=1,
join=True,
daemon=False,
start_method='spawn')它用于产生nprocs给定的进程数。这些进程使用args运行fn。此函数可用于在每个 GPU 上训练模型。
举个例子。假设有一个节点服务器有两个 GPU。
-
可以创建一个函数
fn来处理训练部分。 -
nprocs将等于两个，即 GPU 的数量。
因此，torch.multiprocessing.spawn可通过args在每个GPU上训练函数fn。这可以在每个节点/服务器上完成。
接下来，介绍进程将如何相互协调。
分布式并行DDP
在 DDP 中，模型在每个 GPU 上复制，每个 GPU 由一个进程处理。
DDP 需要知道以下信息：
- 节点数量和每个节点中的 GPU 数量。可以从这些信息中得到World-Size。
- 使用 DDP 进行训练需要进程之间的同步和通信。这是通过
distributed.init_process_group实现的。
pyton
os.environ['MASTER_ADDR'] = '19.16.19.19'
os.environ['MASTER_PORT'] = '8888'
dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=rank)init_process_group
-
需要
master的详细信息。这些已设置为环境变量。它用于所有进程的同步。 -
还需要 world_size 以及进程 rank，注意不是local rank。
可以将所有这些信息作为torch.multiprocessing.spawn中的args传递。rank 可以通过：node_rank × gpus_per_node + local_rank 找到。
注意，local rank 的变化范围是 [0,nprocs]。
分布式采样器
在 DDP 中，每个进程都处理一个 GPU，每个 GPU 使用数据集的单独某块来训练数据集。例如，如果有两个 GPU 和 100 个训练样本，批量大小为 50，那么每个 GPU 将使用 50 个非重叠训练样本。这是通过 PyTorch 提供的DistributedSampler实现的。它确保每个 GPU 都使用数据集的单独子集。它需要 world_size 和 进程的rank/global rank。
在 PyTorch 中使用分布式计算：
python
def get_dataset(world_size, global_rank, batch_size):
transform_train = transforms.Compose([
transforms.Resize((224,224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.456, 0.456, 0.456], std=[0.224, 0.224, 0.224])])
train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, transform=transform_train, download=True)
train_sampler=torch.utils.data.distributed.DistributedSampler(train_dataset,\
num_replicas=world_size,rank=global_rank)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, \
batch_size=batch_size, shuffle=False, num_workers=0, pin_memory=True,\
sampler=train_sampler)
return train_loader使用train_sampler 将不同样本传递给每个 GPU。
接下来，编写一个函数来训练模型：
python
def train(local_rank, args):
## global_rank is the global rank of the process (among all the GPUs (not just on a particular node). )
global_rank = args.node_rank * args.gpus + local_rank
dist.init_process_group(backend='nccl', init_method='env://', world_size=args.world_size, rank=global_rank)
torch.manual_seed(0)
torch.cuda.set_device(local_rank)
batch_size = 64
criterion = nn.CrossEntropyLoss().cuda(local_rank)
model=models.resnet18(pretrained=True)
model.fc=nn.Linear(512, 10)
model.cuda(local_rank)
model = nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
train_loader=get_dataset(args.world_size, global_rank, batch_size)
for epoch in range(10):
for i, (images, labels) in enumerate(train_loader):
images = images.cuda(non_blocking=True)
labels = labels.cuda(non_blocking=True)
# Forward pass
outputs = model(images)
loss = criterion(outputs, labels)
# Backward and optimize
optimizer.zero_grad()
loss.backward()
optimizer.step()
## if this condition is not used, output will print for each process (GPU)
if local_rank == 0:
print('Epoch [{}/{}],Loss: {}'.format(epoch + 1, 10,loss.item()))最后，使用 multiprocessing 的spawn训练函数：
python
parser = argparse.ArgumentParser()
parser.add_argument('--total_nodes', default=1, type=int, help='total number of the nodes')
parser.add_argument('--gpus', default=1, type=int, help='number of gpus per node')
parser.add_argument('--node_rank', default=0, type=int, help='rank of present node (server).')
args = parser.parse_args()
args.world_size = args.gpus * args.total_nodes
os.environ['MASTER_ADDR'] = '19.16.19.19'
os.environ['MASTER_PORT'] = '8888'
mp.spawn(train, nprocs=args.gpus, args=(args,))然后可以通过在每个节点上使用合适的参数运行脚本来使用 DDP。例如，如果有 2 个节点，每个节点有 4 个 GPU：
节点服务器0 将是：
--total_nodes 2 --gpus 4 --node_rank 0
对于节点 1：
--total_nodes 2 --gpus 4 --node_rank 1
请注意，数据集必须存在于每个节点上。
简单的方法
使用torch.distributed.launch可以大大简化上述方法。它将负责spawn过程。我们只需要解析命令行参数：--local_rank，其余的将由该模块处理。修改后的脚本如下：
python
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int, default=-1, help="local_rank for distributed training on gpus")
args = parser.parse_args()
os.environ['MASTER_ADDR'] = '19.16.19.19'
os.environ['MASTER_PORT'] = '8888'
dist.init_process_group(backend='nccl', init_method='env://')
torch.manual_seed(0)
torch.cuda.set_device(args.local_rank)
model=models.resnet18(pretrained=True)
model.fc=nn.Linear(512, 10)
model.cuda(args.local_rank)
model = nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])
batch_size = 64
criterion = nn.CrossEntropyLoss().cuda(args.local_rank)
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
train_loader=get_dataset( batch_size)
for epoch in range(10):
for i, (images, labels) in enumerate(train_loader):
images = images.cuda(non_blocking=True)
labels = labels.cuda(non_blocking=True)
# Forward pass
outputs = model(images)
loss = criterion(outputs, labels)
# Backward and optimize
optimizer.zero_grad()
loss.backward()
optimizer.step()
## if this condition is not used, output will print for each process (GPU)
if args.local_rank == 0:
print('Epoch [{}/{}],Loss: {}'.format(epoch + 1, 10,loss.item()))在节点 0 上，此脚本可写作：
python
python -m torch.distributed.launch --nnodes 2 --node_rank 0 --nproc_per_node=2 train_distributed_2.py在节点 1 上：
python
python -m torch.distributed.launch --nnodes 2 --node_rank 1 --nproc_per_node=2 train_distributed_2.py同样，对于其他节点也是如此。主地址(master address)，也可以作为此模块的参数传递。
python
--master_addr=127.0.0.2 --master_port=29502还有一些关于模型保存以及DDP使用技巧，可以参考：
--master_addr=127.0.0.2 --master_port=29502
参考：
https://shivgahlout.github.io/2021-05-18-distributed-computing/