Merge branch 'develop' of github.com:PaddlePaddle/GraphNet into device_rewrite_sample_pass

Author: lixinqi

graph_net/graph_net_json_file_util.py

@@ -1,17 +1,61 @@
-from pathlib import Path
 import json
+from pathlib import Path
 
 kDimensionGeneralizationPasses = "dimension_generalization_passes"
+kDataTypeGeneralizationPasses = "data_type_generalization_passes"
 kSymbolicDimensionReifier = "symbolic_dimension_reifier"
 
+# Fields for dtype generalization metadata
+kDtypeGeneralizationTargetDtype = "dtype_generalization_target_dtype"
+kDtypeGeneralizationPrecision = "dtype_generalization_precision"
+kDtypeGeneralizationGenerated = "dtype_generalization_generated"
+
 
 def read_json(model_path):
+    """
+    Read JSON from graph_net.json file.
+
+    Args:
+        model_path: Path to model directory
+
+    Returns:
+        Dictionary containing JSON data
+    """
     graph_net_json_file_path = Path(f"{model_path}/graph_net.json")
     return json.loads(graph_net_json_file_path.read_text())
 
 
 def update_json(model_path, field, value):
-    graph_net_json_file_path = Path(f"{model_path}/graph_net.json")
-    graph_net_json = json.loads(graph_net_json_file_path.read_text())
+    """
+    Update a single field in graph_net.json.
+
+    Args:
+        model_path: Path to model directory or graph_net.json file
+        field: Field name to update
+        value: Value to set
+    """
+    if isinstance(model_path, (str, Path)):
+        model_path = Path(model_path)
+        # If it's a file path, use it directly; otherwise assume it's a directory
+        if model_path.suffix == ".json":
+            graph_net_json_file_path = model_path
+        else:
+            graph_net_json_file_path = model_path / "graph_net.json"
+    else:
+        graph_net_json_file_path = Path(f"{model_path}/graph_net.json")
+
+    # Read existing JSON
+    if graph_net_json_file_path.exists():
+        with open(graph_net_json_file_path, "r") as f:
+            graph_net_json = json.load(f)
+    else:
+        graph_net_json = {}
+
+    # Update field
     graph_net_json[field] = value
-    graph_net_json_file_path.write_text(json.dumps(graph_net_json, indent=4))
+
+    # Atomic write: write to temp file then rename
+    temp_path = graph_net_json_file_path.with_suffix(".json.tmp")
+    with open(temp_path, "w") as f:
+        json.dump(graph_net_json, f, indent=4)
+    temp_path.replace(graph_net_json_file_path)

graph_net/test/dtype_gen_test.sh

@@ -0,0 +1,46 @@
+#!/bin/bash
+
+GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(
+os.path.dirname(graph_net.__file__))")
+SAMPLES_ROOT="$GRAPH_NET_ROOT/../samples"
+OUTPUT_DIR="/tmp/dtype_gen_samples"
+mkdir -p "$OUTPUT_DIR"
+
+# Step 1: Initialize dtype generalization passes
+config_json_str_init=$(cat <<EOF
+{
+    "handler_path": "$GRAPH_NET_ROOT/torch/dtype_generalizer.py",
+    "handler_class_name": "InitDataTypeGeneralizationPasses",
+    "handler_config": {
+        "dtype_list": ["float16", "bfloat16"],
+        "model_path_prefix": "$SAMPLES_ROOT"
+    }
+}
+EOF
+)
+CONFIG_INIT=$(echo "$config_json_str_init" | base64 -w 0)
+
+python3 -m graph_net.model_path_handler --model-path "timm/resnet18" --handler-config=$CONFIG_INIT
+python3 -m graph_net.model_path_handler --model-path "transformers-auto-model/opus-mt-en-gmw" --handler-config=$CONFIG_INIT
+
+# Step 2: Apply passes to generate samples
+config_json_str_apply=$(cat <<EOF
+{
+    "handler_path": "$GRAPH_NET_ROOT/torch/dtype_generalizer.py",
+    "handler_class_name": "ApplyDataTypeGeneralizationPasses",
+    "handler_config": {
+        "output_dir": "$OUTPUT_DIR",
+        "model_path_prefix": "$SAMPLES_ROOT",
+        "model_runnable_predicator_filepath": "$GRAPH_NET_ROOT/torch/constraint_util.py",
+        "model_runnable_predicator_class_name": "RunModelPredicator",
+        "model_runnable_predicator_config": {
+            "use_dummy_inputs": true
+        }
+    }
+}
+EOF
+)
+CONFIG_APPLY=$(echo "$config_json_str_apply" | base64 -w 0)
+
+python3 -m graph_net.model_path_handler --model-path "timm/resnet18" --handler-config=$CONFIG_APPLY
+python3 -m graph_net.model_path_handler --model-path "transformers-auto-model/opus-mt-en-gmw" --handler-config=$CONFIG_APPLY

graph_net/torch/dtype_gen_passes/__init__.py

@@ -0,0 +1,3 @@
+from graph_net.torch.dtype_gen_passes.pass_base import (
+    DtypeGeneralizationPass as DtypeGeneralizationPass,
+)

graph_net/torch/dtype_gen_passes/dtype_generalization_pass.py

@@ -0,0 +1,144 @@
+"""
+Concrete implementation of dtype generalization pass.
+
+This pass converts tensor dtypes in FX Graph by:
+1. Converting placeholder nodes (inputs) to target dtype
+2. Converting get_attr nodes (weights) to target dtype, except preserved weights
+3. Inserting .to(dtype) calls where needed
+"""
+
+import torch
+import torch.fx as fx
+from graph_net.torch.dtype_gen_passes.pass_base import DtypeGeneralizationPass
+
+
+class ConcretePass(DtypeGeneralizationPass):
+    """
+    FX Graph pass that converts dtypes of tensors.
+
+    This pass modifies the graph to:
+    - Convert input tensors to target dtype
+    - Convert weight tensors to target dtype (except preserved weights)
+    - Insert dtype conversion nodes where necessary
+    """
+
+    def get_pass_name(self) -> str:
+        return f"dtype_generalization_{self.target_dtype}"
+
+    def need_rewrite(self, gm: fx.GraphModule) -> bool:
+        """
+        Check if graph has float32 tensors that need conversion.
+        """
+        for node in gm.graph.nodes:
+            if self._node_need_rewrite(node):
+                return True
+        return False
+
+    def _node_need_rewrite(self, node: fx.Node) -> bool:
+        """
+        Check if a specific node needs dtype conversion.
+
+        Args:
+            node: FX Node to check
+
+        Returns:
+            True if node should be rewritten
+        """
+        # Check placeholder nodes (inputs)
+        if node.op == "placeholder":
+            return self._is_float32_tensor(node)
+
+        # Check get_attr nodes (weights)
+        if node.op == "get_attr":
+            if self._is_float32_tensor(node):
+                # Only rewrite if not in preserve list
+                attr_name = str(node.target)
+                return not self.should_preserve_weight(attr_name)
+
+        return False
+
+    def rewrite(self, gm: fx.GraphModule) -> fx.GraphModule:
+        """
+        Rewrite the graph to convert dtypes.
+
+        Strategy:
+        1. For each placeholder (input), insert .to(target_dtype) after it
+        2. For each get_attr (weight), insert .to(target_dtype) if not preserved
+        3. Update the graph and recompile
+        """
+        new_graph = fx.Graph()
+        val_map = {}
+
+        def create_placeholder(node: fx.Node) -> fx.Node:
+            """Create a placeholder node with dtype conversion if needed."""
+            new_node = new_graph.node_copy(node, lambda x: val_map.get(x, x))
+            if self._is_float32_tensor(node):
+                return new_graph.call_method("to", args=(new_node, self.torch_dtype))
+            return new_node
+
+        def create_get_attr(node: fx.Node) -> fx.Node:
+            """Create a get_attr node with dtype conversion if needed."""
+            new_node = new_graph.node_copy(node, lambda x: val_map.get(x, x))
+            attr_name = str(node.target)
+            if self._is_float32_tensor(node) and not self.should_preserve_weight(
+                attr_name
+            ):
+                return new_graph.call_method("to", args=(new_node, self.torch_dtype))
+            return new_node
+
+        for node in gm.graph.nodes:
+            if node.op == "placeholder":
+                val_map[node] = create_placeholder(node)
+            elif node.op == "get_attr":
+                val_map[node] = create_get_attr(node)
+            else:
+                new_node = new_graph.node_copy(node, lambda x: val_map.get(x, x))
+                val_map[node] = new_node
+
+        # Replace the graph
+        gm.graph = new_graph
+        gm.recompile()
+
+        return gm
+
+    def _is_float32_tensor(self, node: fx.Node) -> bool:
+        """
+        Check if a node represents a float32 tensor.
+
+        Args:
+            node: FX Node to check
+
+        Returns:
+            True if node is a float32 tensor
+        """
+        # Check tensor_meta if available (most reliable)
+        if "tensor_meta" in node.meta:
+            tensor_meta = node.meta["tensor_meta"]
+            if hasattr(tensor_meta, "dtype"):
+                return tensor_meta.dtype == torch.float32
+
+        # For placeholder and get_attr nodes without metadata,
+        # we need to be conservative and only return True if explicitly float
+        if node.op in ("placeholder", "get_attr"):
+            # Check type annotation if available
+            if node.type is not None:
+                type_str = str(node.type).lower()
+
+                # Explicitly check for integer types - these should NOT be converted
+                integer_types = ["long", "int", "short", "byte", "bool"]
+                if any(int_type in type_str for int_type in integer_types):
+                    return False
+
+                # Only return True if explicitly a floating point tensor
+                # Check for explicit float types: FloatTensor, float32, float16, etc.
+                float_indicators = ["float", "double", "half", "bfloat"]
+                if any(
+                    float_indicator in type_str for float_indicator in float_indicators
+                ):
+                    return True
+
+                # For generic "Tensor" without explicit dtype, be conservative
+                # Don't assume it's float32 - it might be integer
+                return False
+
+        return False

graph_net/torch/dtype_gen_passes/dtype_generalization_pass_bfloat16.py

@@ -0,0 +1,35 @@
+"""
+Dtype generalization pass for bfloat16.
+
+This pass converts float32 tensors to bfloat16.
+"""
+
+from graph_net.torch.dtype_gen_passes.dtype_generalization_pass import (
+    ConcretePass as BaseConcretePass,
+)
+
+# Weights that must remain float32 for numerical stability
+FLOAT32_PRESERVED_WEIGHTS = {
+    "running_mean",
+    "running_var",
+    "num_batches_tracked",
+    "bn_parameters_weight",
+    "bn_parameters_bias",
+    "ln_parameters_weight",
+    "ln_parameters_bias",
+}
+
+
+class ConcretePass(BaseConcretePass):
+    """
+    FX Graph pass that converts dtypes to bfloat16.
+    """
+
+    def __init__(self, *args, **kwargs):
+        # Override target_dtype to bfloat16
+        super().__init__(
+            target_dtype="bfloat16",
+            preserve_weights=FLOAT32_PRESERVED_WEIGHTS,
+            *args,
+            **kwargs,
+        )

graph_net/torch/dtype_gen_passes/dtype_generalization_pass_float16.py

@@ -0,0 +1,35 @@
+"""
+Dtype generalization pass for float16.
+
+This pass converts float32 tensors to float16.
+"""
+
+from graph_net.torch.dtype_gen_passes.dtype_generalization_pass import (
+    ConcretePass as BaseConcretePass,
+)
+
+# Weights that must remain float32 for numerical stability
+FLOAT32_PRESERVED_WEIGHTS = {
+    "running_mean",
+    "running_var",
+    "num_batches_tracked",
+    "bn_parameters_weight",
+    "bn_parameters_bias",
+    "ln_parameters_weight",
+    "ln_parameters_bias",
+}
+
+
+class ConcretePass(BaseConcretePass):
+    """
+    FX Graph pass that converts dtypes to float16.
+    """
+
+    def __init__(self, *args, **kwargs):
+        # Override target_dtype to float16
+        super().__init__(
+            target_dtype="float16",
+            preserve_weights=FLOAT32_PRESERVED_WEIGHTS,
+            *args,
+            **kwargs,
+        )