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是如何做同步的?

相关推荐
程序设计实验室6 分钟前
2025年了,在 Django 之外,Python Web 框架还能怎么选?
python
飞哥数智坊2 小时前
GPT-5-Codex 发布,Codex 正在取代 Claude
人工智能·ai编程
倔强青铜三2 小时前
苦练Python第46天:文件写入与上下文管理器
人工智能·python·面试
虫无涯3 小时前
Dify Agent + AntV 实战:从 0 到 1 打造数据可视化解决方案
人工智能
Dm_dotnet5 小时前
公益站Agent Router注册送200刀额度竟然是真的
人工智能
算家计算5 小时前
7B参数拿下30个世界第一!Hunyuan-MT-7B本地部署教程:腾讯混元开源业界首个翻译集成模型
人工智能·开源
用户2519162427115 小时前
Python之语言特点
python
机器之心5 小时前
LLM开源2.0大洗牌:60个出局,39个上桌,AI Coding疯魔,TensorFlow已死
人工智能·openai
刘立军5 小时前
使用pyHugeGraph查询HugeGraph图数据
python·graphql
Juchecar6 小时前
交叉熵:深度学习中最常用的损失函数
人工智能