Update

Author: JewelRoam

graph_net/paddle/check_redundant_incrementally.py

@@ -110,5 +110,5 @@ def main(args):
         help="Path to GraphNet samples",
     )
     args = parser.parse_args()
-    print(args)
+    print(f"[Check Redundancy Arguments] {args}")
     main(args=args)

graph_net/paddle/validate.py

@@ -12,6 +12,7 @@
 import graph_net
 import os
 import ast
+import astor
 import paddle
 
 
@@ -29,7 +30,7 @@ def _get_sha_hash(content):
     return m.hexdigest()
 
 
-def _extract_forward_source(model_path):
+def _extract_forward_source(model_path, class_name):
     source = None
     with open(f"{model_path}/model.py", "r") as f:
         source = f.read()
@@ -38,18 +39,18 @@ def _extract_forward_source(model_path):
     forward_code = None
 
     for node in tree.body:
-        if isinstance(node, ast.ClassDef) and node.name == "GraphModule":
+        if isinstance(node, ast.ClassDef) and node.name == class_name:
             for fn in node.body:
                 if isinstance(fn, ast.FunctionDef) and fn.name == "forward":
-                    return ast.unparse(fn)
+                    return astor.to_source(fn)
     return None
 
 
 def check_graph_hash(args):
     model_path = args.model_path
     file_path = f"{model_path}/graph_hash.txt"
     if args.dump_graph_hash_key:
-        model_str = _extract_forward_source(model_path)
+        model_str = _extract_forward_source(model_path, class_name="GraphModule")
         assert model_str is not None, f"model_str of {args.model_path} is None."
         new_hash_text = _get_sha_hash(model_str)
 
@@ -140,4 +141,5 @@ def main(args):
         help="Path to GraphNet samples folder. e.g '../../samples'",
     )
     args = parser.parse_args()
+    print(f"[Validate Arguments] {args}")
     main(args=args)

graph_net/torch/extractor.py

@@ -7,6 +7,121 @@
 
 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
+
+
+class GraphExtractor:
+    def __init__(
+        self, name, dynamic, mut_graph_codes=None, placeholder_auto_rename=False
+    ):
+        self.subgraph_counter = 0
+        self.name = name
+        self.dynamic = dynamic
+        self.mut_graph_codes = mut_graph_codes
+        self.placeholder_auto_rename = placeholder_auto_rename
+        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)
+        for item in os.listdir(source_dir):
+            source_path = os.path.join(source_dir, item)
+            if os.path.isfile(source_path):
+                target_path = os.path.join(target_dir, item)
+                shutil.move(source_path, target_path)
+
+    def __call__(self, gm: torch.fx.GraphModule, sample_inputs):
+        # 1. Get model path
+        model_path = os.path.join(self.workspace_path, self.name)
+        os.makedirs(model_path, exist_ok=True)
+
+        if self.subgraph_counter == 0:
+            subgraph_path = model_path
+        else:
+            if self.subgraph_counter == 1:
+                subgraph_0_path = os.path.join(model_path, f"subgraph_0")
+                self.move_files(model_path, subgraph_0_path)
+
+            subgraph_path = os.path.join(
+                model_path, f"subgraph_{self.subgraph_counter}"
+            )
+            os.makedirs(subgraph_path, exist_ok=True)
+
+        self.subgraph_counter += 1
+
+        # 2. Get full params
+        params = {}
+        input_idx = 0
+        unique_id = 0
+
+        def try_rename_placeholder(node):
+            assert node.op == "placeholder"
+            if not self.placeholder_auto_rename:
+                return
+            nonlocal unique_id
+            node.target = f"v{unique_id}"
+            unique_id += 1
+            node.name = f"v{unique_id}"
+            unique_id += 1
+
+        for node in gm.graph.nodes:
+            if node.op == "placeholder":
+                try_rename_placeholder(node)
+                input = sample_inputs[input_idx]
+                if isinstance(input, torch.SymInt):
+                    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 self.mut_graph_codes is not None:
+            assert isinstance(self.mut_graph_codes, list)
+            self.mut_graph_codes.append(gm.code)
+        # 3. Generate and save model code
+        base_code = gm.code
+        # gm.graph.print_tabular()
+        write_code = utils.apply_templates(base_code)
+        with open(os.path.join(subgraph_path, "model.py"), "w") as fp:
+            fp.write(write_code)
+
+        # 4. Save metadata
+        metadata = {
+            "framework": "torch",
+            "num_devices_required": 1,
+            "num_nodes_required": 1,
+            "dynamic": bool(self.dynamic),
+            "model_name": self.name,
+        }
+        with open(os.path.join(subgraph_path, "graph_net.json"), "w") as f:
+            json.dump(metadata, f, indent=4)
+
+        # 5. Save tensor metadata
+        # Adapt to different input structures (e.g., single tensor vs. dict/tuple of tensors)
+        converted = utils.convert_state_and_inputs(params, [])
+        utils.save_converted_to_text(converted, file_path=subgraph_path)
+        utils.save_constraints_text(
+            converted,
+            file_path=os.path.join(subgraph_path, "input_tensor_constraints.py"),
+        )
+
+        print(
+            f"Graph and tensors for '{self.name}' extracted successfully to: {model_path}"
+        )
+
+        return gm.forward
 
 
 def extract(name, dynamic=True, mut_graph_codes=None, placeholder_auto_rename=False):
@@ -80,109 +195,11 @@ def forward(self, s0 : torch.SymInt, L_x_ : torch.Tensor):
 
     def wrapper(model: torch.nn.Module):
         assert isinstance(model, torch.nn.Module), f"{type(model)=}"
-
-        class GraphExtractor:
-            def __init__(self):
-                self.subgraph_counter = 0
-
-            def move_files(self, source_dir, target_dir):
-                os.makedirs(target_dir, exist_ok=True)
-                for item in os.listdir(source_dir):
-                    source_path = os.path.join(source_dir, item)
-                    if os.path.isfile(source_path):
-                        target_path = os.path.join(target_dir, item)
-                        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)
-                os.makedirs(model_path, exist_ok=True)
-
-                if self.subgraph_counter == 0:
-                    subgraph_path = model_path
-                else:
-                    if self.subgraph_counter == 1:
-                        subgraph_0_path = os.path.join(model_path, f"subgraph_0")
-                        self.move_files(model_path, subgraph_0_path)
-
-                    subgraph_path = os.path.join(
-                        model_path, f"subgraph_{self.subgraph_counter}"
-                    )
-                    os.makedirs(subgraph_path, exist_ok=True)
-
-                self.subgraph_counter += 1
-
-                # 2. Get full params
-                params = {}
-                input_idx = 0
-                unique_id = 0
-
-                def try_rename_placeholder(node):
-                    assert node.op == "placeholder"
-                    if not placeholder_auto_rename:
-                        return
-                    nonlocal unique_id
-                    node.target = f"v{unique_id}"
-                    unique_id += 1
-                    node.name = f"v{unique_id}"
-                    unique_id += 1
-
-                for node in gm.graph.nodes:
-                    if node.op == "placeholder":
-                        try_rename_placeholder(node)
-                        input = sample_inputs[input_idx]
-                        if isinstance(input, torch.SymInt):
-                            input = torch.tensor(4)
-                        params[node.target] = input
-                        input_idx += 1
-                assert input_idx == len(sample_inputs)
-                if mut_graph_codes is not None:
-                    assert isinstance(mut_graph_codes, list)
-                    mut_graph_codes.append(gm.code)
-                # 3. Generate and save model code
-                base_code = gm.code
-                # gm.graph.print_tabular()
-                write_code = utils.apply_templates(base_code)
-                with open(os.path.join(subgraph_path, "model.py"), "w") as fp:
-                    fp.write(write_code)
-
-                # 4. Save metadata
-                metadata = {
-                    "framework": "torch",
-                    "num_devices_required": 1,
-                    "num_nodes_required": 1,
-                    "dynamic": bool(dynamic),
-                    "model_name": name,
-                }
-                with open(os.path.join(subgraph_path, "graph_net.json"), "w") as f:
-                    json.dump(metadata, f, indent=4)
-
-                # 5. Save tensor metadata
-                # Adapt to different input structures (e.g., single tensor vs. dict/tuple of tensors)
-                converted = utils.convert_state_and_inputs(params, [])
-                utils.save_converted_to_text(converted, file_path=subgraph_path)
-                utils.save_constraints_text(
-                    converted,
-                    file_path=os.path.join(
-                        subgraph_path, "input_tensor_constraints.py"
-                    ),
-                )
-
-                print(
-                    f"Graph and tensors for '{name}' extracted successfully to: {model_path}"
-                )
-
-                return gm.forward
-
-        extractor = GraphExtractor()
+        extractor = GraphExtractor(
+            name, dynamic, mut_graph_codes, placeholder_auto_rename
+        )
         # return torch.compile(backend=extractor, dynamic=dynamic)
         compiled_model = torch.compile(model, backend=extractor, dynamic=dynamic)
-
         return compiled_model
 
     def decorator(module_class):

paddle_samples/PaddleNLP/bert-base-cased/graph_hash.txt

@@ -1 +1 @@
-f2b5a332b1b19703e7ccfb450de96c9c12244144c7b9d305d20587f772fb6672
+517608d4d2699e09c6171648da38a4f924556cf25abd97875599acfdda5807e4

paddle_samples/PaddleNLP/bert-large-cased/graph_hash.txt

@@ -1 +1 @@
-02fa10efca360c8ba7818c367cdeb9979e2af8c72cf489913396a1f241bbad07
+2a46a550da3ca0bd5aa6157a26aff525a3bc69ff8f67fe35b4424303c12e2820

paddle_samples/PaddleNLP/convbert-base/graph_hash.txt

@@ -1 +1 @@
-482fd9e9f201b45c2ce0b22b3037878aa3d139cc203fb35c781fd470140ec962
+e13d4b5e10e7aadcf05e891979bb73813fb3c4c1407b2688fb6ac8f849cdcee0

paddle_samples/PaddleNLP/convbert-medium-small/graph_hash.txt

@@ -1 +1 @@
-e013c0a1d9173f7db5ed91398ad65fa43154e3bc8ce2e15c2d5a6637ddec61d8
+2511edee7164b3327d5efcce7879c5a19a19aec8a86e74e233ae83db0807ed46

paddle_samples/PaddleNLP/convbert-small/graph_hash.txt

@@ -1 +1 @@
-55b1fcce22aee360f71154396a1f528446cae70ebd991927c0abf6c06016d201
+c1e76a465ae2ac6d1cb568acb5f17db4bca92d6d0239061cd319f2d591ba82b9

paddle_samples/PaddleNLP/ernie-1.0/graph_hash.txt

@@ -1 +1 @@
-2ea1f7f9bb52a294ff9fb5fd9876b9e9ed8b4af2fdb6cce93985eedfe50c7a94
+b23ce390b79f214cdbd74ea52c32d6dc141d93b179a7bf75f94bb12e8bd91561

paddle_samples/PaddleNLP/ernie-2.0-base-zh/graph_hash.txt

@@ -1 +1 @@
-40fde6163a995d989050cf8b78b44132b4b62ce218f604dad67aff1f4f5a56f0
+94a3256e834ecd7e836da57b44da751d75ef9e095b04ac00abc37a5e18a01390