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

相关推荐
蹦蹦跳跳真可爱5899 分钟前
Python----OpenCV(几何变换--图像平移、图像旋转、放射变换、图像缩放、透视变换)
开发语言·人工智能·python·opencv·计算机视觉
蹦蹦跳跳真可爱58913 分钟前
Python----循环神经网络(Transformer ----Layer-Normalization(层归一化))
人工智能·python·rnn·transformer
m0_6256865515 分钟前
Day58
python
夜阳朔18 分钟前
Conda环境激活失效问题
人工智能·后端·python
小Lu的开源日常22 分钟前
AI模型太多太乱?用 OpenRouter,一个接口全搞定!
人工智能·llm·api
m0_7231402340 分钟前
Python训练营-Day49
开发语言·python
北风toto1 小时前
python学习DataFrame数据结构
数据结构·python·学习
亿牛云爬虫专家1 小时前
微服务化采集平台:可扩展性与容错机制
python·微服务·架构·爬虫代理·扩展性·新浪财经·财经新闻
傻啦嘿哟2 小时前
Python爬虫动态IP代理报错全解析:从问题定位到实战优化
爬虫·python·tcp/ip
mit6.8242 小时前
[Meetily后端框架] Whisper转录服务器 | 后端服务管理脚本
c++·人工智能·后端·python