Skip to content

Add backward graph extraction for fusible and typical subgraphs #675

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
fangfangssj wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

Add backward graph extraction for fusible and typical subgraphs #675

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
89 changes: 79 additions & 10 deletions graph_net/tools/generate_subgraph_dataset.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,18 +29,20 @@ DIM_GENERALIZED_FUSIBLE_SUBGRAPH_DIR=$DECOMPOSE_WORKSPACE/11_dimension_generaliz
RENAMED_DIM_GENERALIZED_FUSIBLE_SUBGRAPH_DIR=$DECOMPOSE_WORKSPACE/12_renamed_dimension_generalized_fusible_subgraphs
DEDUP_DIM_GENERALIZED_FUSIBLE_SUBGRAPH_DIR=$DECOMPOSE_WORKSPACE/13_deduplicated_dimension_generalized_fusible_subgraphs
DTYPE_GENERALIZED_FUSIBLE_SUBGRAPH_DIR=$DECOMPOSE_WORKSPACE/14_dtype_generalized_fusible_subgraphs
FUSIBLE_SUBGRAPH_UNITTEST_DIR=$DECOMPOSE_WORKSPACE/15_fusible_subgraphs_unittests
BACKWARD_GRAPH_OUTPUT_DIR=$DECOMPOSE_WORKSPACE/15_backward_graph_extracted
FUSIBLE_SUBGRAPH_UNITTEST_DIR=$DECOMPOSE_WORKSPACE/16_fusible_subgraphs_unittests

# typical_subgraphs
DIM_GENERALIZED_TYPICAL_SUBGRAPH_DIR=$DECOMPOSE_WORKSPACE/2-08_dimension_generalized_typical_subgraphs
RENAMED_DIM_GENERALIZED_TYPICAL_SUBGRAPH_DIR=$DECOMPOSE_WORKSPACE/2-09_renamed_dimension_generalized_typical_subgraphs
DEDUP_DIM_GENERALIZED_TYPICAL_SUBGRAPH_DIR=$DECOMPOSE_WORKSPACE/2-10_deduplicated_dimension_generalized_typical_subgraphs
DTYPE_GENERALIZED_TYPICAL_SUBGRAPH_DIR=$DECOMPOSE_WORKSPACE/2-11_dtype_generalized_typical_subgraphs
TYPICAL_SUBGRAPH_UNITTEST_DIR=$DECOMPOSE_WORKSPACE/2-12_typical_kernelbench_unittests
BACKWARD_GRAPH_TYPICAL_OUTPUT_DIR=$DECOMPOSE_WORKSPACE/2-13_backward_graph_extracted
TYPICAL_SUBGRAPH_UNITTEST_DIR=$DECOMPOSE_WORKSPACE/2-14_typical_kernelbench_unittests

mkdir -p "$DECOMPOSE_WORKSPACE"

model_list="$GRAPH_NET_ROOT/graph_net/config/torch_samples_list.txt"
model_list="$GRAPH_NET_ROOT/graph_net/config/small10_torch_samples_list.txt"
DB_PATH=$DECOMPOSE_WORKSPACE/small100_torch_samples.db

device_rewrited_sample_list=${DECOMPOSE_WORKSPACE}/device_rewrited_sample_list.txt
Expand Down Expand Up @@ -401,8 +403,28 @@ EOF
)
}

function backward_graph_extractor() {
echo ">>> [13] Extract backward graphs for samples under ${DTYPE_GENERALIZED_FUSIBLE_SUBGRAPH_DIR}."
echo ">>>"
python3 -m graph_net.model_path_handler \
--model-path-list ${dtype_generalized_subgraphs_list} \
--handler-config=$(base64 -w 0 <<EOF
{
"handler_path": "$GRAPH_NET_ROOT/graph_net/torch/sample_pass/backward_graph_extractor.py",
"handler_class_name": "BackwardGraphExtractorPass",
"handler_config": {
"model_path_prefix": "$DTYPE_GENERALIZED_FUSIBLE_SUBGRAPH_DIR",
"output_dir": "$BACKWARD_GRAPH_OUTPUT_DIR",
"device": "cuda",
"resume": ${RESUME}
}
}
EOF
)
}

function generate_unittests() {
echo ">>> [13] Generate unittests for subgraph samples under ${DTYPE_GENERALIZED_FUSIBLE_SUBGRAPH_DIR}. "
echo ">>> [14] Generate unittests for subgraph samples under ${BACKWARD_GRAPH_OUTPUT_DIR}. "
echo ">>>"
python3 -m graph_net.model_path_handler \
--model-path-list ${dtype_generalized_subgraphs_list} \
Expand All @@ -412,7 +434,7 @@ function generate_unittests() {
"handler_class_name": "AgentUnittestGeneratorPass",
"handler_config": {
"framework": "torch",
"model_path_prefix": "${DTYPE_GENERALIZED_FUSIBLE_SUBGRAPH_DIR}",
"model_path_prefix": "${BACKWARD_GRAPH_OUTPUT_DIR}",
"output_dir": "${FUSIBLE_SUBGRAPH_UNITTEST_DIR}",
"device": "cuda",
"generate_main": false,
Expand Down Expand Up @@ -516,8 +538,28 @@ EOF
)
}

function backward_graph_extractor_typical() {
echo ">>> [2-13] Extract backward graphs for typical subgraph samples under ${DTYPE_GENERALIZED_TYPICAL_SUBGRAPH_DIR}."
echo ">>>"
python3 -m graph_net.model_path_handler \
--model-path-list ${dtype_generalized_typical_subgraph_list} \
--handler-config=$(base64 -w 0 <<EOF
{
"handler_path": "$GRAPH_NET_ROOT/graph_net/torch/sample_pass/backward_graph_extractor.py",
"handler_class_name": "BackwardGraphExtractorPass",
"handler_config": {
"model_path_prefix": "$DTYPE_GENERALIZED_TYPICAL_SUBGRAPH_DIR",
"output_dir": "$BACKWARD_GRAPH_TYPICAL_OUTPUT_DIR",
"device": "cuda",
"resume": ${RESUME}
}
}
EOF
)
}

function generate_unittest_for_typical_subgraphs() {
echo ">>> [2-13] Generate unittests for subgraph samples under ${DTYPE_GENERALIZED_TYPICAL_SUBGRAPH_DIR}. "
echo ">>> [2-14] Generate unittests for subgraph samples under ${BACKWARD_GRAPH_TYPICAL_OUTPUT_DIR}. "
echo ">>>"
python3 -m graph_net.model_path_handler \
--model-path-list ${dtype_generalized_typical_subgraph_list} \
Expand All @@ -527,7 +569,7 @@ function generate_unittest_for_typical_subgraphs() {
"handler_class_name": "AgentUnittestGeneratorPass",
"handler_config": {
"framework": "torch",
"model_path_prefix": "${DTYPE_GENERALIZED_TYPICAL_SUBGRAPH_DIR}",
"model_path_prefix": "${BACKWARD_GRAPH_TYPICAL_OUTPUT_DIR}",
"output_dir": "${TYPICAL_SUBGRAPH_UNITTEST_DIR}",
"device": "cuda",
"generate_main": false,
Expand Down Expand Up @@ -582,6 +624,9 @@ function generate_fusible_subgraphs() {
dtype_generalizer 2>&1 | tee ${DECOMPOSE_WORKSPACE}/log_dtype_generalizer_${suffix}.txt
generate_generalized_subgraph_list ${DTYPE_GENERALIZED_FUSIBLE_SUBGRAPH_DIR} ${dtype_generalized_subgraphs_list}

# extract backward graphs (train_forward, train_backward, eval_forward)
backward_graph_extractor 2>&1 | tee ${DECOMPOSE_WORKSPACE}/log_backward_graph_extractor_${suffix}.txt

# generate kernelbench format unittest
generate_unittests 2>&1 | tee ${DECOMPOSE_WORKSPACE}/log_unittests_${suffix}.txt
}
Expand All @@ -602,6 +647,9 @@ function generate_typical_subgraphs() {
dtype_generalizer_for_typical_subgraphs 2>&1 | tee ${DECOMPOSE_WORKSPACE}/log_dtype_generalizer_typical_subgraphs_${suffix}.txt
generate_generalized_subgraph_list ${DTYPE_GENERALIZED_TYPICAL_SUBGRAPH_DIR} ${dtype_generalized_typical_subgraph_list}

# extract backward graphs (train_forward, train_backward, eval_forward)
backward_graph_extractor_typical 2>&1 | tee ${DECOMPOSE_WORKSPACE}/log_backward_graph_extractor_typical_${suffix}.txt

# generate kernelbench format unittest
generate_unittest_for_typical_subgraphs 2>&1 | tee ${DECOMPOSE_WORKSPACE}/log_unittests_typical_subgraphs_${suffix}.txt
}
Expand All @@ -613,11 +661,14 @@ function main() {
python ${GRAPH_NET_ROOT}/sqlite/init_db.py --db_path ${DB_PATH} 2>&1 | tee sqlite/logs/init_db_${timestamp}.log
insert_graph_sample ${GRAPH_NET_ROOT} "github_torch_samples" "full_graph" ${model_list}

# Generate common decomposition (must run first)
do_common_generalzation_and_decompose

generate_fusible_subgraphs
insert_graph_sample ${DEDUP_DIM_GENERALIZED_FUSIBLE_SUBGRAPH_DIR} "github_torch_samples" "fusible_graph" ${deduplicated_fusible_subgraphs_list}

generate_typical_subgraphs
insert_graph_sample ${DEDUP_TYPICAL_SUBGRAPH_DIR} "github_torch_samples" "typical_graph" ${deduplicated_typical_subgraph_list}
insert_graph_sample ${DEDUP_TYPICAL_SUBGRAPH_DIR} "github_torch_samples" "typical_graph" ${deduplicated_typical_subgraph_list}
}

summary() {
Expand Down Expand Up @@ -700,9 +751,18 @@ summary() {
done
echo ""

num_backward_graph_subgraphs=`find ${BACKWARD_GRAPH_OUTPUT_DIR} -name "model.py" | wc -l`
echo "- [Fusible - 6] backward graph extraction: successed=${num_backward_graph_subgraphs}"
for mode in eval_forward train_forward train_backward
do
num_backward_graph_subgraphs_mode=`find ${BACKWARD_GRAPH_OUTPUT_DIR}/${mode} -name "model.py" | wc -l`
echo " ${mode}, successed=${num_backward_graph_subgraphs_mode}"
done
echo ""

num_successed_unittests=`find ${FUSIBLE_SUBGRAPH_UNITTEST_DIR} -name "*_test.py" | wc -l`
unittest_successed_percent=$((num_successed_unittests * 100 / num_dtype_generalized_subgraphs))
echo "- [Fusible - 6] generate unittest: successed=${num_successed_unittests}, percent=${unittest_successed_percent}%"
unittest_successed_percent=$((num_successed_unittests * 100 / num_backward_graph_subgraphs))
echo "- [Fusible - 7] generate unittest: successed=${num_successed_unittests}, percent=${unittest_successed_percent}%"
for dtype in float32 float16 bfloat16
do
num_successed_unittests=`find ${FUSIBLE_SUBGRAPH_UNITTEST_DIR}/${dtype} -name "*_test.py" | wc -l`
Expand Down Expand Up @@ -734,6 +794,15 @@ summary() {
echo " ${index}, successed=${num_deduplicated_typical_subgraphs_index}"
done
echo ""

num_backward_graph_typical_subgraphs=`find ${BACKWARD_GRAPH_TYPICAL_OUTPUT_DIR} -name "model.py" | wc -l`
echo "- [Typical - 4] backward graph extraction: successed=${num_backward_graph_typical_subgraphs}"
for mode in eval_forward train_forward train_backward
do
num_backward_graph_typical_mode=`find ${BACKWARD_GRAPH_TYPICAL_OUTPUT_DIR}/${mode} -name "model.py" | wc -l`
echo " ${mode}, successed=${num_backward_graph_typical_mode}"
done
echo ""
}

main
Expand Down
28 changes: 15 additions & 13 deletions graph_net/torch/sample_pass/backward_graph_extractor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,36 +16,40 @@


class BackwardGraphExtractor:
def __init__(self, model_name, model_path, output_dir, device):
def __init__(self, model_name, model_path, output_dir, device, rel_model_path):
self.model_name = model_name
self.model_path = model_path
self.output_dir = output_dir
self.device = device
self.rel_model_path = rel_model_path

def __call__(self):
module, forward_inputs = get_torch_module_and_inputs(
self.model_path, use_dummy_inputs=False, device=self.device
)
module.train()

original_name_dir = os.path.join(self.output_dir, self.model_name)
if not os.path.exists(original_name_dir):
shutil.copytree(self.model_path, original_name_dir)
eval_forward_dir = os.path.join(
self.output_dir, "eval_forward", self.rel_model_path
)
if not os.path.exists(eval_forward_dir):
shutil.copytree(self.model_path, eval_forward_dir)

forward_inputs = self.set_requires_grad_for_forward_inputs(
self.model_path, module, forward_inputs
)
gm_holder, backward_inputs = self.capture_graph(module, forward_inputs)
self.get_extractor("forward")(gm_holder["forward_gm"], forward_inputs)
self.get_extractor("backward")(gm_holder["backward_gm"], backward_inputs)
self.get_extractor("train_forward")(gm_holder["forward_gm"], forward_inputs)
self.get_extractor("train_backward")(gm_holder["backward_gm"], backward_inputs)

def get_extractor(self, suffix):
workspace_path = os.path.join(self.output_dir, suffix, self.rel_model_path)
return BuiltinGraphExtractor(
name=f"{self.model_name}_{suffix}",
dynamic=False,
mut_graph_codes=[],
placeholder_auto_rename=False,
workspace_path=self.output_dir,
workspace_path=workspace_path,
)

def capture_graph(self, module, forward_inputs):
Expand Down Expand Up @@ -180,13 +184,11 @@ def resume(self, rel_model_path: str):
model_path_prefix = Path(self.config["model_path_prefix"])
model_name = f"{os.path.basename(rel_model_path)}"
model_path = model_path_prefix / rel_model_path
output_dir = (
Path(self.config["output_dir"])
/ os.path.dirname(rel_model_path)
/ model_name
)
output_dir = Path(self.config["output_dir"])
device = self._choose_device(self.config["device"])
extractor = BackwardGraphExtractor(model_name, model_path, output_dir, device)
extractor = BackwardGraphExtractor(
model_name, model_path, output_dir, device, rel_model_path
)
extractor()

def _choose_device(self, device) -> str:
Expand Down
Loading