多卡分布式训练修改为单卡训练(为了本地调试)
·

思考源于项目: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 训练的修改
- 模型实例的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~
更多推荐
所有评论(0)