From f81f3484e586d2726d21c7f7b1fb98ce2def458c Mon Sep 17 00:00:00 2001 From: WHoutstanding Date: Mon, 2 Feb 2026 12:40:03 +0800 Subject: [PATCH 1/4] fix output_dir of dtype generalization pass --- graph_net/test/dtype_gen_test.sh | 22 ------------- .../torch/sample_pass/dtype_generalizer.py | 33 +++++++++++++++---- 2 files changed, 26 insertions(+), 29 deletions(-) diff --git a/graph_net/test/dtype_gen_test.sh b/graph_net/test/dtype_gen_test.sh index 31528ef5a..c8283b8ae 100755 --- a/graph_net/test/dtype_gen_test.sh +++ b/graph_net/test/dtype_gen_test.sh @@ -43,25 +43,3 @@ EOF ) -# Step 3: Valiation -SUCCESS_CNT=0 -FAIL_CNT=0 - -for model_path in "$OUTPUT_DIR"/*; do - echo "[VALIDATE] $model_path" - - output=$(python -m graph_net.torch.validate \ - --model-path "$model_path" 2>&1) - - if echo "$output" | grep -q "Validation success, model_path="; then - echo "SUCCESS" - ((SUCCESS_CNT++)) - else - echo "FAIL" - ((FAIL_CNT++)) - fi -done - -echo "====================" -echo "SUCCESS $SUCCESS_CNT" -echo "FAIL $FAIL_CNT" \ No newline at end of file diff --git a/graph_net/torch/sample_pass/dtype_generalizer.py b/graph_net/torch/sample_pass/dtype_generalizer.py index 3ed56b32e..6258bfef0 100644 --- a/graph_net/torch/sample_pass/dtype_generalizer.py +++ b/graph_net/torch/sample_pass/dtype_generalizer.py @@ -308,17 +308,20 @@ def resume(self, model_path: str) -> List[str]: """ # Apply model_path_prefix if provided if self.model_path_prefix: - model_path = str(Path(self.model_path_prefix) / model_path) + model_path_abs = 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) + dtype_pass_names = self._read_dtype_pass_names(model_path_abs) if not dtype_pass_names: - logging.warning(f"No dtype passes found in {model_path}/graph_net.json") + logging.warning(f"No dtype passes found in {model_path_abs}/graph_net.json") return [] # Parse the computation graph - traced_model = parse_immutable_model_path_into_sole_graph_module(model_path) + traced_model = parse_immutable_model_path_into_sole_graph_module(model_path_abs) + + # Copy the originl sample + self._copy_sample(model_path) # Generate samples for each pass generated_samples = [] @@ -388,8 +391,7 @@ def _apply_pass_and_generate( gm_modified = dtype_pass.rewrite(gm_copy) # Generate output directory - model_name = Path(model_path).name - output_sample_dir = Path(self.output_dir) / f"{model_name}_{dtype}" + output_sample_dir = Path(self.output_dir) / dtype / model_path output_sample_dir.mkdir(parents=True, exist_ok=True) # Write modified model.py @@ -399,7 +401,7 @@ def _apply_pass_and_generate( f.write(write_code) # Copy metadata files - for fname in ["graph_net.json", "weight_meta.py", "input_meta.py"]: + for fname in ["graph_net.json", "weight_meta.py", "input_meta.py", "input_tensor_constraints.py", "graph_hash.txt"]: src = Path(model_path) / fname if src.exists(): shutil.copy(src, output_sample_dir / fname) @@ -429,6 +431,23 @@ def _update_sample_metadata(self, sample_dir: Path, dtype: str) -> None: update_json(graph_net_json_path, kDtypeGeneralizationPrecision, dtype) update_json(graph_net_json_path, kDtypeGeneralizationGenerated, True) + def _copy_sample(self, model_path: str) -> None: + """ + Copy a sample. + + Args: + model_path: Original model path + """ + # Generate output directory + output_sample_dir = Path(self.output_dir) / f"float32" / model_path + output_sample_dir.mkdir(parents=True, exist_ok=True) + + # Copy all files + for fname in ["graph_net.json", "weight_meta.py", "input_meta.py", "model.py", + "input_tensor_constraints.py", "graph_hash.txt"]: + src = Path(model_path) / fname + if src.exists(): + shutil.copy(src, output_sample_dir / fname) class MultiDtypeFilter: """ From bed396ca358be90a9d97fa62c28c06e771998e3e Mon Sep 17 00:00:00 2001 From: WHoutstanding Date: Mon, 2 Feb 2026 12:45:53 +0800 Subject: [PATCH 2/4] fix CodeStyle --- .../torch/sample_pass/dtype_generalizer.py | 21 +++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/graph_net/torch/sample_pass/dtype_generalizer.py b/graph_net/torch/sample_pass/dtype_generalizer.py index 6258bfef0..1fd21e327 100644 --- a/graph_net/torch/sample_pass/dtype_generalizer.py +++ b/graph_net/torch/sample_pass/dtype_generalizer.py @@ -401,7 +401,13 @@ def _apply_pass_and_generate( f.write(write_code) # Copy metadata files - for fname in ["graph_net.json", "weight_meta.py", "input_meta.py", "input_tensor_constraints.py", "graph_hash.txt"]: + for fname in [ + "graph_net.json", + "weight_meta.py", + "input_meta.py", + "input_tensor_constraints.py", + "graph_hash.txt", + ]: src = Path(model_path) / fname if src.exists(): shutil.copy(src, output_sample_dir / fname) @@ -439,16 +445,23 @@ def _copy_sample(self, model_path: str) -> None: model_path: Original model path """ # Generate output directory - output_sample_dir = Path(self.output_dir) / f"float32" / model_path + output_sample_dir = Path(self.output_dir) / "float32" / model_path output_sample_dir.mkdir(parents=True, exist_ok=True) # Copy all files - for fname in ["graph_net.json", "weight_meta.py", "input_meta.py", "model.py", - "input_tensor_constraints.py", "graph_hash.txt"]: + for fname in [ + "graph_net.json", + "weight_meta.py", + "input_meta.py", + "model.py", + "input_tensor_constraints.py", + "graph_hash.txt", + ]: src = Path(model_path) / fname if src.exists(): shutil.copy(src, output_sample_dir / fname) + class MultiDtypeFilter: """ Filter for graphs that cannot support dtype generalization. From f0d8fb31b5a67230be3aeebe07c0ff5eb115df82 Mon Sep 17 00:00:00 2001 From: WHoutstanding Date: Mon, 2 Feb 2026 16:29:44 +0800 Subject: [PATCH 3/4] gen graph_hash.txt and subgraph_sources.json --- .../torch/sample_pass/dtype_generalizer.py | 80 +++++++++++++------ 1 file changed, 56 insertions(+), 24 deletions(-) diff --git a/graph_net/torch/sample_pass/dtype_generalizer.py b/graph_net/torch/sample_pass/dtype_generalizer.py index 1fd21e327..55feadd01 100644 --- a/graph_net/torch/sample_pass/dtype_generalizer.py +++ b/graph_net/torch/sample_pass/dtype_generalizer.py @@ -42,6 +42,8 @@ from graph_net.sample_pass.sample_pass import SamplePass from graph_net.sample_pass.resumable_sample_pass_mixin import ResumableSamplePassMixin +from graph_net.hash_util import get_sha256_hash + # Weights that must remain float32 for numerical stability FLOAT32_PRESERVED_WEIGHTS = { "running_mean", @@ -296,7 +298,7 @@ def sample_handled(self, rel_model_path: str) -> bool: def __call__(self, rel_model_path: str): self.resumable_handle_sample(rel_model_path) - def resume(self, model_path: str) -> List[str]: + def resume(self, rel_model_path: str) -> List[str]: """ Apply dtype passes to generate new samples. @@ -308,27 +310,36 @@ def resume(self, model_path: str) -> List[str]: """ # Apply model_path_prefix if provided if self.model_path_prefix: - model_path_abs = str(Path(self.model_path_prefix) / model_path) + abs_model_path = str(Path(self.model_path_prefix) / rel_model_path) # Read pass names from graph_net.json - dtype_pass_names = self._read_dtype_pass_names(model_path_abs) + dtype_pass_names = self._read_dtype_pass_names(abs_model_path) if not dtype_pass_names: - logging.warning(f"No dtype passes found in {model_path_abs}/graph_net.json") + logging.warning(f"No dtype passes found in {abs_model_path}/graph_net.json") return [] # Parse the computation graph - traced_model = parse_immutable_model_path_into_sole_graph_module(model_path_abs) + traced_model = parse_immutable_model_path_into_sole_graph_module(abs_model_path) # Copy the originl sample - self._copy_sample(model_path) + files_copied = [ + "model.py", + "graph_hash.txt", + "graph_net.json", + "weight_meta.py", + "input_meta.py", + "input_tensor_constraints.py", + "subgraph_sources.json", + ] + self._copy_sample_files(rel_model_path, "float32", files_copied) # Generate samples for each pass generated_samples = [] for pass_name in dtype_pass_names: try: sample_dir = self._apply_pass_and_generate( - model_path, traced_model, pass_name + rel_model_path, traced_model, pass_name ) generated_samples.append(sample_dir) logging.info(f"Generated sample: {sample_dir}") @@ -400,17 +411,20 @@ def _apply_pass_and_generate( with open(output_sample_dir / "model.py", "w") as f: f.write(write_code) + # Write modified graph_hash.txt + model_hash = get_sha256_hash(model_code) + with open(output_sample_dir / "graph_hash.txt", "w") as f: + f.write(model_hash) + # Copy metadata files - for fname in [ + files_copied = [ "graph_net.json", "weight_meta.py", "input_meta.py", "input_tensor_constraints.py", - "graph_hash.txt", - ]: - src = Path(model_path) / fname - if src.exists(): - shutil.copy(src, output_sample_dir / fname) + "subgraph_sources.json", + ] + self._copy_sample_files(model_path, dtype, files_copied) # Update graph_net.json with dtype information self._update_sample_metadata(output_sample_dir, dtype) @@ -437,7 +451,32 @@ def _update_sample_metadata(self, sample_dir: Path, dtype: str) -> None: update_json(graph_net_json_path, kDtypeGeneralizationPrecision, dtype) update_json(graph_net_json_path, kDtypeGeneralizationGenerated, True) - def _copy_sample(self, model_path: str) -> None: + # def _copy_sample(self, model_path: str) -> None: + # """ + # Copy a sample. + + # Args: + # model_path: Original model path + # """ + # # Generate output directory + # output_sample_dir = Path(self.output_dir) / "float32" / model_path + # output_sample_dir.mkdir(parents=True, exist_ok=True) + + # # Copy all files + # for fname in [ + # "graph_net.json", + # "weight_meta.py", + # "input_meta.py", + # "model.py", + # "input_tensor_constraints.py", + # "graph_hash.txt", + # ]: + # src = Path(model_path) / fname + # if src.exists(): + # shutil.copy(src, output_sample_dir / fname) + def _copy_sample_files( + self, rel_model_path: str, dtype: str, files_copied: list + ) -> None: """ Copy a sample. @@ -445,19 +484,12 @@ def _copy_sample(self, model_path: str) -> None: model_path: Original model path """ # Generate output directory - output_sample_dir = Path(self.output_dir) / "float32" / model_path + output_sample_dir = Path(self.output_dir) / dtype / rel_model_path output_sample_dir.mkdir(parents=True, exist_ok=True) # Copy all files - for fname in [ - "graph_net.json", - "weight_meta.py", - "input_meta.py", - "model.py", - "input_tensor_constraints.py", - "graph_hash.txt", - ]: - src = Path(model_path) / fname + for fname in files_copied: + src = Path(rel_model_path) / fname if src.exists(): shutil.copy(src, output_sample_dir / fname) From a24934d03f516190cdf4fbf28517841128d29b19 Mon Sep 17 00:00:00 2001 From: WHoutstanding Date: Mon, 2 Feb 2026 16:34:18 +0800 Subject: [PATCH 4/4] fix comment of dtype_generalizer.py --- .../torch/sample_pass/dtype_generalizer.py | 29 ++----------------- 1 file changed, 3 insertions(+), 26 deletions(-) diff --git a/graph_net/torch/sample_pass/dtype_generalizer.py b/graph_net/torch/sample_pass/dtype_generalizer.py index 55feadd01..a8fec4d3f 100644 --- a/graph_net/torch/sample_pass/dtype_generalizer.py +++ b/graph_net/torch/sample_pass/dtype_generalizer.py @@ -451,43 +451,20 @@ def _update_sample_metadata(self, sample_dir: Path, dtype: str) -> None: update_json(graph_net_json_path, kDtypeGeneralizationPrecision, dtype) update_json(graph_net_json_path, kDtypeGeneralizationGenerated, True) - # def _copy_sample(self, model_path: str) -> None: - # """ - # Copy a sample. - - # Args: - # model_path: Original model path - # """ - # # Generate output directory - # output_sample_dir = Path(self.output_dir) / "float32" / model_path - # output_sample_dir.mkdir(parents=True, exist_ok=True) - - # # Copy all files - # for fname in [ - # "graph_net.json", - # "weight_meta.py", - # "input_meta.py", - # "model.py", - # "input_tensor_constraints.py", - # "graph_hash.txt", - # ]: - # src = Path(model_path) / fname - # if src.exists(): - # shutil.copy(src, output_sample_dir / fname) def _copy_sample_files( self, rel_model_path: str, dtype: str, files_copied: list ) -> None: """ - Copy a sample. + Copy files of sample. Args: - model_path: Original model path + rel_model_path: relative model path """ # Generate output directory output_sample_dir = Path(self.output_dir) / dtype / rel_model_path output_sample_dir.mkdir(parents=True, exist_ok=True) - # Copy all files + # Copy files of original sample for fname in files_copied: src = Path(rel_model_path) / fname if src.exists():