在这里插入图片描述

思考源于项目:https://github.com/OFA-Sys/Chinese-CLIP/blob/master/
谢谢作者。

一、多卡分布式训练的依赖

import torch.distributed as dist

在训练命令中的配置:

# Number of GPUs per GPU worker
GPUS_PER_NODE=2
# Number of GPU workers, for single-worker training, please set to 1
WORKER_CNT=1
# The ip address of the rank-0 worker, for single-worker training, please set to localhost
export MASTER_ADDR=XXX.XXX.XXX.XXX
# The port for communication
export MASTER_PORT=8514
# The rank of this worker, should be in {0, ..., WORKER_CNT-1}, for single-worker training, please set to 0
export RANK=0 

torchrun  --nproc_per_node=${GPUS_PER_NODE} --nnodes=${WORKER_CNT} --node_rank=${RANK} \
          --master_addr=${MASTER_ADDR} --master_port=${MASTER_PORT}  main.py \

切换为单卡训练则不需要进行配置直接:

python main.py ...

二、多卡切单卡的修改项

2.1 关于初始化的设置

args.local_device_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(args.local_device_rank)
args.device = torch.device("cuda", args.local_device_rank)
dist.init_process_group(backend="nccl")
args.rank = dist.get_rank()
args.world_size = dist.get_world_size()
修改为:
args.local_device_rank = 0
args.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
torch.cuda.set_device(args.device)
args.rank = 0
args.world_size = 1

2.2 分布式并行训练封装

model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_device_rank], find_unused_parameters=True)

这行代码可以直接注释掉,不需要并行封装

2.3 训练的修改

  1. 模型实例的module方法
m = model.module

这行代码也注释掉,这是torch.nn.parallel.DistributedDataParallel类才有的方法
2. 分布式的聚合方法

dist.all_reduce(cumulative_loss, op=dist.ReduceOp.SUM)
dist.all_reduce(cumulative_acc, op=dist.ReduceOp.SUM)
dist.all_reduce(num_elements, op=dist.ReduceOp.SUM)

这几行也要注释掉。


如此,即可切换多卡到单卡训练,特别由于远程调试看数据,但网络不通畅不得不使用本地训练的情况。Enjoy~

Logo

北京人形旗下天工造物具身智能开源社区,聚焦具身天工与慧思开物两大平台

更多推荐