DDP分布式训练的官方demo及相关知识

2023-12-28 16:52:58

参考链接

'''
1.分布式训练:并行化,跨进程和跨集群的计算
2.torch.distributed.init_process_group() 来初始化进程组,需要指定worker的通信机制,
    一般为nccl(NVIDIA推出),1个进程对应1个gpu.
    nccl是nvidia的显卡通信方式,用于把模型参数、梯度传递到训练的每个节点上。
3.DDP让进程组的所有worker进行通信
4.torch.utils.data.DataLoader中的batch_size指的是每个进程下的batch_size。也就是说,
    总batch_size是这里的batch_size再乘以并行数(world_size)。

#########################################################################
rank是指在整个分布式任务中进程的序号;local_rank是指在一个node上进程的相对序号
nnodes是指物理节点数量 (node:物理节点,可以是一台机器也可以是一个容器,节点内部可以有多个GPU )
node_rank是物理节点的序号
nproc_per_node是指每个物理节点上面进程的数量
########################################
上一个运算题: 每个node包含16个GPU,且nproc_per_node=8,nnodes=3,机器的node_rank=5,请问word_size是多少?
答案:word_size = 3*8 = 24
结论:word_size = nproc_per_node * nnodes
##########################################################################

多进程组启动方法:torch.distributed.launch  或 torchrun
常用启动方法示例:python3 -m torch.distributed.launch --nproc_per_node 2 main.py
'''

import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP


class ToyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))


def demo_basic():
    dist.init_process_group('nccl')
    rank = dist.get_rank()
    print(f'running on {rank}')

    # get the number of GPUs available
    n_gpus = torch.cuda.device_count()
    # get every rank
    device_id = rank % n_gpus
    model = ToyModel().to(device_id)
    '''
    多卡多线程,设置broadcast_buffers=False,会报错:
    'RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:3 and cpu! 
    (when checking argument for argument mat1 in method wrapper_CUDA_addmm)'
    示例
    ddp_model = DDP(model, broadcast_buffers=False)    
    '''
    ddp_model = DDP(model, device_ids=[device_id])

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
    optimizer.zero_grad()
    inputs = torch.randn(3, 10)
    labels = torch.randn(3, 10).to(device_id)
    output = ddp_model(inputs)
    loss = loss_fn(output, labels)
    loss.backward()
    optimizer.step()


if __name__ == '__main__':
    demo_basic()

文章来源:https://blog.csdn.net/Dreaming_of_you/article/details/135272604
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。