Thanks for the report @rsn870! Our model_engine.backward() is a wrapper around loss.backward() that also manages gradient accumulation, gradient reduction for data parallelism, timing, etc. You are seeing differences because additional arguments like retain_graph are not passed along to loss.backward().

I would think that we can simply add a **backward_kwargs argument to our model_engine.backward() that forwards to loss.backward(**backward_kwargs).

@tjruwase / @samyam / @jeffra : do you have any thoughts on this?

Hi , if one does loss.backward() instead of model_engine.backward(loss), so it will be wrong because gradients will not be averaged and therefore at every GPU, the weights will differ eventually leading us to train N models on N GPUs with batch_size/N samples per iteration.
Also please tell me if there is any workaround for retain_graph=True for now?

Please correct me if I am wrong

If the model_engine() is based on DDP() for example I think computed loss should be synchronized among all nodes ?

Hi , if one does loss.backward() instead of model_engine.backward(loss), so it will be wrong because gradients will not be averaged and therefore at every GPU, the weights will differ eventually leading us to train N models on N GPUs with batch_size/N samples per iteration.

Correct, you need model_engine.backward() to reduce gradients and perform other DeepSpeed bookkeeping.

If the model_engine() is based on DDP() for example I think computed loss should be synchronized among all nodes ?

It is not based on DDP directly, so without model_engine.backward() you don't get synchronization.

Also please tell me if there is any workaround for retain_graph=True for now?

Not currently, DeepSpeed needs to propagate additional arguments to our backward interface. I self-assigned but will not be able to get to it until hopefully next week. We are very welcoming of PRs if you would like to contribute :-). I would be happy to help guide through the process.

@ShadenSmith it would be great if we can add retain_graph and other arguments.

@ShadenSmith is this issue still open? If yes can I pick this up?

@asit2898 yes it is, apologies for us not getting to it quickly. PRs are very much welcome.

Thanks a lot @ShadenSmith, I'll work on this issue. This would be my first commit, so any pointers would be truly helpful!