fix autocast (#235)

Author: fangfangssj

graph_net/torch/extractor.py

@@ -7,6 +7,9 @@
 
 torch._dynamo.config.capture_scalar_outputs = True
 torch._dynamo.config.capture_dynamic_output_shape_ops = True
+torch._dynamo.config.capture_sparse_compute = True
+torch._dynamo.config.raise_on_ctx_manager_usage = False
+torch._dynamo.config.allow_rnn = True
 
 
 def extract(name, dynamic=True, mut_graph_codes=None, placeholder_auto_rename=False):
@@ -84,6 +87,11 @@ def wrapper(model: torch.nn.Module):
         class GraphExtractor:
             def __init__(self):
                 self.subgraph_counter = 0
+                self.workspace_path = os.environ.get("GRAPH_NET_EXTRACT_WORKSPACE")
+                if not self.workspace_path:
+                    raise EnvironmentError(
+                        "Environment variable 'GRAPH_NET_EXTRACT_WORKSPACE' is not set."
+                    )
 
             def move_files(self, source_dir, target_dir):
                 os.makedirs(target_dir, exist_ok=True)
@@ -94,13 +102,8 @@ def move_files(self, source_dir, target_dir):
                         shutil.move(source_path, target_path)
 
             def __call__(self, gm: torch.fx.GraphModule, sample_inputs):
-                # 1. Get workspace path
-                workspace_path = os.environ.get("GRAPH_NET_EXTRACT_WORKSPACE")
-                if not workspace_path:
-                    raise EnvironmentError(
-                        "Environment variable 'GRAPH_NET_EXTRACT_WORKSPACE' is not set."
-                    )
-                model_path = os.path.join(workspace_path, name)
+                # 1. Get model path
+                model_path = os.path.join(self.workspace_path, name)
                 os.makedirs(model_path, exist_ok=True)
 
                 if self.subgraph_counter == 0:
@@ -140,6 +143,15 @@ def try_rename_placeholder(node):
                             input = torch.tensor(4)
                         params[node.target] = input
                         input_idx += 1
+
+                    if node.op == "call_function" and hasattr(node.target, "__name__"):
+                        if node.target.__name__ in [
+                            "_enter_autocast",
+                            "_exit_autocast",
+                        ]:
+                            node.replace_all_uses_with(node.args[0])
+                            gm.graph.erase_node(node)
+
                 assert input_idx == len(sample_inputs)
                 if mut_graph_codes is not None:
                     assert isinstance(mut_graph_codes, list)

samples/transformers-auto-model/Qwen1.5-0.5B/graph_hash.txt

@@ -0,0 +1 @@
+fdf650e0bdc3326eccec3edf5adcc6683a1d516b7d8192a39e3fd1d27bfc423f

samples/transformers-auto-model/Qwen1.5-0.5B/graph_net.json

@@ -0,0 +1,7 @@
+{
+    "framework": "torch",
+    "num_devices_required": 1,
+    "num_nodes_required": 1,
+    "dynamic": false,
+    "model_name": "Qwen/Qwen1.5-0.5B"
+}