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

相关推荐
查理零世1 分钟前
【算法】 区间合并(附蓝桥杯真题) python
python·算法·蓝桥杯
机器之心4 分钟前
提前免费!百度连发两款模型,我们实测:能听歌看电影,还会蛐蛐人
人工智能
lihuayong23 分钟前
RAG的工作原理以及案例列举
人工智能·rag·文本向量化·检索增强生成·语义相似度
果冻人工智能28 分钟前
Google 发布 Gemma 3 —— 你需要了解的内容
人工智能
人工智能研究所32 分钟前
使用OpenCV与Python编写自己的俄罗斯方块小游戏
开发语言·python·opencv
DDD小小小宇宙32 分钟前
python列表基础知识
开发语言·windows·python
-一杯为品-42 分钟前
【动手学深度学习】#2线性神经网络
人工智能·深度学习·神经网络
SZ1701102311 小时前
语音识别 FireRedASR-AED模型主要特点
人工智能·语音识别
@黄色海岸1 小时前
【sklearn 05】sklearn功能模块
人工智能·python·sklearn