【Hackathon 9th Sprint No.87】feat: implement MultiDtypeGenerator for low-precision sample generation (#396)

* feat: implement MultiDtypeGenerator for low-precision sample generation

- Add MultiDtypeGenerator class to generate float16/bfloat16/float8 samples from float32 samples
- Automatically convert input and weight tensor dtypes
- Keep batch_norm parameters (running_mean, running_var, scale, bias) as float32
- Add torch.autocast context manager to model.py forward method
- Update graph_net.json metadata with dtype information
- Add sample validation mechanism
- Add MultiDtypeFilter for filtering invalid graphs
- Add command-line tool for batch generation
- Support test_compiler evaluation and Agent code generation workflows

* refactor: optimize MultiDtypeGenerator code style and reduce nesting

- Refactor _modify_weight_meta to use single-pass traversal
- Extract _should_keep_weight_float32 helper method
- Simplify _ensure_autocast_import to reduce duplicate file writes
- Clean up unused variables in _validate_sample
- Apply black code formatting

* feat: implement MultiDtypeGenerator with FX Graph passes

- Add MultiDtypeGenerator for generating low-precision (float16/bfloat16/float8) samples
- Implement FX Graph passes for dtype conversion (no AST/string manipulation)
- Preserve BatchNorm/LayerNorm buffers as float32 for numerical stability
- Add autocast metadata for runtime mixed-precision support
- Thread-safe extraction with proper locking
- Comprehensive error handling and device detection
- Validated on CV (ResNet18) and NLP (BERT) samples

Architecture:
- graph_net/torch/multi_dtype_generator.py: Main generator with RunModelDecorator
- graph_net/torch/multi_dtype_passes/: FX Graph pass implementations
  - pass_base.py: Base class for dtype conversion passes
  - dtype_conversion_pass.py: Concrete dtype conversion implementation
  - autocast_wrapper_pass.py: Autocast metadata injection
  - pass_mgr.py: Pass manager

Features:
- Real FX Graph passes (not string replacement)
- Smart weight preservation (BN/LN buffers stay float32)
- Automatic device type detection
- Compatible with test_compiler and Agent workflows
- Extensible filter system for unsupported operations

Fixes:
- Correct val_map semantics in FX Graph rewriting
- Proper dtype detection (avoid false positives on int tensors)
- CPU fallback when CUDA unavailable
- Import organization and error handling

* fix: resolve ruff linter errors in multi_dtype_passes

- Add explicit re-export in __init__.py for DtypeConversionPass
- Use TYPE_CHECKING for forward reference in pass_mgr.py
- Remove unused torch import in autocast_wrapper_pass.py
- Fix type annotation to use type[DtypeConversionPass]

* feat: add data type generalization passes with FX Graph architecture

- Add InitDataTypeGeneralizationPasses class for testing and initializing dtype conversion passes
- Implement DtypeConversionPass base class and ConcretePass for dtype conversion
- Add AutocastWrapperPass for handling operators not supporting low precision
- Store applicable pass names in graph_net.json under kDataTypeGeneralizationPasses
- Use FX Graph passes instead of AST/string manipulation for robustness
- Add atomic file writing for graph_net.json to prevent corruption
- Add comprehensive weight preservation logic for BatchNorm/LayerNorm parameters
- Add RunModelPredicator integration for testing graph runnability
- Refine exception handling with specific exception types
- Auto-detect device type for autocast (cuda/cpu)
- Organize passes in multi_dtype_passes/ directory with proper management

This feature enables generating low-precision (float16/bfloat16/float8) samples
from float32 computation graphs, following the initializer pattern where passes
are tested offline and applied dynamically at runtime.

* Refactor dtype generalization to two-step architecture per reviewer feedback

* Fix integer tensor type detection in dtype generalization pass

* Resolve merge conflicts: merge dtype and develop branches

* Fix code formatting: apply black formatter

* Fix code formatting: add missing blank line

* Refactor: fix update_json API and add dtype generalization field constants

* Fix: create separate pass files for each dtype to match graph_net.json

* Refactor: simplify test script and extract helper functions in dtype pass

Author: Dayuxiaoshui

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,
+        )