support torchscript backend

frontend/bytecode_writter.py

@@ -234,7 +234,6 @@ def fix_constants(instructions: List[Instruction],
                 const_list.append(entry)
                 const_set.add(entry)
             inst.arg = const_list.index(entry)
-    print("const_list", const_list)
     code_options["co_consts"] = tuple((x[1] for x in const_list))
 
 

frontend/fx_graph.py

@@ -70,6 +70,35 @@ def eager_due_to_inductor_bug(node: torch.fx.Node) -> bool:
     elif backend == 'xla':
         return torch._dynamo.backends.torchxla.aot_torchxla_trace_once(
             gm, example_inputs)
+    elif backend == 'script':
+        import os, importlib, re, random
+        random_number = str(random.randint(0, 1000000))
+        os.makedirs('tmp/fx_module_' + random_number, exist_ok=True)
+        gm.to_folder('tmp/fx_module_' + random_number)
+
+        module = importlib.import_module('tmp.fx_module_' + random_number)
+        model = module.FxModule().cuda().eval()
+        real_inputs = []
+        for x in example_inputs:
+            if x.dtype == torch.float32:
+                real_inputs.append(
+                    torch.rand(*x.shape,
+                               dtype=x.dtype,
+                               layout=x.layout,
+                               device=x.device))
+            elif x.dtype == torch.int64:
+                real_inputs.append(
+                    torch.randint(0,
+                                  2,
+                                  size=x.shape,
+                                  dtype=x.dtype,
+                                  layout=x.layout,
+                                  device=x.device))
+            else:
+                raise NotImplementedError
+        with torch.no_grad():
+            script_model = torch.jit.trace(model, real_inputs)
+            return script_model
     else:
         raise RuntimeError(f"Unknown backend: {backend}")
 

frontend/guard_tracker.py

@@ -143,10 +143,9 @@ def add_submodule(self, module: torch.nn.Module) -> None:
         # self.written = True # not mark as written as graph break may happen
 
     def add_subparam(self, param: torch.nn.Parameter) -> None:
-        new_param_name = "__external_param__" + str(len(self.subparam_paths))
+        new_param_name = "external_param__" + str(len(self.subparam_paths))
         self.root.register_parameter(new_param_name, param)
         self.subparam_paths[param] = new_param_name
-        # self.written = True # not mark as written as graph break may happen
 
     def as_node_args_kwargs(
         self, args: list[Any], kwargs: dict[str, Any]

frontend/tracer.py

@@ -87,8 +87,9 @@ def preprocess_frame(
             frame_cache = get_frame_cache(frame_id)
             frame_cache.update_code(frame.f_code, frame_id, is_callee)
             new_code, code_map = frame_cache.get_new_code(is_callee)
-            print("bytecode to run:")
-            print(format_insts(code_map.guard_insts))
+            if is_debug:
+                print("bytecode to run:")
+                print(format_insts(code_map.guard_insts))
             trace_func = get_trace_func(frame_id)
 
         except Exception as e:
@@ -102,12 +103,14 @@ def postprocess_frame(frame: FrameType, frame_id: int) -> None:
             from .bytecode_writter import SHOULD_NOT_CALL_REWRITE
             if SHOULD_NOT_CALL_REWRITE:
                 raise ValueError("should not call postprocess")
-            print(f"postprocess frame {frame.f_code.co_filename}")
+            if is_debug:
+                print(f"postprocess frame {frame.f_code.co_filename}")
             set_frame_root(frame_id, f)
             frame_cache = get_frame_cache(frame_id)
             frame_cache.update_code(frame.f_code, frame_id, is_callee)
         except Exception as e:
-            print("exception in postprocess:", e, type(e))
+            if is_debug:
+                print("exception in postprocess:", e, type(e))
             print(traceback.format_exc())
             raise e
         return

frontend/variables/__init__.py

@@ -92,7 +92,6 @@ def make_var_from_value(
                                       fx_graph, extract_code_at_start)
     else:
         # NOTE: use any instead of iteartor_var to represent iterator with unknown source due to the hardness of getting iterable and num_iters
-        print("generate any for", value, type(value), extract_code_at_start)
         return AnyVar.from_value(value, need_guard_check, helper_functions,
                                  fx_graph, extract_code_at_start)
 

test/test_model_deberta.py

@@ -1157,7 +1157,6 @@ def get_model():
     config.num_hidden_layers = 2
     config.return_dict = False
     model = DebertaModel(config).to(device)
-    print("model type", type(model))
     return model