diff --git a/graph_net/paddle/naive_graph_decomposer.py b/graph_net/paddle/graph_decomposer.py similarity index 100% rename from graph_net/paddle/naive_graph_decomposer.py rename to graph_net/paddle/graph_decomposer.py diff --git a/graph_net/subgraph_decompose_and_evaluation_step.py b/graph_net/subgraph_decompose_and_evaluation_step.py index 2d5da20f2..20ae6dbb2 100755 --- a/graph_net/subgraph_decompose_and_evaluation_step.py +++ b/graph_net/subgraph_decompose_and_evaluation_step.py @@ -186,7 +186,7 @@ def run_decomposer_for_single_model( "decorator_path": f"{graphnet_root}/graph_net/{framework}/extractor.py", "decorator_config": { "name": model_name, - "custom_extractor_path": f"{graphnet_root}/graph_net/{framework}/naive_graph_decomposer.py", + "custom_extractor_path": f"{graphnet_root}/graph_net/{framework}/graph_decomposer.py", "custom_extractor_config": { "output_dir": output_dir, "split_positions": split_positions, diff --git a/graph_net/test/chain_naive_graph_decomposer_test.sh b/graph_net/test/chain_naive_graph_decomposer_test.sh index 6b220a992..d6dc43392 100644 --- a/graph_net/test/chain_naive_graph_decomposer_test.sh +++ b/graph_net/test/chain_naive_graph_decomposer_test.sh @@ -11,7 +11,7 @@ decorator_config_json_str=$(cat < "$temp_model_list" +model_list="$GRAPH_NET_ROOT/graph_net/config/small100_torch_samples_list.txt" python3 -m graph_net.torch.typical_sequence_split_points \ - --model-list "$temp_model_list" \ + --model-list "$model_list" \ --device "cuda" \ --window-size 10 \ + --fold-policy default \ + --fold-times 10 \ --output-json "$DECOMPOSE_PATH/split_results.json" -while IFS= read -r MODEL_PATH_IN_SAMPLES; do - if [[ -n "$MODEL_PATH_IN_SAMPLES" ]]; then - MODEL_FULL_PATH="$GRAPH_NET_ROOT/$MODEL_PATH_IN_SAMPLES" - MODEL_NAME=$(basename "$MODEL_PATH_IN_SAMPLES") - - echo "== Decomposing $MODEL_PATH_IN_SAMPLES. ==" - - decomposer_config_json_str=$(cat < "$DECOMPOSE_PATH/${MODEL_NAME}_validation.log" 2>&1 +python3 -m graph_net.model_path_handler \ + --model-path-list $model_list \ + --handler-config=$DECOMPOSE_CONFIG \ + --use-subprocess - echo "== Finished processing $MODEL_PATH_IN_SAMPLES. ==" - fi -done < $temp_model_list - -rm -f "$temp_model_list" +test_compiler_config_json_str=$(cat <> $DECOMPOSE_PATH/combined.log +python3 -m graph_net.torch.test_compiler \ + --allow-list $model_list \ + --compiler range_decomposer_validator \ + --device cuda \ + --config $TEST_COMPILER_CONFIG \ + --model-path-prefix $GRAPH_NET_ROOT \ + > "$DECOMPOSE_PATH/validation.log" 2>&1 python3 -m graph_net.plot_ESt \ - --benchmark-path "$DECOMPOSE_PATH/combined.log" \ + --benchmark-path "$DECOMPOSE_PATH/validation.log" \ --output-dir "$DECOMPOSE_PATH" \ No newline at end of file diff --git a/graph_net/torch/backend/range_decomposer_backend.py b/graph_net/torch/backend/range_decomposer_backend.py index 56a044cf8..b265d660b 100644 --- a/graph_net/torch/backend/range_decomposer_backend.py +++ b/graph_net/torch/backend/range_decomposer_backend.py @@ -44,7 +44,7 @@ def __call__(self, model: torch.nn.Module) -> torch.nn.Module: "decorator_config": { "name": model_name, "custom_extractor_path": str( - self.graph_net_root / "torch/naive_graph_decomposer.py" + self.graph_net_root / "torch/graph_decomposer.py" ), "custom_extractor_config": { "output_dir": str(model_output_dir), diff --git a/graph_net/torch/fully_fusible_subgraph_extractor.py b/graph_net/torch/fully_fusible_subgraph_extractor.py index 73c06c488..13384650a 100644 --- a/graph_net/torch/fully_fusible_subgraph_extractor.py +++ b/graph_net/torch/fully_fusible_subgraph_extractor.py @@ -86,7 +86,7 @@ def _build_decompose_config( "decorator_path": f"{graph_net_root}/torch/extractor.py", "decorator_config": { "name": f"{self.name}", - "custom_extractor_path": f"{graph_net_root}/torch/naive_graph_decomposer.py", + "custom_extractor_path": f"{graph_net_root}/torch/graph_decomposer.py", "custom_extractor_config": { "output_dir": temp_dir, "split_positions": self.config["split_positions"], diff --git a/graph_net/torch/naive_graph_decomposer.py b/graph_net/torch/graph_decomposer.py similarity index 71% rename from graph_net/torch/naive_graph_decomposer.py rename to graph_net/torch/graph_decomposer.py index 7a298c974..2723f59bf 100644 --- a/graph_net/torch/naive_graph_decomposer.py +++ b/graph_net/torch/graph_decomposer.py @@ -1,5 +1,6 @@ import os import torch +import json from graph_net.torch.decompose_util import convert_to_submodules_graph from graph_net.torch.extractor import GraphExtractor as BuiltinGraphExtractor import graph_net.imp_util as imp_util @@ -7,6 +8,12 @@ from graph_net.torch.fx_graph_parse_util import parse_sole_graph_module +def load_json(file_path): + with open(file_path, "r", encoding="utf-8") as file: + data_dict = json.load(file) + return data_dict + + class GraphExtractor: """ Used by graph_net.torch.run_model @@ -151,6 +158,83 @@ def fn(submodule, seq_no): return fn +class RangeDecomposerExtractor: + """ + Used by graph_net.model_path_handler + """ + + def __init__(self, config: dict = None): + if config is None: + config = {} + self.config = self._make_config(**config) + + def _make_config( + self, + split_results_path=None, + group_head_and_tail=False, + chain_style=False, + output_dir="./tmp/naive_decomposer_dir", + filter_path=None, + filter_config=None, + post_extract_process_path=None, + post_extract_process_class_name=None, + post_extract_process_config=None, + model_path_prefix="", + **kwargs, + ): + if os.path.isfile(split_results_path) and split_results_path.endswith(".json"): + pass + else: + raise ValueError( + f"split_results_path should be a valid JSON file path, but got {split_results_path=}" + ) + if post_extract_process_config is None: + post_extract_process_config = {} + return { + "split_results_path": split_results_path, + "group_head_and_tail": group_head_and_tail, + "chain_style": chain_style, + "output_dir": output_dir, + "filter_path": filter_path, + "filter_config": filter_config if filter_config is not None else {}, + "post_extract_process_path": post_extract_process_path, + "post_extract_process_class_name": post_extract_process_class_name, + "post_extract_process_config": post_extract_process_config, + "model_path_prefix": model_path_prefix, + } + + def __call__(self, rel_model_path): + model_path = os.path.join(self.config["model_path_prefix"], rel_model_path) + split_results = load_json(self.config["split_results_path"]) + split_positions = split_results[os.path.basename(rel_model_path)][ + "split_points" + ] + config = { + "split_positions": split_positions, + "group_head_and_tail": self.config.get("group_head_and_tail", False), + "chain_style": self.config.get("chain_style", False), + } + module, inputs = get_torch_module_and_inputs(model_path) + gm = parse_sole_graph_module(module, inputs) + rewrited_gm: torch.fx.GraphModule = convert_to_submodules_graph( + gm, + submodule_hook=self.get_naive_decomposer_extractor(model_path), + **config, + ) + rewrited_gm(*inputs) + + def get_naive_decomposer_extractor(self, model_path): + def fn(submodule, seq_no): + return NaiveDecomposerExtractorModule( + config=self.config, + parent_graph_name=os.path.basename(model_path), + submodule=submodule, + seq_no=seq_no, + ) + + return fn + + class NaiveDecomposerExtractorModule(torch.nn.Module): def __init__( self, diff --git a/graph_net/torch/typical_sequence_split_points.py b/graph_net/torch/typical_sequence_split_points.py index 9a8fb6445..e158218c7 100644 --- a/graph_net/torch/typical_sequence_split_points.py +++ b/graph_net/torch/typical_sequence_split_points.py @@ -82,8 +82,12 @@ def get_input_dict(self, model_path: str, device: str) -> Dict[str, torch.Tensor class SplitAnalyzer: - def __init__(self, window_size: int = 10): + def __init__( + self, window_size: int = 10, fold_policy: str = "default", fold_times: int = 0 + ): self.window_size = window_size + self.fold_policy = fold_policy + self.fold_times = fold_times def _resolve_token_to_ops( self, tid, num_primitives, token_id2primitive_id, symbol_map @@ -169,7 +173,9 @@ def analyze(self, model_paths_file: str, device: str) -> Dict[str, Dict]: return {} rp_parser = RpExprParser( - window_size=self.window_size, fold_policy="default", fold_times=0 + window_size=self.window_size, + fold_policy=self.fold_policy, + fold_times=self.fold_times, ) rp_expr, token_id2primitive_id = rp_parser(inputs_seqs) rp_expr.try_unwrap_body_of_sole_symbol_token() @@ -253,7 +259,11 @@ def _print_analysis(self, name, path, splits, total_len, full_ops): def main(args): - analyzer = SplitAnalyzer(window_size=args.window_size) + analyzer = SplitAnalyzer( + window_size=args.window_size, + fold_policy=args.fold_policy, + fold_times=args.fold_times, + ) results = analyzer.analyze(args.model_list, args.device) if args.output_json: with open(args.output_json, "w") as f: @@ -279,6 +289,18 @@ def main(args): parser.add_argument( "--window-size", type=int, default=10, help="Window size for RP Parser." ) + parser.add_argument( + "--fold-policy", + type=str, + default="default", + help="Policy for split analysis, one of 'default' or 'longest'", + ) + parser.add_argument( + "--fold-times", + type=int, + default=0, + help="How many times to fold tokens. If 0, then no folding is done.", + ) parser.add_argument( "--output-json", type=str,