PyTorch多GPU训练时同步梯度是mean还是sum?

PyTorch 通过两种方式可以进行多GPU训练: DataParallel, DistributedDataParallel. 当使用DataParallel的时候, 梯度的计算结果和在单卡上跑是一样的, 对每个数据计算出来的梯度进行累加. 当使用DistributedDataParallel的时候, 每个卡单独计算梯度, 然后多卡的梯度再进行平均.

下面是实验验证:

DataParallel

python 复制代码
import torch
import os
import torch.nn as nn

def main():
    model = nn.Linear(2, 3).cuda()
    model = torch.nn.DataParallel(model, device_ids=[0, 1])
    input = torch.rand(2, 2)
    labels = torch.tensor([[1, 0, 0], [0, 1, 0]]).cuda()
    (model(input) * labels).sum().backward()
    print('input', input)
    print([p.grad for p in model.parameters()])


if __name__=="__main__":
    main()

执行CUDA_VISIBLE_DEVICES=0,1 python t.py可以看到输出, 代码中对两个样本分别求梯度, 梯度等于样本的值, DataParallel把两个样本的梯度累加起来在不同GPU中同步.

bash 复制代码
input tensor([[0.4362, 0.4574],
        [0.2052, 0.2362]])
[tensor([[0.4363, 0.4573],
        [0.2052, 0.2362],
        [0.0000, 0.0000]], device='cuda:0'), tensor([1., 1., 0.], device='cuda:0')]

DistributedDataParallel

python 复制代码
import torch
import os
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP


def example(rank, world_size):
    # create default process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    # create local model
    model = nn.Linear(2, 3).to(rank)
    print('model param', 'rank', rank, [p for p in model.parameters()])
    # construct DDP model
    ddp_model = DDP(model, device_ids=[rank])
    print('ddp model param', 'rank', rank, [p for p in ddp_model.parameters()])
    # forward pass
    input = torch.randn(1, 2).to(rank)
    outputs = ddp_model(input)
    labels = torch.randn(1, 3).to(rank) * 0
    labels[0, rank] = 1
    # backward pass
    (outputs * labels).sum().backward()
    print('rank', rank, 'grad', [p.grad for p in ddp_model.parameters()])
    print('rank', rank, 'input', input, 'outputs', outputs)
    print('rank', rank, 'labels', labels)
    # update parameters
    optimizer.step()

def main():
    world_size = 2
    mp.spawn(example,
        args=(world_size,),
        nprocs=world_size,
        join=True)

if __name__=="__main__":
    # Environment variables which need to be
    # set when using c10d's default "env"
    # initialization mode.
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "29504"
    main()

执行CUDA_VISIBLE_DEVICES=0,1 python t1.py可以看到输出, 代码中对两个样本分别求梯度, 梯度等于样本的值, 最终的梯度是各个GPU的梯度的平均.

bash 复制代码
model param rank 0 [Parameter containing:
tensor([[-0.4819,  0.0253],
        [ 0.0858,  0.2256],
        [ 0.5614,  0.2702]], device='cuda:0', requires_grad=True), Parameter containing:
tensor([-0.0090,  0.4461, -0.3493], device='cuda:0', requires_grad=True)]
model param rank 1 [Parameter containing:
tensor([[-0.3737,  0.3062],
        [ 0.6450,  0.2930],
        [-0.2422,  0.2089]], device='cuda:1', requires_grad=True), Parameter containing:
tensor([-0.5868,  0.2106, -0.4461], device='cuda:1', requires_grad=True)]
ddp model param rank 1 [Parameter containing:
tensor([[-0.4819,  0.0253],
        [ 0.0858,  0.2256],
        [ 0.5614,  0.2702]], device='cuda:1', requires_grad=True), Parameter containing:
tensor([-0.0090,  0.4461, -0.3493], device='cuda:1', requires_grad=True)]
ddp model param rank 0 [Parameter containing:
tensor([[-0.4819,  0.0253],
        [ 0.0858,  0.2256],
        [ 0.5614,  0.2702]], device='cuda:0', requires_grad=True), Parameter containing:
tensor([-0.0090,  0.4461, -0.3493], device='cuda:0', requires_grad=True)]
rank 1 grad [tensor([[ 0.2605,  0.1631],
        [-0.0934, -0.5308],
        [ 0.0000,  0.0000]], device='cuda:1'), tensor([0.5000, 0.5000, 0.0000], device='cuda:1')]
rank 0 grad [tensor([[ 0.2605,  0.1631],
        [-0.0934, -0.5308],
        [ 0.0000,  0.0000]], device='cuda:0'), tensor([0.5000, 0.5000, 0.0000], device='cuda:0')]
rank 1 input tensor([[-0.1868, -1.0617]], device='cuda:1') outputs tensor([[ 0.0542,  0.1906, -0.7411]], device='cuda:1',
       grad_fn=<AddmmBackward0>)
rank 0 input tensor([[0.5209, 0.3261]], device='cuda:0') outputs tensor([[-0.2518,  0.5644,  0.0314]], device='cuda:0',
       grad_fn=<AddmmBackward0>)
rank 1 labels tensor([[-0., 1., -0.]], device='cuda:1')
rank 0 labels tensor([[1., 0., -0.]], device='cuda:0')
相关推荐
Blossom.11815 分钟前
机器学习在智能供应链中的应用:需求预测与物流优化
人工智能·深度学习·神经网络·机器学习·计算机视觉·机器人·语音识别
Gyoku Mint22 分钟前
深度学习×第4卷:Pytorch实战——她第一次用张量去拟合你的轨迹
人工智能·pytorch·python·深度学习·神经网络·算法·聚类
zzywxc78724 分钟前
AI大模型的技术演进、流程重构、行业影响三个维度的系统性分析
人工智能·重构
点控云25 分钟前
智能私域运营中枢:从客户视角看 SCRM 的体验革新与价值重构
大数据·人工智能·科技·重构·外呼系统·呼叫中心
zhaoyi_he33 分钟前
多模态大模型的技术应用与未来展望:重构AI交互范式的新引擎
人工智能·重构
葫三生1 小时前
如何评价《论三生原理》在科技界的地位?
人工智能·算法·机器学习·数学建模·量子计算
m0_751336392 小时前
突破性进展:超短等离子体脉冲实现单电子量子干涉,为飞行量子比特奠定基础
人工智能·深度学习·量子计算·材料科学·光子器件·光子学·无线电电子
美狐美颜sdk6 小时前
跨平台直播美颜SDK集成实录:Android/iOS如何适配贴纸功能
android·人工智能·ios·架构·音视频·美颜sdk·第三方美颜sdk
DeepSeek-大模型系统教程6 小时前
推荐 7 个本周 yyds 的 GitHub 项目。
人工智能·ai·语言模型·大模型·github·ai大模型·大模型学习
有Li6 小时前
通过具有一致性嵌入的大语言模型实现端到端乳腺癌放射治疗计划制定|文献速递-最新论文分享
论文阅读·深度学习·分类·医学生