Pytorch ddp切换forward函数 验证ddp是否生效

DDP及其在pytorch中应用

ddp默认调用forward函数,有些模型无法使用forward函数,可以对模型包装一下。

python 复制代码
class modelWraper(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, *args, **kwargs):
        return self.model.rlhf(*args, **kwargs)

有时ddp跑起来,不确定是否生效,loss backward后不同rank进程的梯度应该一样的,可以通过print 梯度确认。

python 复制代码
loss.backward()
grad_flag = raw_model.lm_head.weight.grad[0,:3]
print(f"grad {ddp_rank} {grad_flag}")


grad 1 tensor([2.9296e-04, 6.2223e-05, 1.0089e-03], device='cuda:1')
grad 0 tensor([2.9296e-04, 6.2223e-05, 1.0089e-03], device='cuda:0')

pytorch分布式系列2------DistributedDataParallel是如何做同步的?

相关推荐
朱剑君13 分钟前
用Python打造AI玩家:挑战2048,谁与争锋
人工智能·python
东临碣石821 小时前
【AI论文】MM-Eureka:基于规则的大规模强化学习探索视觉“啊哈”时刻
人工智能
JM丫1 小时前
python基础
笔记·python
陌陌6232 小时前
Prompt优化 COT/COD
人工智能·prompt
訾博ZiBo2 小时前
AI日报 - 2025年3月17日
人工智能
树莓集团2 小时前
树莓科技集团董事长:第五代产业园运营模式的深度剖析与展望
大数据·人工智能·科技·物联网·百度
晴空对晚照2 小时前
【论文阅读】AlexNet——深度学习奠基作之一
论文阅读·人工智能·深度学习
天才测试猿2 小时前
接口自动化测试用例
自动化测试·软件测试·python·测试工具·测试用例·接口测试·postman
程序设计实验室3 小时前
DeepSeek+Claude强强联手,使用AI驱动DjangoStarter 3.1框架升级
python·django·djangostarter
孔令飞3 小时前
16 | 实现简洁架构的 Store 层
人工智能·ai·云原生·golang·kubernetes