Fix integer tensor type detection in dtype generalization pass

Dayuxiaoshui · Dayuxiaoshui · commit 7da86c651ba1 · 2025-12-08T10:00:01.000Z
diff --git a/graph_net/graph_net_json_file_util.py b/graph_net/graph_net_json_file_util.py
@@ -1,2 +1,33 @@
 kDimensionGeneralizationPasses = "dimension_generalization_passes"
 kDataTypeGeneralizationPasses = "data_type_generalization_passes"
+
+import json
+from pathlib import Path
+from typing import Union
+
+
+def update_json(json_path: Union[str, Path], updates: dict) -> None:
+    """
+    Atomically update a JSON file with the given updates.
+
+    Args:
+        json_path: Path to the JSON file
+        updates: Dictionary of key-value pairs to update
+    """
+    json_path = Path(json_path)
+    
+    # Read existing JSON
+    if json_path.exists():
+        with open(json_path, "r") as f:
+            metadata = json.load(f)
+    else:
+        metadata = {}
+    
+    # Apply updates
+    metadata.update(updates)
+    
+    # Atomic write: write to temp file then rename
+    temp_path = json_path.with_suffix(".json.tmp")
+    with open(temp_path, "w") as f:
+        json.dump(metadata, f, indent=4)
+    temp_path.replace(json_path)
diff --git a/graph_net/test/dtype_gen_test.sh b/graph_net/test/dtype_gen_test.sh
@@ -31,27 +31,28 @@ echo ""
 echo "[1/2] Testing CV model: timm/resnet18"
 config_json_str_init=$(cat <<EOF
 {
-    "decorator_path": "$GRAPH_NET_ROOT/torch/dtype_generalizer.py",
-    "decorator_class_name": "InitDataTypeGeneralizationPasses",
-    "decorator_config": {
-        "dtype_list": ["float16", "bfloat16"]
+    "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.torch.run_model \
-    --model-path "$SAMPLES_ROOT/timm/resnet18" \
-    --decorator-config="$CONFIG_INIT" || echo "Warning: CV model test failed"
+python3 -m graph_net.model_path_handler \
+    --model-path "timm/resnet18" \
+    --handler-config="$CONFIG_INIT" || echo "Warning: CV model test failed"
 
 echo ""
 
 # Test on an NLP model (BERT-like)
 echo "[2/2] Testing NLP model: transformers-auto-model/opus-mt-en-gmw"
-python3 -m graph_net.torch.run_model \
-    --model-path "$SAMPLES_ROOT/transformers-auto-model/opus-mt-en-gmw" \
-    --decorator-config="$CONFIG_INIT" || echo "Warning: NLP model test failed"
+python3 -m graph_net.model_path_handler \
+    --model-path "transformers-auto-model/opus-mt-en-gmw" \
+    --handler-config="$CONFIG_INIT" || echo "Warning: NLP model test failed"
 
 echo ""
 echo "Step 1 completed. Pass names written to graph_net.json"
@@ -67,69 +68,37 @@ echo ""
 
 config_json_str_apply=$(cat <<EOF
 {
-    "decorator_path": "$GRAPH_NET_ROOT/torch/dtype_generalizer.py",
-    "decorator_class_name": "ApplyDataTypeGeneralizationPasses",
-    "decorator_config": {
-        "output_dir": "$OUTPUT_DIR"
+    "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)
 
 echo "[1/2] Generating CV samples..."
-python3 -m graph_net.torch.run_model \
-    --model-path "$SAMPLES_ROOT/timm/resnet18" \
-    --decorator-config="$CONFIG_APPLY" || echo "Warning: CV generation failed"
+python3 -m graph_net.model_path_handler \
+    --model-path "timm/resnet18" \
+    --handler-config="$CONFIG_APPLY" || echo "Warning: CV generation failed"
 
 echo ""
 
 echo "[2/2] Generating NLP samples..."
-python3 -m graph_net.torch.run_model \
-    --model-path "$SAMPLES_ROOT/transformers-auto-model/opus-mt-en-gmw" \
-    --decorator-config="$CONFIG_APPLY" || echo "Warning: NLP generation failed"
+python3 -m graph_net.model_path_handler \
+    --model-path "transformers-auto-model/opus-mt-en-gmw" \
+    --handler-config="$CONFIG_APPLY" || echo "Warning: NLP generation failed"
 
 echo ""
 echo "Step 2 completed. Generated samples in: $OUTPUT_DIR"
-echo ""
-
-# ============================================
-# Verification
-# ============================================
-echo "=========================================="
-echo "Verification"
-echo "=========================================="
-echo ""
-
-if [ -d "$OUTPUT_DIR" ]; then
-    echo "Generated samples:"
-    ls -lh "$OUTPUT_DIR"
-    echo ""
-    
-    # Count generated samples
-    SAMPLE_COUNT=$(find "$OUTPUT_DIR" -mindepth 1 -maxdepth 1 -type d | wc -l)
-    echo "Total samples generated: $SAMPLE_COUNT"
-    
-    if [ $SAMPLE_COUNT -gt 0 ]; then
-        echo ""
-        echo "✓ Test PASSED: Successfully generated $SAMPLE_COUNT low-precision samples"
-        echo ""
-        echo "You can now use these samples for:"
-        echo "  - test_compiler evaluation"
-        echo "  - Agent code generation"
-        echo "  - Performance benchmarking"
-    else
-        echo ""
-        echo "✗ Test WARNING: No samples were generated"
-        echo "  This might be normal if models don't support dtype conversion"
-    fi
-else
-    echo "✗ Test FAILED: Output directory not created"
-    exit 1
-fi
-
 echo ""
 echo "=========================================="
 echo "Test Complete"
 echo "=========================================="
-
diff --git a/graph_net/torch/dtype_gen_passes/dtype_generalization_pass.py b/graph_net/torch/dtype_gen_passes/dtype_generalization_pass.py
@@ -129,16 +129,25 @@ def _is_float32_tensor(self, node: fx.Node) -> bool:
                 return tensor_meta.dtype == torch.float32
 
         # For placeholder and get_attr nodes without metadata,
-        # conservatively assume they might be float32
-        # This is safe because:
-        # 1. .to() on non-float tensors is a no-op for most cases
-        # 2. Integer tensors (like input_ids) won't be affected
+        # 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)
-                # Only return True if it's explicitly a floating point tensor
-                if "Tensor" in type_str and "int" not in type_str.lower():
+                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
diff --git a/graph_net/torch/dtype_generalizer.py b/graph_net/torch/dtype_generalizer.py
@@ -19,14 +19,18 @@
 
 import torch.fx as fx
 
-from graph_net.graph_net_json_file_util import kDataTypeGeneralizationPasses
+from graph_net.graph_net_json_file_util import (
+    kDataTypeGeneralizationPasses,
+    update_json,
+)
 from graph_net.torch.constraint_util import RunModelPredicator
 from graph_net.torch.fx_graph_cache_util import (
     parse_immutable_model_path_into_sole_graph_module,
 )
 from graph_net.torch.fx_graph_serialize_util import serialize_graph_module_to_str
 from graph_net.torch.dtype_gen_passes.pass_mgr import get_dtype_generalization_pass
 from graph_net.torch import utils
+from graph_net.imp_util import load_module
 
 
 # Weights that must remain float32 for numerical stability
@@ -51,12 +55,14 @@ class InitDataTypeGeneralizationPasses:
     Config format:
         {
             "dtype_list": ["float16", "bfloat16"],
+            "model_path_prefix": "",
         }
     """
 
     def __init__(self, config: Dict[str, Any]):
         self.config = config
         self.dtype_list = config.get("dtype_list", ["float16", "bfloat16"])
+        self.model_path_prefix = config.get("model_path_prefix", "")
 
         # Validate dtypes
         valid_dtypes = {"float16", "bfloat16", "float8"}
@@ -71,8 +77,12 @@ def __call__(self, model_path: str) -> None:
         Initialize dtype passes for the given model.
 
         Args:
-            model_path: Path to the model directory
+            model_path: Path to the model directory (may be relative to model_path_prefix)
         """
+        # Apply model_path_prefix if provided
+        if self.model_path_prefix:
+            model_path = str(Path(self.model_path_prefix) / model_path)
+        
         # Parse the computation graph
         traced_model = parse_immutable_model_path_into_sole_graph_module(model_path)
 
@@ -176,18 +186,7 @@ def _save_dtype_pass_names(
             model_path: Path to model directory
         """
         graph_net_json_path = Path(model_path) / "graph_net.json"
-
-        with open(graph_net_json_path, "r") as f:
-            metadata = json.load(f)
-
-        metadata[kDataTypeGeneralizationPasses] = dtype_pass_names
-
-        # Atomic write: write to temp file then rename
-        temp_path = graph_net_json_path.with_suffix(".json.tmp")
-        with open(temp_path, "w") as f:
-            json.dump(metadata, f, indent=4)
-
-        temp_path.replace(graph_net_json_path)
+        update_json(graph_net_json_path, {kDataTypeGeneralizationPasses: dtype_pass_names})
 
 
 class ApplyDataTypeGeneralizationPasses:
@@ -200,6 +199,10 @@ class ApplyDataTypeGeneralizationPasses:
     Config format:
         {
             "output_dir": "/path/to/output",
+            "model_path_prefix": "",
+            "model_runnable_predicator_filepath": "...",
+            "model_runnable_predicator_class_name": "...",
+            "model_runnable_predicator_config": {...},
         }
     """
 
@@ -208,17 +211,41 @@ def __init__(self, config: Dict[str, Any]):
         self.output_dir = config.get("output_dir")
         if not self.output_dir:
             raise ValueError("output_dir is required in config")
+        
+        self.model_path_prefix = config.get("model_path_prefix", "")
+        
+        # model_runnable_predicator is required to ensure generated code is runnable
+        if "model_runnable_predicator_filepath" not in config:
+            raise ValueError(
+                "model_runnable_predicator_filepath is required in config. "
+                "Generated code must be validated."
+            )
+        self.model_runnable_predicator = self._make_model_runnable_predicator(config)
+    
+    def _make_model_runnable_predicator(self, config: Dict[str, Any]):
+        """Create model runnable predicator from config."""
+        module = load_module(config["model_runnable_predicator_filepath"])
+        cls = getattr(
+            module,
+            config.get("model_runnable_predicator_class_name", "RunModelPredicator"),
+        )
+        predicator_config = config.get("model_runnable_predicator_config", {})
+        return cls(predicator_config)
 
     def __call__(self, model_path: str) -> List[str]:
         """
         Apply dtype passes to generate new samples.
 
         Args:
-            model_path: Path to the original model directory
+            model_path: Path to the original model directory (may be relative to model_path_prefix)
 
         Returns:
             List of generated sample directories
         """
+        # Apply model_path_prefix if provided
+        if self.model_path_prefix:
+            model_path = str(Path(self.model_path_prefix) / model_path)
+        
         # Read pass names from graph_net.json
         dtype_pass_names = self._read_dtype_pass_names(model_path)
 
@@ -316,6 +343,13 @@ def _apply_pass_and_generate(
         # Update graph_net.json with dtype information
         self._update_sample_metadata(output_sample_dir, dtype)
 
+        # Validate generated sample (required - generated code must be runnable)
+        if not self.model_runnable_predicator(str(output_sample_dir)):
+            raise RuntimeError(
+                f"Generated sample failed validation: {output_sample_dir}"
+            )
+        logging.info(f"Generated sample validated: {output_sample_dir}")
+
         return str(output_sample_dir)
 
     def _update_sample_metadata(self, sample_dir: Path, dtype: str) -> None:
@@ -327,20 +361,14 @@ def _update_sample_metadata(self, sample_dir: Path, dtype: str) -> None:
             dtype: Target dtype
         """
         graph_net_json_path = sample_dir / "graph_net.json"
-
-        with open(graph_net_json_path, "r") as f:
-            metadata = json.load(f)
-
-        # Add dtype information
-        metadata["dtype"] = dtype
-        metadata["precision"] = dtype
-        metadata["generated_from_dtype_generalization"] = True
-
-        # Atomic write
-        temp_path = graph_net_json_path.with_suffix(".json.tmp")
-        with open(temp_path, "w") as f:
-            json.dump(metadata, f, indent=4)
-        temp_path.replace(graph_net_json_path)
+        update_json(
+            graph_net_json_path,
+            {
+                "dtype": dtype,
+                "precision": dtype,
+                "generated_from_dtype_generalization": True,
+            },
+        )
 
 
 class MultiDtypeFilter:
diff --git a/samples/timm/resnet18/graph_net.json b/samples/timm/resnet18/graph_net.json
@@ -3,5 +3,9 @@
     "num_devices_required": 1,
     "num_nodes_required": 1,
     "source": "timm",
-    "heuristic_tag": "computer_vision"
+    "heuristic_tag": "computer_vision",
+    "data_type_generalization_passes": [
+        "dtype_generalization_pass_float16",
+        "dtype_generalization_pass_bfloat16"
+    ]
 }
