Training Technology One : Distribution

本节内容以BLIP的分布式训练代码为蓝本介绍分布式训练的过程.

本文采用DDP作为分布式框架, 库为torch.distributed.

训练命令

python -m torch.distributed.launch <分布式参数> <文件名> <文件参数>

常用的分布式参数有:

1. nnodes : 节点的数量;
2. node_rank : 节点的序号;
3. nproc_per_node : 节点中显卡的数量.

通常我们采用单机多卡的模式,故只用传输第三个参数.

当命令下发之后, 其会创建一些环境变量, 即os.environ的参数:

1. WORLD_SIZE: os.environ["WORLD_SIZE"], 所有进程的数量;
2. LOCAL_RANK:os.environ["LOCAL_RANK"], 每张显卡在自己主机中的序号;
3. RANK:os.environ["RANK"], 进程的序号, 通常1个GPU对应一个进程.

不同的卡上共享脚本,但是采样的数据,以及local_rank有所不同.

初始化

获得分布式参数

最重要的是得到local_rank, 这可以通过命令行的方式获得.

从torch1.10开始，官方建议使用环境变量的方式来获取local_rank, 在后期版本中，会移除命令行的方式。

if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
       args.rank = int(os.environ["RANK"])
       args.world_size = int(os.environ['WORLD_SIZE'])
       args.gpu = int(os.environ['LOCAL_RANK'])
elif 'SLURM_PROCID' in os.environ:
	   args.rank = int(os.environ['SLURM_PROCID'])
	   args.gpu = args.rank % torch.cuda.device_count()
else:
	   print('Not using distributed mode')
	   args.distributed = False
	   return

设定GPU

torch.cuda.set_device(args.gpu)

在进程内部设定可见的GPU, 这样 to('cuda')即可.

初始化

torch.distributed.init_process_group(
backend,
init_method,
world_size,
rank
)

backend表示后端方式, 对于GPU的分布式训练选择NCCL, 对于CPU的分布式训练选择GLOO.
init_method表示初始化方法:

1. init_method='tcp://ip:port'： 通过指定rank 0（即：MASTER进程）的IP和端口，各个进程进行信息交换。 需指定 rank 和 world_size 这两个参数。
2. init_method='file://path'：通过所有进程都可以访问共享文件系统来进行信息共享。需要指定rank和world_size参数。
3. init_method=env://：从环境变量中读取分布式的信息(os.environ)，主要包括 MASTER_ADDR, MASTER_PORT, RANK, WORLD_SIZE。 其中，rank和world_size可以选择手动指定，否则从环境变量读取。

一般选择env://.

同步

torch.distributed.barrier

在分布式训练过程中，多个进程需要进行数据的同步和通信。torch.distributed.barrier函数可以用来实现进程的同步，确保所有进程达到一个统一的同步点后再继续执行后续的代码。

输出控制(非常规流程,技巧)*

def setup_for_distributed(is_master):
    """
    This function disables printing when not in master process
    """
    import builtins as __builtin__
    builtin_print = __builtin__.print

    def print(*args, **kwargs):
        force = kwargs.pop('force', False)
        if is_master or force:
            builtin_print(*args, **kwargs)

    __builtin__.print = print

setup_for_distributed(args.rank == 0)        

分布式判断(非常规流程,技巧)*

1. 设定flag, 详情见上文代码;
2. torch.ditributed.is_available()+torch.distributed.is_initialized()

数据集

dataset的格式不变, 重点考虑采样器的设计.

torch.utils.data.DistributedSampler(
dataset,
num_replicas,		# 总进程数
rank,						# 当前进程等级
shuffle,
seed,
drop_last				# 是否丢弃数据的尾部
)

对train, 我们将shuffle设为True, 反之为False.

模型

model_without_ddp = model
if args.distributed:
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
    model_without_ddp = model.module   

分布式封装的模型用于训练, 反之用于测试及存储.

训练

注意事项-1: 采样

train_loader.sampler.set_epoch(epoch)

注意事项-2: 只对主进程的模型进行测试与保存

注意事项-3: 每一个epoch后添加torch.distributed.barrier()