Repository: microsoft/DeepSpeed

Issue:
Difference between loss.backward() and model_engine.backward(loss) ?
Hi , 

I have tried out both loss.backward() and model_engine.backward(loss) for my code. There are several subtle differences that I have observed , for one retain_graph = True does not work for model_engine.backward(loss) . This is creating a problem since buffers are not being retained every time I run the code for some reason. 

Please look into this if you could.




Discussion:
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](https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/pt/deepspeed_light.py#L736). 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!

Resolving patch:
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -1711,11 +1711,14 @@ def backward(self,
                  loss,
                  allreduce_gradients=True,
                  release_loss=False,
+                 retain_graph=False,
                  scale_wrt_gas=True):
         r"""Execute backward pass on the loss
         Arguments:
             loss: Torch tensor on which to execute backward propagation
             allreduce_gradients: is deprecated, ignored, and will soon be removed'
+            retain_graph: bool, default: false
+                forward on user defined choice of retain_graph
         """
 
         see_memory_usage("Engine before backward", force=self.memory_breakdown())
@@ -1751,9 +1754,9 @@ def backward(self,
         self._start_timers(self.engine_timers.backward_inner_timers)
 
         if self.zero_optimization():
-            self.optimizer.is_gradient_accumulation_boundary = (
-                self.is_gradient_accumulation_boundary())
-            self.optimizer.backward(loss)
+            self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary(
+            )
+            self.optimizer.backward(loss, retain_graph=retain_graph)
         elif self.amp_enabled():
             # AMP requires delaying unscale when inside gradient accumulation boundaries
             # https://nvidia.github.io/apex/advanced.html#gradient-accumulation-across-iterations
@@ -1761,19 +1764,19 @@ def backward(self,
             with amp.scale_loss(loss,
                                 self.optimizer,
                                 delay_unscale=delay_unscale) as scaled_loss:
-                scaled_loss.backward()
+                scaled_loss.backward(retain_graph=retain_graph)
         elif self.fp16_enabled():
             if self.eigenvalue_enabled():
                 self.optimizer.backward(loss, create_graph=True, retain_graph=True)
             else:
-                self.optimizer.backward(loss)
+                self.optimizer.backward(loss, retain_graph=retain_graph)
         elif self.bfloat16_enabled():
             self.optimizer.backward(loss)
         else:
             if self.eigenvalue_enabled():
                 loss.backward(create_graph=True, retain_graph=True)
             else:
-                loss.backward()
+                loss.backward(retain_graph=retain_graph)
 
         self._stop_timers(self.engine_timers.backward_inner_timers)
 


Test patch: