unnecessary gradients in case of multiple optimizers / engines

[kracwarlock]

GAN training is currently broken. Pytorch Lightning also had the same issue described in detail: Lightning-AI/pytorch-lightning#591, Lightning-AI/pytorch-lightning#592

In a 2 engine setup, my expectation was that running backward() in one of the engines would only accumulate gradients in the engine on which backward() is called. However, this is not the current behaviour. The Lightning folks have recently fixed it with this PR: Lightning-AI/pytorch-lightning#603 where backward() will only compute gradients on that engine’s parameters (imho there should be a flag to set this behavior that defaults to True). It would be great to have this fix in DeepSpeed as well.

[tjruwase]

@kracwarlock Thanks for sharing this issue. We will look into as soon as possible, but we also open to PRs.

[eltonzheng]

@kracwarlock, thanks for sharing this issue, if it is about "fake_out = D(G(z))", and want D.backward() won't call into G.backward(), how about changing it to "fake_out = D(G(z).detach())"? If this is not the issue you mentioned, could you paste some sample code?

[kracwarlock]

i think there are 2 cases of interest for D_fake = self.D(x_gen)

• Discriminator pass: we want to prevent accumulation of gradients on the Generator. This can be solved by replacing with D_fake = self.D(x_gen.detach()). This is the fix you pointed out and this would work.
• Generator pass: we want to prevent gradient accumulation on the Discriminator, but we want them to accumulate on the Generator. D_fake = self.D(x_gen.detach()) would not do that. with torch.no_grad(): D_fake = self.D(x_gen) would also not work. The only solution I can think of would be to switch off requires grad for D before the backward() and turn it on afterwards.

I have done that manually in my code right now. However, I think PR#603 I pointed out earlier has a better solution. If you have 2 optimizers then backward() only accumulates gradients of the corresponding parameter.