目前大家基本都在使用DistributedDataParallel(简称DDP)用来训练,该方法主要用于分布式训练,但也可以用在单机多卡。
第一步:初始化分布式环境,主要用来帮助进程间通信
torch.distributed.init_process_group(backend='nccl')
第二步:负责创建 args.local_rank 变量,并接受 torch.distributed.launch 注入的值
归根到底是创建一个变量,来接收torch.distributed.launch 注入。
目前代码中常见的两种方式:
local_rank = int(os.environ["LOCAL_RANK"]) # 这种是从自定义config文件中获取LOCAL_RANK
另外一种是parser:
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int, default=-1)
args = parser.parse_args()
第三步:每个进程根据自己的local_rank设置应该使用的GPU
torch.cuda.set_device(args.local_rank)
torch.manual_seed(hps.train.seed) # 非必要,尽可能固定种子
第四步:分布式数据和模型
# 分布式数据
train_sampler = DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset, sampler=train_sampler, batch_size=batch_size)
# 分布式模型
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)
第五步:运行
nproc_per_node表示每个node有多少个进程,每个GPU对应一个进程。
nnodes表示使用几个节点,一个节点对应一台电脑。
python -m torch.distributed.launch --nproc_per_node=2 --nnodes=1 train.py
解释
-
torch.distributed.launch参数解析(终端运行命令的参数):
--> python -m torch.distributed.launch --help
usage: launch.py [-h] [--nnodes NNODES] [--node_rank NODE_RANK]
[--nproc_per_node NPROC_PER_NODE] [--master_addr MASTER_ADDR] [--master_port MASTER_PORT]
[--use_env] [-m] [--no_python] [--logdir LOGDIR]
training_script ...
- nnodes:节点的数量,通常一个节点对应一个主机,方便记忆,直接表述为主机
- node_rank:节点的序号,从0开始
- nproc_per_node:一个节点中显卡的数量
- master_addr:master节点的ip地址,也就是0号主机的IP地址,该参数是为了让 其他节点 知道0号节点的位,来将自己训练的参数传送过去处理
- master_port:master节点的port号,在不同的节点上master_addr和master_port的设置是一样的,用来进行通信
- torch.ditributed.launch相关环境变量解析(代码中os.environ中的参数):
- WORLD_SIZE:os.environ["WORLD_SIZE"]所有进程的数量
- LOCAL_RANK:os.environ["LOCAL_RANK"]每张显卡在自己主机中的序号,从0开始
- RANK:os.environ["RANK"]进程的序号,一般是1个gpu对应一个进程
-
一般WORLD_SIZE和RANK和LOCAL_RANK在一起声明
rank = dist.get_rank()
local_rank = int(os.environ["LOCAL_RANK"])
n_gpus = dist.get_world_size() -
多机和单机的一些概念问题
- rank
多机多卡:代表某一台机器
单机多卡:代表某一块GPU - world_size
多机多卡:代表有几台机器
单机多卡:代表有几块GPU - local_rank
多机多卡:代表某一块GPU的编号
单机多卡:代表某一块GPU的编号