Repository: pytorch/PiPPy

Issue:
[Bug] Pipe.from_tracing(transformers.T5Model()) crashes with 'GraphModule' object has no attribute 'decoder'
```
import inspect

import transformers.utils.fx as fx
from pippy.IR import MultiUseParameterConfig, Pipe
from transformers import *

model = T5Model(T5Config())
print(model)

input_names = model.dummy_inputs.keys()
sig = inspect.signature(model.forward)
concrete_args = {p.name: p.default for p in sig.parameters.values() if p.name not in input_names}

hf_tracer = fx.HFTracer()

model_pipe = Pipe.from_tracing(model, MultiUseParameterConfig.TRANSMIT, tracer=hf_tracer,
                               concrete_args=concrete_args)

```

```
T5Model(
  (shared): Embedding(32128, 512)
  (encoder): T5Stack(...)
  (decoder): T5Stack(...)
  ...
```

```
Traceback (most recent call last):
  File "/Users/pbelevich/PycharmProjects/PiPPy/test/hf_t5_test.py", line 15, in <module>
    model_pipe = Pipe.from_tracing(model, MultiUseParameterConfig.TRANSMIT, tracer=hf_tracer,
  File "/Users/pbelevich/PycharmProjects/PiPPy/pippy/IR.py", line 606, in from_tracing
    return Pipe._from_traced(mod, traced, multi_use_param_spec, loss_fn, **kwargs)
  File "/Users/pbelevich/PycharmProjects/PiPPy/pippy/IR.py", line 457, in _from_traced
    mod_itr = getattr(mod_itr, atom)
  File "/Users/pbelevich/miniconda3/envs/PiPPy/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1186, in __getattr__
    raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'GraphModule' object has no attribute 'decoder'
```

Maybe related to https://github.com/jamesr66a/PiPPy/issues/44


Discussion:

Resolving patch:
diff --git a/pippy/IR.py b/pippy/IR.py
--- a/pippy/IR.py
+++ b/pippy/IR.py
@@ -204,7 +204,7 @@ def __init__(self, split_gm : torch.fx.GraphModule, qualname_mapping : Dict[str,
                    (node.op, node.target) == ('call_method', 'backward') or
                    (node.op, node.target) == ('call_function', stage_backward) or
                    (node.op, node.target) == ('call_function', torch.add) or
-                   (node.op, node.target) == ('call_function', sync_barrier))
+                   (node.op, node.target) == ('call_function', sync_barrier)), node
 
         # Detect replicated parameters so we know that we have to do an additional allreduce
         # before applying the optimizer
@@ -356,6 +356,8 @@ def split_callback(n : torch.fx.Node):
         # TODO: what does split do with module invocations? does it move the modules
         # into the submodules?
         split = split_module(traced, mod, split_callback, qualname_map)
+        # a (custom) tracer can produce dead code like orphan get_attr nodes
+        split.graph.eliminate_dead_code()
 
         # peephole to remove pipe_split
         for submodule in split.modules():
@@ -413,6 +415,8 @@ def move_param_to_callee(root, callee_name, param_val, use_idx, is_buffer):
 
             return get_attr
 
+        to_delete = list()  # a list of nodes for deferral deletion
+
         for node in split.graph.nodes:
             if node.op == 'get_attr' and len(node.users) == 1:
                 user = list(node.users)[0]
@@ -429,7 +433,11 @@ def move_param_to_callee(root, callee_name, param_val, use_idx, is_buffer):
 
                 move_param_to_callee(split, user.target, param_val, use_idx, is_buffer)
 
-                delattr(mod_itr, atoms[-1])
+                to_delete.append((mod_itr, atoms))
+
+        # deferral deletion
+        for (mod_itr, atoms) in to_delete:
+            delattr(mod_itr, atoms[-1])
 
         split.graph.lint()
         split.recompile()
@@ -559,14 +567,12 @@ def from_tracing(mod : torch.nn.Module, multi_use_param_spec : Optional[MultiUse
                      tracer=None, **kwargs):
         # TODO: abstract partitioning policy
 
-        # TODO(pbelevich): Here we can add support for custom torch.fx tracers
-        #  https://github.com/jamesr66a/PiPPy/issues/15
-
         global _pipeline_tracer
         old__pipeline_tracer = _pipeline_tracer
         _pipeline_tracer = tracer or torch.fx.Tracer()
         try:
             # TODO: tracing policy
+            mod = copy.deepcopy(mod)  # because further pipe building activities can modify mod
             graph = _pipeline_tracer.trace(mod, **kwargs)
             traced = torch.fx.GraphModule(mod, graph)
         finally:


Test patch:
diff --git a/test/hf_test.py b/test/hf_test.py
--- a/test/hf_test.py
+++ b/test/hf_test.py
@@ -1,13 +1,13 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates
 import inspect
+import sys
+import unittest
 
 import torch
 
 import transformers.utils.fx as fx
 from pippy.IR import MultiUseParameterConfig, Pipe, annotate_split_points, PipeSplitWrapper
 from transformers import *
-import sys
-import unittest
 
 
 def albert_splitter(model) -> int:
@@ -82,7 +82,7 @@ def transformer_splitter(model) -> int:
             annotate_split_points(model, {f'encoder.block.{i}': PipeSplitWrapper.SplitPoint.BEGINNING})
         for i in range(model.config.num_decoder_layers):
             annotate_split_points(model, {f'decoder.block.{i}': PipeSplitWrapper.SplitPoint.BEGINNING})
-        return model.config.num_layers + model.config.num_decoder_layers
+        return model.config.num_layers + model.config.num_decoder_layers + 1
     elif isinstance(model, GPTNeoPreTrainedModel):
         if isinstance(model, GPTNeoModel):
             for i in range(model.config.num_layers):
@@ -136,30 +136,30 @@ def roberta_splitter(model: GPT2PreTrainedModel) -> int:
 num_choices = 3
 seq_length = 32
 
+
 class HFModelsTest(unittest.TestCase):
     pass
 
+
 for _model_cls in fx._SUPPORTED_MODELS:
     def scope(model_cls, replicate):
         def test_case(self):
-            if model_cls in [T5Model, T5ForConditionalGeneration]:  # https://github.com/jamesr66a/PiPPy/issues/57
-                self.skipTest('Known bad model')
             splitter = splitters.get(model_cls.base_model_prefix)
             print(f"Testing {model_cls.__name__} ", file=sys.stderr)
             assert splitter is not None
             config_cls = model_cls.config_class
             config = config_cls()
             if model_cls in [GPT2ForSequenceClassification, GPTNeoForSequenceClassification,
-                            GPTJForSequenceClassification] or model_cls.__name__.startswith("Roberta"):
+                             GPTJForSequenceClassification] or model_cls.__name__.startswith("Roberta"):
                 config.pad_token_id = 0
             model = model_cls(config)
             model.eval()
-            # print(model)
 
             submodules_cnt = splitter(model)
-            # print(model)
 
             input_names = model.dummy_inputs.keys()
+            if model_cls.__name__.startswith("T5"):
+                input_names -= ["decoder_attention_mask"]
             sig = inspect.signature(model.forward)
             concrete_args = {p.name: p.default for p in sig.parameters.values() if p.name not in input_names}
 
@@ -167,8 +167,7 @@ def test_case(self):
 
             multi_use_param_config = MultiUseParameterConfig.REPLICATE if replicate else MultiUseParameterConfig.TRANSMIT
             model_pipe = Pipe.from_tracing(model, multi_use_param_config, tracer=hf_tracer,
-                                        concrete_args=concrete_args)
-            # print(model_pipe)
+                                           concrete_args=concrete_args)
             assert submodules_cnt == len(list(model_pipe.split_gm.children()))
 
             if model_cls.__name__.endswith('MultipleChoice'):
@@ -177,8 +176,13 @@ def test_case(self):
                 input = torch.zeros(bs, seq_length, dtype=torch.long)
             else:
                 input = torch.zeros(bs, seq_length, dtype=torch.long).random_(model.config.vocab_size)
-            model_output = model(input)
-            model_pipe_output = model_pipe(input)
+
+            if model_cls.__name__.startswith("T5"):
+                model_output = model(input_ids=input, decoder_input_ids=input)
+                model_pipe_output = model_pipe(input_ids=input, decoder_input_ids=input)
+            else:
+                model_output = model(input)
+                model_pipe_output = model_pipe(input)
 
             def recursive_value_check(out, ref_out):
                 if isinstance(out, torch.Tensor):
@@ -201,6 +205,7 @@ def recursive_value_check(out, ref_out):
 
         return test_case
 
+
     setattr(HFModelsTest, f'test_{_model_cls.__name__}_transmit', scope(_model_cls, False))
     setattr(HFModelsTest, f'test_{_model_cls.__name__}_replicate', scope(_model_cls, True))
 
diff --git a/test/test_ir.py b/test/test_ir.py
--- a/test/test_ir.py
+++ b/test/test_ir.py
@@ -301,8 +301,9 @@ def test_from_tracing_preserves_buffer(self):
         ec = ExampleCode()
         pipe = Pipe.from_tracing(ec)
         assert 'moved_buffer' in dict(pipe.split_gm.submod_1.named_buffers())
-        # NB: identity comparison
-        assert ec.buffer is pipe.split_gm.submod_1.moved_buffer
+        # NB: identity comparison is not possible because pipe has a deepcopy of the original parameters
+        assert ec.buffer is not pipe.split_gm.submod_1.moved_buffer
+        torch.testing.assert_allclose(ec.buffer, pipe.split_gm.submod_1.moved_buffer)
 
     def test_annotate_split_points_end(self):
         class Foo(torch.nn.Module):