[tgif][multiforward] allow codegen to generate different func name (#111446)

842974287 · pytorchmergebot · commit 675df7520a04 · 2023-10-19T21:19:30.000Z
Summary: see Shiyan's design doc for ATM TS publish weights dedupe https://fb.quip.com/HnUVAjUMaXMQ Test Plan: tested in N4454041 after D50341352 that multiforward method is working for ts model Differential Revision: D45750812 Pull Request resolved: #111446 Approved by: https://github.com/842974287
diff --git a/torch/fx/graph.py b/torch/fx/graph.py
@@ -311,17 +311,18 @@ def _parse_stack_trace(stack_trace: str):
 class CodeGen:
     def __init__(self):
         self._body_transformer: Optional[TransformCodeFunc] = None
+        self._func_name: str = "forward"
 
     def gen_fn_def(self, free_vars: List[str], maybe_return_annotation: str) -> str:
         """
         Given the free variables and a return annotation, generates the beginning of the FX function.
-        By default, `gen_fn_def(['a', 'b'], '') == 'def forward(a, b):'`
+        By default, `gen_fn_def(['a', 'b'], '') == 'def {self._func_name}(a, b):'`
         """
         # If the original function didn't have self as its first argument, we
         # would have added it.
         if len(free_vars) == 0 or free_vars[0] != 'self':
             free_vars.insert(0, 'self')
-        return f"def forward({', '.join(free_vars)}){maybe_return_annotation}:"
+        return f"def {self._func_name}({', '.join(free_vars)}){maybe_return_annotation}:"
 
     def generate_output(self, output_args: Argument) -> str:
         """
diff --git a/torch/fx/graph_module.py b/torch/fx/graph_module.py
@@ -6,7 +6,7 @@
 import traceback
 import warnings
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Set, Type, Union
+from typing import Any, Callable, Dict, List, Optional, Set, Type, Union
 
 import torch
 import torch.nn as nn
@@ -85,12 +85,20 @@ def _exec_with_source(src: str, globals: Dict[str, Any], co_fields=None):
 
 
 def _forward_from_src(src: str, globals: Dict[str, Any], co_fields=None):
+    return _method_from_src(
+        method_name="forward", src=src, globals=globals, co_fields=co_fields
+    )
+
+
+def _method_from_src(
+    method_name: str, src: str, globals: Dict[str, Any], co_fields=None
+) -> Callable:
     # avoid mutating the passed in dict
     globals_copy = globals.copy()
     _exec_with_source(src, globals_copy, co_fields)
-    forward_fn = globals_copy["forward"]
-    del globals_copy["forward"]
-    return forward_fn
+    fn = globals_copy[method_name]
+    del globals_copy[method_name]
+    return fn
 
 
 def _format_import_statement(name: str, obj: Any, importer: Importer) -> str:
diff --git a/torch/fx/passes/split_utils.py b/torch/fx/passes/split_utils.py
@@ -1,6 +1,6 @@
 import copy
 from dataclasses import dataclass, field
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Type, Union
 
 import torch.fx
 from torch.fx._compatibility import compatibility
@@ -59,7 +59,10 @@ class Component:
 
 @compatibility(is_backward_compatible=False)
 def split_by_tags(
-    gm: torch.fx.GraphModule, tags: List[str], return_fqn_mapping: bool = False
+    gm: torch.fx.GraphModule,
+    tags: List[str],
+    return_fqn_mapping: bool = False,
+    GraphModuleCls: Type[torch.fx.GraphModule] = torch.fx.GraphModule,
 ) -> Union[torch.fx.GraphModule, Tuple[torch.fx.GraphModule, Dict[str, str]]]:
     """
     Splits a GraphModule using tags on its graph nodes. We honor the order of
@@ -287,7 +290,7 @@ def remap_func(x):
         if x.op == "get_attr":
             setattr(main_root, x.name, getattr_recursive(gm, x.target))  # type: ignore[arg-type]
 
-    result_gm = torch.fx.GraphModule(main_root, main_g)
+    result_gm = GraphModuleCls(main_root, main_g)
     if return_fqn_mapping:
         return result_gm, orig_to_split_fqn_mapping
 
