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

本文主要是介绍Pytorch ddp切换forward函数 验证ddp是否生效，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

DDP及其在pytorch中应用

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

class modelWraper(nn.Module):def __init__(self, model):super().__init__()self.model = modeldef 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')

pytorch分布式系列2——DistributedDataParallel是如何做同步的？

这篇关于Pytorch ddp切换forward函数 验证ddp是否生效的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！

---