PyTorch 分布式训练实现(DP/DDP/torchrun/多机多卡)

###	第一步：构建模型
'''
module 需要分发的模型
device_ids 可分发的gpu，默认分发到所有看见GPU（环境变量设置的）
output_device 结果输出设备 通常设置成逻辑gpu的第一个
'''
module=your_simple_net() #你的模型
Your_Parallel_Net=torch.nn.DataParallel(module,device_ids=None,output_device=None)

### 第二步：数据迁移
inputs=inputs.to(device)	
labels=labels.to(device)	
#此处的device通常应为模型输出的output_device，否则无法计算loss

#启动方式，shell中运行：
python -m torch.distributed.launch \
       --nnodes 1 \
       --nproc_per_node=4 \
       YourScript.py
# nnodes: 表示有多少个节点，可以通俗的理解为有多少台机器
# nproc_per_node 表示每个节点上有多少个进程，每个进程一般独占一块GPU

########################## 	第1步	 ##########################
#初始化
'''
在启动器为我们启动python脚本后，在执行过程中，启动器会将当前进程的（其实就是 GPU的）index 通过参数传递给 python，
我们可以这样获得当前进程的 index：即通过参数 local_rank 来告诉我们当前进程使用的是哪个GPU，
用于我们在每个进程中指定不同的device
'''
parse.add_argument('--local_rank',type=int)
args=parser.parse_args()
local_rank=args.local_rank
torch.cuda.set_device(local_rank)
'''
init_process_group用于初始化GPU通信方式（NCCL）和参数的获取方式（env代表通过环境变量）
gpu使用nccl最快，gloo为cpu分布式训练，mpu则需要重新编码
init_method 指定如何初始化进程组的 URL。 
默认及推荐为'env://' 其他初始化方式与多机多卡有关（not sure，挖个坑）
'''
torch.distributed.init_process_group('nccl'，init_method='env://')
device = torch.device(f'cuda:{args.local_rank}')

##########################	第2步  ##########################
#处理Dataloader
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset,shuffle=True)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)
#torch.utils.data.DataLoader中的shuffle应该设置为False（默认），因为打乱的任务交给了sampler
##########################	第3步  ##########################
#模型的初始化
model=torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])
'''
使用 DistributedDataParallel 包装模型，
它能帮助我们为不同 GPU 上求得的梯度进行allreduce（即汇总不同 GPU 计算所得的梯度，并同步计算结果）。
allreduce 后不同 GPU 中模型的梯度均为 allreduce 之前各 GPU 梯度的均值。
''''
##########################	第4步  ##########################
#同DP，进行inputs、labels的设备转移

# local_rank参数应当从环境变量中读取，而不是通过参数传递。

### BEFORE
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()

local_rank = args.local_rank

### NOW
import os
local_rank = int(os.environ["LOCAL_RANK"])

#运行脚本
torchrun train_script.py #除了--use_env参数，其他torch.distributed.launch所使用的参数均可使用，
			 #如nnodes、nproc_per_node

########################## 	第1步	 ##########################
#初始化
rank = int(os.environ["RANK"])
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(rank % torch.cuda.device_count())
dist.init_process_group(backend="nccl")
device = torch.device("cuda", local_rank)
########################## 	第2步	 ##########################
#模型定义
model = model.to(device)
model = DDP(model, device_ids=[local_rank], output_device=local_rank)

#数据集操作与DDP一致

#####运行
'''
exmaple: 2 node, 8 GPUs per node (16GPUs)
需要在两台机器上分别运行脚本
注意细节：node_rank master 为 0 
机器1
>>> python -m torch.distributed.launch \
    --nproc_per_node=8 \
    --nnodes=2 \
    --node_rank=0 \
    --master_addr="master的ip" \
    --master_port=xxxxx \
    YourScript.py
机器2
>>> python -m torch.distributed.launch \
    --nproc_per_node=8 \
    --nnodes=2 \
    --node_rank=1 \
    --master_addr="master的ip" \
    --master_port=xxxxx \
    YourScript.py

'''