Pytorch ddp切换forward函数 验证ddp是否生效-CSDN博客
阿里云国内75折 回扣 微信号:monov8 |
阿里云国际,腾讯云国际,低至75折。AWS 93折 免费开户实名账号 代冲值 优惠多多 微信号:monov8 飞机:@monov6 |
ddp默认调用forward函数有些模型无法使用forward函数可以对模型包装一下。
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 梯度确认。
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')