Repository: pytorch/PiPPy

Issue:
Update HF BERT test to forward-backward
![Screen Shot 2022-04-16 at 11 50 41 PM](https://user-images.githubusercontent.com/1160355/163699578-5011adbc-ac6e-4bde-ab5f-d77a4600e462.png)



Discussion:

Resolving patch:
diff --git a/pippy/IR.py b/pippy/IR.py
--- a/pippy/IR.py
+++ b/pippy/IR.py
@@ -17,7 +17,7 @@
 # 5. Add parameter movement to split_module
 
 # TODO: move to a separate file with runtime artifacts
-def stage_backward(stage_output, output_grads, input_values, stage_info : str):
+def stage_backward(stage_output, output_grads, input_values, stage_info : str, outputs_with_grads_idxs : List[int]):
     """
     Given the input value(s) and the corresponding gradient for those/that input
     value(s), compute and accumulate gradients for all parameter values (leaves
@@ -30,6 +30,9 @@ def friendly_debug_info(v):
         else:
             return str(v)
     try:
+        stage_output_with_grads = [stage_output[i] for i in outputs_with_grads_idxs]
+        output_grads_with_grads = [output_grads[i] for i in outputs_with_grads_idxs]
+
         # stage_output may be a composite datatype like dict. Extract all individual
         # tensor values here
         stage_output_tensors = []
@@ -46,7 +49,7 @@ def extract_tensors_with_grads(output_val, grad_val):
                 if grad_val is None:
                     return
                 assert isinstance(grad_val, (tuple, list)), f'grad_value expected to have type {type(output_val)} but got {type(grad_val)}'
-                # assert len(output_val) == len(grad_val) Investigate why it fails! TODO(https://github.com/pytorch/PiPPy/issues/135)
+                assert len(output_val) == len(grad_val)
                 for ov, gv in zip(output_val, grad_val):
                     extract_tensors_with_grads(ov, gv)
             elif isinstance(output_val, dict):
@@ -60,7 +63,7 @@ def extract_tensors_with_grads(output_val, grad_val):
                 # Output is a non-tensor type; just ignore it
                 pass
 
-        extract_tensors_with_grads(stage_output, output_grads)
+        extract_tensors_with_grads(stage_output_with_grads, output_grads_with_grads)
 
         torch.autograd.backward(stage_output_tensors, grad_tensors=output_grad_tensors)
 
@@ -199,11 +202,13 @@ def add_to_live_nodes(n):
             if node.op == 'call_module':
                 output_grads : Union[Tuple[Optional[torch.fx.Node], ...], Optional[torch.fx.Node]]
                 if node in tuples:
-                    stage_output = tuple(n for n in tuples[node] if n in live_nodes)
-                    output_grads = tuple(val_to_grad[n] for n in tuples[node] if n in live_nodes)
+                    stage_output = tuples[node]
+                    output_grads = tuple(val_to_grad.get(n, None) for n in tuples[node])
+                    outputs_with_grads_idxs = [i for i, n in enumerate(tuples[node]) if n in live_nodes]
                 else:
                     stage_output = node,
                     output_grads = val_to_grad[node]
+                    outputs_with_grads_idxs = [0]
 
                 output_grads = (output_grads,) if not isinstance(output_grads, tuple) else output_grads
 
@@ -211,6 +216,7 @@ def add_to_live_nodes(n):
                     'stage_output' : stage_output,
                     'output_grads' : output_grads,
                     'input_values' : list(node.all_input_nodes),
+                    'outputs_with_grads_idxs' : outputs_with_grads_idxs,
                 })
                 # Insert backward stage debug info
                 kwargs_copy = dict(grad_call.kwargs)


Test patch: