Repository: pytorch/PiPPy

Issue:
_null_coalesce_accumulate after the last stage_backward before sync_barrier produces "list index out of range"
1. Use T5 with 6 encoders and 36 decoders:
```
Layers:
0: encoder's embedding
1-6: encoders
7: decoder's embedding
8-43: decoders
44: lm_head
```
2. Add `attention_mask` to the inputs before tracing
3. Add split points before the following layers: [3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 44]
4. The result pipe has extra `_null_coalesce_accumulate` after the last `stage_backward` before `sync_barrier`

```
    stage_backward_15 = pippy_IR_stage_backward(stage_output = (getitem, getitem_1), output_grads = (getitem_93, _null_coalesce_accumulate_33), input_values = [input_ids, attention_mask], outputs_with_grads_idxs = [0, 1], stage_info = 'stage_backward_15 for stage %submod_0 : [#users=2] = call_module[target=submod_0](args = (%input_ids, %attention_mask), kwargs = {})');  getitem = getitem_1 = getitem_93 = _null_coalesce_accumulate_33 = input_ids = attention_mask = None
    getitem_95 = stage_backward_15[0]
    getitem_96 = stage_backward_15[1];  stage_backward_15 = None
    getitem_97 = getitem_95[0]
    getitem_98 = getitem_95[1];  getitem_95 = None
    _null_coalesce_accumulate_34 = pippy_IR__null_coalesce_accumulate(getitem_90, getitem_98);  getitem_90 = getitem_98 = None
    sync_barrier = pippy_IR_sync_barrier({'loss': getitem_6, 'logits': getitem_7, 'encoder_last_hidden_state': getitem_4}, [getitem_9, getitem_13, getitem_19, getitem_25, getitem_31, getitem_37, getitem_43, getitem_49, getitem_55, getitem_61, getitem_67, getitem_73, getitem_79, getitem_85, getitem_92, getitem_96]);  getitem_6 = getitem_7 = getitem_4 = getitem_9 = getitem_13 = getitem_19 = getitem_25 = getitem_31 = getitem_37 = getitem_43 = getitem_49 = getitem_55 = getitem_61 = getitem_67 = getitem_73 = getitem_79 = getitem_85 = getitem_92 = getitem_96 = None
    return sync_barrier
```
`_null_coalesce_accumulate_34` leads to this:
```
Traceback (most recent call last):
  File "/Users/pbelevich/miniconda3/envs/PiPPy/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/Users/pbelevich/PycharmProjects/PiPPy/examples/slurm/hf/t5/pippy_t5.py", line 293, in run_worker
    run_master(args, pp_ranks_per_dp_group[rank])
  File "/Users/pbelevich/PycharmProjects/PiPPy/examples/slurm/hf/t5/pippy_t5.py", line 221, in run_master
    pipe_driver: PipelineDriverBase = schedules[args.schedule](t5_pipe, args_chunk_spec, kwargs_chunk_spec,
  File "/Users/pbelevich/PycharmProjects/PiPPy/pippy/PipelineDriver.py", line 1224, in __init__
    self._init_remote_executors()
  File "/Users/pbelevich/PycharmProjects/PiPPy/pippy/PipelineDriver.py", line 907, in _init_remote_executors
    node.meta['fw_stage'] = executor_descriptors[bw_idx].name
IndexError: list index out of range
```

[The repro code](https://gist.github.com/pbelevich/913af2a597a4058de2d30a12ed3d6ecf)


Discussion:

Resolving patch:
diff --git a/pippy/IR.py b/pippy/IR.py
--- a/pippy/IR.py
+++ b/pippy/IR.py
@@ -114,7 +114,7 @@ def _insert_stage_symbolic_backward(g : pippy.fx.Graph, output_loss_value_spec):
     val_to_grad : Dict[pippy.fx.Node, Optional[pippy.fx.Node]] = {loss_node : None}
 
     def assign_or_accumulate_grad(forward_node, grad_value):
-        if forward_node in val_to_grad:
+        if forward_node in val_to_grad and forward_node.op != 'placeholder':
             grad_value = g.call_function(_null_coalesce_accumulate, (val_to_grad[forward_node], grad_value))
         val_to_grad[forward_node] = grad_value
 
diff --git a/pippy/PipelineDriver.py b/pippy/PipelineDriver.py
--- a/pippy/PipelineDriver.py
+++ b/pippy/PipelineDriver.py
@@ -595,7 +595,7 @@ def extract_value_ref_args(arg):
                 fut.wait()
 
         with self.value_store_lock:
-            assert output_unique_key not in self.value_store
+            assert output_unique_key not in self.value_store, f"output_unique_key = {output_unique_key}"
             self.value_store[output_unique_key] = RefcountedFuture(future, output_refcount)
 
         return ValueReference(self.stage_id, output_unique_key)
diff --git a/pippy/microbatch.py b/pippy/microbatch.py
--- a/pippy/microbatch.py
+++ b/pippy/microbatch.py
@@ -29,7 +29,8 @@ def shard_dict_of_args(args_dict, args_chunk_spec, num_chunks, _debug_mask_minib
 
     arg_specs = []
 
-    assert len(args_dict) == len(args_chunk_spec)
+    assert len(args_dict) == len(args_chunk_spec),\
+        f'args_dict.keys() = {list(args_dict.keys())} args_chunk_spec.keys() = {list(args_chunk_spec.keys())}'
     for arg_key, arg in args_dict.items():
         chunk_spec = args_chunk_spec[arg_key]
         flat, spec = tree_flatten(arg)


Test patch: