关于多卡训练Bart的问题

多卡预训练中文bart。用了不到7000条数据,8张卡一起跑,batchsize调到最低了。按理说内存是够的,但是还是爆内存了。

GPUS_PER_NODE=8
MASTER_ADDR=localhost
MASTER_PORT=6000
NNODES=1
NODE_RANK=0
WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))

DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE --nnodes $NNODES --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT"

python -m torch.distributed.launch $DISTRIBUTED_ARGS \
        pretrain_bart.py \
       --num-layers 12 \
       --hidden-size 1024 \
       --num-attention-heads 16 \
       --micro-batch-size 1 \
       --global-batch-size 8 \

上面是脚本的部分代码,运行后,开始训练不到5秒钟,就报超内存

RuntimeError: CUDA out of memory. Tried to allocate 42.00 MiB (GPU 2; 9.78 GiB total capacity; 7.63 GiB already allocated; 4.56 MiB free; 7.66 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

这个错误报了8次,我感觉多卡并没有跑起来,有人能帮忙看看怎么解决吗?

通常的多卡训练是指每个显卡上都加载一样的模型,然后把 batch 平分到多卡上计算梯度后汇总,你报错在 gpu2,基本上断定多卡没问题,而是显存不够应对 BART 的大小。如果你要拆分模型以减轻显存消耗,那需要其他的方式。

你现在可以看看 fp16 混合精度训练或者 ddp_shard 等方式,减少显存消耗。