训练命令

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

常用的分布式参数有:

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

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

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

WORLD_SIZE: os.environ["WORLD_SIZE"], 所有进程的数量;
LOCAL_RANK:os.environ["LOCAL_RANK"], 每张显卡在自己主机中的序号;
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 = Falsereturn

设定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表示初始化方法:

init_method='tcp://ip:port'：通过指定rank 0（即：MASTER进程）的IP和端口，各个进程进行信息交换。需指定 rank 和 world_size 这两个参数。
init_method='file://path'：通过所有进程都可以访问共享文件系统来进行信息共享。需要指定rank和world_size参数。
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__.printdef 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)

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

设定flag, 详情见上文代码;
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