Repository: pytorch/PiPPy

Issue:
Support custom `torch.fx` tracers in `Pipe.from_tracing`
Some users may want to make use of a custom `torch.fx` tracer when pipelining their model. An example is HF with their custom tracer:

https://github.com/huggingface/transformers/blob/f65fe3663a6c62975a9c04654703252644c9a652/src/transformers/utils/fx.py#L233

We should generalize the interface to `Pipe.from_tracing` to allow users to use a custom tracer: https://github.com/jamesr66a/PiPPy/blob/527af1fd8123d35bd81b9fe304a8d0ed29c9fd8d/pippy/IR.py#L382


Discussion:

Resolving patch:
diff --git a/pippy/IR.py b/pippy/IR.py
--- a/pippy/IR.py
+++ b/pippy/IR.py
@@ -116,7 +116,8 @@ def assign_or_accumulate_grad(forward_node, grad_value):
 _pipeline_tracer = None
 
 def pipe_split():
-    if _pipeline_tracer is not None:
+    # hasattr(_pipeline_tracer, 'graph') is a workaround to support HFTracer
+    if _pipeline_tracer is not None and hasattr(_pipeline_tracer, 'graph'):
         _pipeline_tracer.graph.call_function(pipe_split, (), {})
 
 class MultiUseParameterConfig(Enum):
@@ -375,9 +376,9 @@ def is_leaf_module(self, *args, **kwargs):
         return Pipe(gm, Pipe._hack_build_qualname_mapping(old=seq, new=gm), num_stages, has_loss_and_backward)
 
     @staticmethod
-    def from_traced(mod : torch.nn.Module, traced : torch.fx.GraphModule,
-                    multi_use_param_spec : Optional[MultiUseParamSpec] = None,
-                    loss_fn : Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, **kwargs):
+    def _from_traced(mod: torch.nn.Module, traced: torch.fx.GraphModule,
+                     multi_use_param_spec: Optional[MultiUseParamSpec] = None,
+                     loss_fn: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, **kwargs):
         part_idx = 0
 
         def split_callback(n : torch.fx.Node):
@@ -576,7 +577,8 @@ def move_param_to_callee(callee, param_val, use_idx):
 
     @staticmethod
     def from_tracing(mod : torch.nn.Module, multi_use_param_spec : Optional[MultiUseParamSpec] = None,
-                     loss_fn : Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, **kwargs):
+                     loss_fn : Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,
+                     tracer=None, **kwargs):
         # TODO: abstract partitioning policy
 
         # TODO(pbelevich): Here we can add support for custom torch.fx tracers
@@ -584,7 +586,7 @@ def from_tracing(mod : torch.nn.Module, multi_use_param_spec : Optional[MultiUse
 
         global _pipeline_tracer
         old__pipeline_tracer = _pipeline_tracer
-        _pipeline_tracer = torch.fx.Tracer()
+        _pipeline_tracer = tracer or torch.fx.Tracer()
         try:
             # TODO: tracing policy
             graph = _pipeline_tracer.trace(mod, **kwargs)
@@ -592,7 +594,7 @@ def from_tracing(mod : torch.nn.Module, multi_use_param_spec : Optional[MultiUse
         finally:
             _pipeline_tracer = old__pipeline_tracer
 
-        return Pipe.from_traced(mod, traced, multi_use_param_spec, loss_fn, **kwargs)
+        return Pipe._from_traced(mod, traced, multi_use_param_spec, loss_fn, **kwargs)
 
 
 class PipeSplitWrapper(torch.nn.Module):


Test patch: