有时候训练模型为了提高速度，单纯的提高batch size 已经没用了, 这时候我们会使用多卡训练，这里我们使用DDP（Distributed Data Parallel）来实现多卡训练。

这里是单机多卡, 毕竟我们集群上也没几台机器…

DDP 比较麻烦, 因此这里记录一下如何使用DDP进行多卡训练。

# 这里只是一个大致的样例，具体的代码需要根据实际情况进行修改
# 有些常见的包我也不写了

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

parser = argparse.ArgumentParser(description="PyTorch DDP Example")
# 这个是我的习惯, 常常是为了适配单卡和多卡的场景
parser.add_argument("--is_multi_gpu", type=bool, default=False)
# 下面这个是必要的, 这个参数不需要手动传入
parser.add_argument("--local-rank", type=int, default=-1)
# 好像2.0 开始的torch 才是 local-rank
# 1.x 的是 local_rank
args = parser.parse_args()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# 指定gpu0为主GPU
is_multi_gpu = args.is_multi_gpu
# 后面全是用这种方式去进行适配单卡和多卡
if is_multi_gpu == True:
    device = args.local_rank
    torch.cuda.set_device(args.local_rank)
    dist.init_process_group(backend='nccl')

# 首先加载数据集
train_dataset = ...
train_dataloader = DataLoader(train_dataset, batch_size=..., shuffle=True, num_workers=...)
if is_multi_gpu == True:
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
    train_dataloader = DataLoader(train_dataset, batch_size=..., sampler=train_sampler, num_workers=...)


# 超級重要!
# 如果你要load 模型, 那么你需要在这里加上这个
# model.load_state_dict(torch.load("model.pth", map_location="cpu"))
# 不然你的卡一會多占用很多份的顯存,份數為 GPU 數量 - 1
model = ...
model = model.to(device)
if is_multi_gpu == True:
    model = DDP(model, device_ids=[args.local_rank], output_device=args.local_rank)
    # 一般会加上另一个参数(因为有些参数不会回归)
    model = DDP(model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True)



optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)

for epoch in range(epochs):
    if is_multi_gpu:
        train_dataloader.sampler.set_epoch(epoch)
    train_bar = tqdm(train_dataloader, desc="训练")
    model.train()
    for step, batch in enumerate(train_bar):
        ...
        ...
        ...
        train_bar.update(1)

# 结束后要销毁
if is_multi_gpu == True:
    dist.destroy_process_group()

# 如果要保存模型应该保存的是 model.module 而不是 model
# torch.save(model.module.state_dict(), "model.pth")

大致是上面这样一个流程，当然具体的代码需要根据实际情况进行修改. 之前遇到一个 bug, 就是在集群上使用多卡进行训练的时候, 卡越多, 训练的速度越慢, 这是因为在集群上训练的时候需要设置一个参数, #SBATCH --ntasks-per-node=8, 这里的8是因为我使用了8张卡, 如果你使用的是4张卡, 那么这里就是4.

原因是你使用的几张卡, 他就会启动几个进程, 但是之前默认的线程数是1, 所以好像会有问题, 可能会阻塞? maybe, 我也不是很清楚, 但是设置了这个参数之后, 训练速度就恢复正常了, 每增加一张卡都能提高一些训练速度.

然后是启动训练的脚本

# 我这里实验用环境的是
# torch 2.0.1
# 这里的 8 是因为我使用了8张卡
# --local-rank 不需要手动传入
python -m torch.distributed.launch --nproc_per_node 8 xxx.py --is_multi_gpu True

# 由于启动后, 会启动和显卡数一样多的线程, 因此会导致你线程里面一旦输出了什么, 就会输出 8份的内容,保存文件也是一样, 会保存8份(虽然会覆盖)‘
# 因此需要在你的代码里面加上这个
if (is_multi_gpu == True and dist.get_rank() == 1) or is_multi_gpu == False:
    pass
# 这样就只会输出一份了
# 当我要保存什么, 或者输出什么不想输出多份的时候就用这个

题外话, 我这个代码需要用多线程进行处理, 所以可以在代码里面加

CPU_THREAD_NUM = 8
os.environ["OMP_NUM_THREADS"] = "{}".format(CPU_THREAD_NUM)
torch.set_num_threads(CPU_THREAD_NUM)
# 然后你启动后的每个线程可以用 8 个CPU线程
# 至少我的会快一点

然后抢节点前需要看一下节点的CPU资源, 希望别因为CPU需要的太多, 导致一直等, 毕竟卡更重要一些.

# gpu002 可以换成你想看的节点
scontrol show node gpu002 -o | grep -oP 'CPUAlloc=\K[^ ]+|CPUTot=\K[^ ]+' | awk 'BEGIN { OFS="=" } { if (NR % 2 == 0) { print "已分配", $0 } else { print "总共", $0 } }'

你可能会需要?

# register
# model.register_buffer('step', torch.tensor(0))

# DDP 会封装原模型的 forward
# 如果你的模型里面有其他的操作, 比如生成操作, 那么你需要手动调用
# model.module.原模型的函数()