上次我们了解了DDP的原理和multiprocessing启动的数据并行，这次我们介绍用更流行的torchrun方式启动数据并行。可以理解为 torchrun 是 torch.multiprocessing as mp 的封装版本，但更精确地说，torchrun 封装和简化了分布式多进程训练的管理细节。torchrun 管理了多进程的启动、配置和通信，它主要是为了简化分布式训练的启动过程，而 torch.multiprocessing 是一种更底层的手动启动多进程的方法。
两种启动方式
torch.multiprocessing启动方式
这是利用 Python 的多进程模块 torch.multiprocessing 启动多进程的方式。常用于分布式数据并行（Distributed Data Parallel，DDP）训练。核心原理是通过 Python 的 multiprocessing 模块启动多个进程，每个进程负责一个 GPU 的计算。用户需要自己管理进程启动、初始化通信组（如 dist.init_process_group），并且在代码中通常需要明确调用每个进程的目标函数。
torchrun启动方式
torchrun 是 PyTorch 从 1.9.0 开始引入的一个推荐的启动分布式训练的工具（前身是 torch.distributed.launch）。它是用于启动分布式训练的一个命令行工具。torchrun 自动管理了多进程启动、环境变量配置、world_size（进程数）、rank（进程编号）等。使用 torchrun 时，你只需要在 Python 脚本中调用 torch.distributed 相关函数（如 init_process_group），而不需要显式管理多进程的启动。它的原理是通过传递环境变量（如 RANK, WORLD_SIZE 等）来实现进程间的通信和同步。
优缺点对比
torch.multiprocessing (mp):
- 优点： 更灵活，可以手动管理多进程启动和配置，适合自定义复杂的分布式训练场景。
- 缺点： 需要用户手动处理进程管理、通信配置等细节，容易出错且配置繁琐。
torchrun:
- 优点： 更易用，自动管理多进程启动和配置，减少用户手动设置，适合大部分常见的分布式训练任务。
- 缺点： 某些特殊场景下可能不够灵活，不易于实现极端的定制需求。
代码
torchrun方式自动管理环境变量配置，因此我们不需要像之前那样配置world_size等参数了。
def prepare():
parser= argparse.ArgumentParser()
parser.add_argument('--gpu', default= '2,3')
parser.add_argument('--epochs', type= int, default= 10000, help= 'epoch number')
parser.add_argument('--lr', type= float, default= 1e-4, help= 'learning rate')
parser.add_argument('--weight_decay', type= float, default= 5e-5, help= 'weight decay')
parser.add_argument('--batch_size', type= int, default= 256, help= 'batch size number')
parser.add_argument('--mnist_file_path', type= str, default= '../data/mnist_test.csv')
parser.add_argument('--seed', type= int, default= 1206)
parser.add_argument('--master_addr', type= str, default= 'localhost')
parser.add_argument('--port', type= str, default= '1206')
parser.add_argument('--pt_path', type= str, default= '../pt/')
parser.add_argument('--patience', type= int, default= 5)
args= parser.parse_args()
# setting environment variables to enable DDP
os.environ['CUDA_VISIBLE_DEVICES']= args.gpu # visible gpu devices
return args
因为是torchrun启动方式，初始化进程组那init_method要使用环境变量方式。
def init_ddp(local_rank):
# after this setup, tensors can be moved to GPU via 'a= a.cuda()' rather than 'a= a.to(local_rank)'
torch.cuda.set_device(local_rank)
os.environ['RANK']= str(local_rank)
dist.init_process_group(backend= 'nccl', init_method= 'env://') # init communication mode
相比mp启动方式，main函数也不需要传入指定的LOCAL_RANK了，
def main(args):
local_rank= int(os.environ['LOCAL_RANK'])
# init ddp
init_ddp(local_rank)
# load data
lbls, imgs= load_mnist_data(args.mnist_file_path)
lbls= lbls.long()
imgs= imgs.float().view(-1, 1, 28, 28)
lbl_tr, lbl_val, lbl_te, img_tr, img_val, img_te= split_data(args.seed, lbls, imgs, 0.8, 0.1)
dataset_tr, dataset_val, dataset_te= TensorDataset(img_tr, lbl_tr), TensorDataset(img_val, lbl_val), TensorDataset(img_te, lbl_te)
train_sampler, val_sampler, test_sampler= DistributedSampler(dataset_tr), DistributedSampler(dataset_val), DistributedSampler(dataset_te)
# generator control random behavior (such as data shuffling) to ensure that the order of data on each GPU is different but controllable.
train_loader= DataLoader(dataset= dataset_tr, batch_size= args.batch_size, shuffle= False, sampler= train_sampler, generator= get_ddp_generator())
valid_loader= DataLoader(dataset= dataset_val, batch_size= args.batch_size, shuffle= False, sampler= val_sampler)
test_loader= DataLoader(dataset= dataset_te, batch_size= args.batch_size, shuffle= False, sampler= test_sampler)
# init net
net= ConvNet().cuda()
net= nn.SyncBatchNorm.convert_sync_batchnorm(net)
net= nn.parallel.DistributedDataParallel(net, device_ids= [local_rank])
loss_fn= nn.CrossEntropyLoss().cuda()
opt= torch.optim.SGD(net.parameters(), args.lr)
scaler= GradScaler() # mixed precision training
earlystopping= EarlyStop(args.pt_path, 'cnn.pt', args.patience)
stop_flag_tensor= torch.tensor([0], dtype= torch.long).cuda()
# train
for epoch in range(args.epochs):
# changing the sampling starting point to ensure that the order of data in each round is different, the model's generalization ability can be improved.
train_loader.sampler.set_epoch(epoch)
train(epoch, net, train_loader, loss_fn, opt, scaler)
# valid
acc= valid(net, valid_loader)
if local_rank== 0:
earlystopping.check_early_stopping(-acc, net, opt, scaler)
print(f'epoch: {epoch}, early_stopping.time: {earlystopping.time}, acc: {acc}')
stop_flag_tensor.fill_(int(earlystopping.flag))
dist.all_reduce(stop_flag_tensor, op= dist.ReduceOp.SUM)
if stop_flag_tensor> 0:break
# destory the process group
dist.destroy_process_group()
其他的train和valid部分与mp启动方式的相同，这儿就不再写出了，需要的可以看下面的完整代码。命令行也只需要执行下面命令就可以启动数据并行。standalone表示单机模式，nproc_per_node表示该计算结点上进程的个数，后面是文件路径，输出指定文件，报错指定文件。
torchrun --standalone --nproc_per_node=2 /data_new/gyh/240817ParallelTrain/code/main4ddp_torchrun.py 1> result.txt 2>err.txt
如图1所示，我们成功用torchrun方式实现了数据并行。
图1 torchrun方式启动的数据并行
完整代码
完整代码和数据可以见我的GitHub.
小结
-
torchrun的启动方式，它于torch1.9推出，是比较常用的数据并行的启动方式。相比mp启动方式，它以环境变量方式传递参数，封装了分布式多进程训练和管理方面的细节，减少了参数的配置。 -
代码上，
torchrun方式的数据并行的实现也简单。
参考
[1] chatgpt.
