调试pytorch DDP训练代码

背景

Pytorch提供了Distributed Data Parallel (DDP)工具以便在多机多卡上并行训练，并提供了torchrun指令来启动。然而，torchrun指令启动不便于debug。可以通过修改成等价mp.spawn启动方式先debug，完成后再转回torchrun指令启动正式训练。

流程

假设原始DDP训练代码是：

import torch.distributed as dist

def main():
	args.local_rank = int(os.environ["LOCAL_RANK"])
    args.world_size = int(os.environ["WORLD_SIZE"]