From 27ab5240965fc9cc0ec74ff48356abcbf098bd74 Mon Sep 17 00:00:00 2001 From: oahzxl Date: Fri, 6 Jan 2023 11:07:57 +0800 Subject: [PATCH] refactor structure --- .../chunk_codegen.py | 41 ++++++++----------- .../evoformer}/evoformer.py | 0 .../evoformer}/initializer.py | 0 {evoformer => autochunk/evoformer}/kernel.py | 0 {evoformer => autochunk/evoformer}/msa.py | 0 {evoformer => autochunk/evoformer}/ops.py | 0 .../evoformer}/triangle.py | 0 .../openfold}/checkpointing.py | 0 {openfold => autochunk/openfold}/dropout.py | 0 {openfold => autochunk/openfold}/evoformer.py | 0 {openfold => autochunk/openfold}/msa.py | 0 .../openfold}/outer_product_mean.py | 0 .../openfold}/pair_transition.py | 0 .../openfold}/primitives.py | 0 .../openfold}/tensor_utils.py | 0 .../openfold}/triangular_attention.py | 0 .../triangular_multiplicative_update.py | 0 autochunk_benchmark.py | 18 ++++---- chunk_codegen_run.py => autochunk_test.py | 4 +- 19 files changed, 29 insertions(+), 34 deletions(-) rename chunk_codegen.py => autochunk/chunk_codegen.py (98%) rename {evoformer => autochunk/evoformer}/evoformer.py (100%) rename {evoformer => autochunk/evoformer}/initializer.py (100%) rename {evoformer => autochunk/evoformer}/kernel.py (100%) rename {evoformer => autochunk/evoformer}/msa.py (100%) rename {evoformer => autochunk/evoformer}/ops.py (100%) rename {evoformer => autochunk/evoformer}/triangle.py (100%) rename {openfold => autochunk/openfold}/checkpointing.py (100%) rename {openfold => autochunk/openfold}/dropout.py (100%) rename {openfold => autochunk/openfold}/evoformer.py (100%) rename {openfold => autochunk/openfold}/msa.py (100%) rename {openfold => autochunk/openfold}/outer_product_mean.py (100%) rename {openfold => autochunk/openfold}/pair_transition.py (100%) rename {openfold => autochunk/openfold}/primitives.py (100%) rename {openfold => autochunk/openfold}/tensor_utils.py (100%) rename {openfold => autochunk/openfold}/triangular_attention.py (100%) rename {openfold => autochunk/openfold}/triangular_multiplicative_update.py (100%) rename chunk_codegen_run.py => autochunk_test.py (97%) diff --git a/chunk_codegen.py b/autochunk/chunk_codegen.py similarity index 98% rename from chunk_codegen.py rename to autochunk/chunk_codegen.py index 41fcb5a3c..7a5d06689 100644 --- a/chunk_codegen.py +++ b/autochunk/chunk_codegen.py @@ -1967,13 +1967,11 @@ def _replace_reshape_size(context, node_name, reshape_size_dict): def emit_code_with_chunk( body, - ckpt_func, nodes, emit_node_func, delete_unused_value_func, - meta_nodes, - meta_graph, - max_memory=None, + chunk_region_search, + chunk_infos ): """Emit code with nested activation checkpoint When we detect some of the node.activation_checkpoint is a List, we will use @@ -1988,23 +1986,19 @@ def emit_code_with_chunk( """ node_list = list(nodes) - # find the chunk regions - chunk_region_search = ChunkRegionSearch(meta_graph, max_memory) - chunk_search = chunk_region_search.search_region() - - chunk_regions = [i["region"] for i in chunk_search] + chunk_regions = [i["region"] for i in chunk_infos] chunk_starts = [i[0] for i in chunk_regions] chunk_ends = [i[1] for i in chunk_regions] - chunk_inputs = [i["inputs"] for i in chunk_search] - chunk_inputs_non_chunk = [i["inputs_non_chunk"] for i in chunk_search] - chunk_inputs_dim = [i["inputs_dim"] for i in chunk_search] + chunk_inputs = [i["inputs"] for i in chunk_infos] + chunk_inputs_non_chunk = [i["inputs_non_chunk"] for i in chunk_infos] + chunk_inputs_dim = [i["inputs_dim"] for i in chunk_infos] chunk_inputs_names = [j.name for i in chunk_inputs for j in i] + [ j.name for i in chunk_inputs_non_chunk for j in i ] - chunk_outputs = [i["outputs"][0] for i in chunk_search] - chunk_outputs_dim = [i["outputs_dim"] for i in chunk_search] + chunk_outputs = [i["outputs"][0] for i in chunk_infos] + chunk_outputs_dim = [i["outputs_dim"] for i in chunk_infos] node_list = chunk_region_search.index_tracer.reorder_node_list(node_list) node_idx = 0 @@ -2022,7 +2016,7 @@ def emit_code_with_chunk( chunk_inputs[region_idx], chunk_outputs[region_idx], chunk_outputs_dim[region_idx], - chunk_search[region_idx]["chunk_size"], + chunk_infos[region_idx]["chunk_size"], ) ) @@ -2041,14 +2035,14 @@ def emit_code_with_chunk( # ones like if "ones_like" in node.name: meta_node = chunk_region_search.index_tracer.node_list[node_idx] - chunk_dim = chunk_search[region_idx]["node_chunk_dim"][meta_node][ + chunk_dim = chunk_infos[region_idx]["node_chunk_dim"][meta_node][ "chunk_dim" ] if _get_node_shape(meta_node)[chunk_dim] != 1: source_node = meta_node.args[0].args[0] if ( - source_node not in chunk_search[region_idx]["node_chunk_dim"] - or chunk_search[region_idx]["node_chunk_dim"][source_node][ + source_node not in chunk_infos[region_idx]["node_chunk_dim"] + or chunk_infos[region_idx]["node_chunk_dim"][source_node][ "chunk_dim" ] is None @@ -2060,7 +2054,7 @@ def emit_code_with_chunk( body[-1], node.args[0].name, node.args[0].name + chunk_slice ) body[-1] = _replace_reshape_size( - body[-1], node.name, chunk_search[region_idx]["reshape_size"] + body[-1], node.name, chunk_infos[region_idx]["reshape_size"] ) body[-1] = " " + body[-1] delete_unused_value_func(node, body, chunk_inputs_names) @@ -2092,6 +2086,9 @@ if CODEGEN_AVAILABLE: self.meta_graph = meta_graph self.max_memory = max_memory self.meta_node = list(meta_graph.graph.nodes) + # find the chunk regions + self.chunk_region_search = ChunkRegionSearch(meta_graph, max_memory) + self.chunk_infos = self.chunk_region_search.search_region() def _gen_python_code( self, nodes, root_module: str, namespace: _Namespace @@ -2323,13 +2320,11 @@ if CODEGEN_AVAILABLE: # will use nested type of activation checkpoint codegen emit_code_with_chunk( body, - ckpt_func, nodes, emit_node, delete_unused_values, - self.meta_node, - self.meta_graph, - self.max_memory, + self.chunk_region_search, + self.chunk_infos ) if len(body) == 0: diff --git a/evoformer/evoformer.py b/autochunk/evoformer/evoformer.py similarity index 100% rename from evoformer/evoformer.py rename to autochunk/evoformer/evoformer.py diff --git a/evoformer/initializer.py b/autochunk/evoformer/initializer.py similarity index 100% rename from evoformer/initializer.py rename to autochunk/evoformer/initializer.py diff --git a/evoformer/kernel.py b/autochunk/evoformer/kernel.py similarity index 100% rename from evoformer/kernel.py rename to autochunk/evoformer/kernel.py diff --git a/evoformer/msa.py b/autochunk/evoformer/msa.py similarity index 100% rename from evoformer/msa.py rename to autochunk/evoformer/msa.py diff --git a/evoformer/ops.py b/autochunk/evoformer/ops.py similarity index 100% rename from evoformer/ops.py rename to autochunk/evoformer/ops.py diff --git a/evoformer/triangle.py b/autochunk/evoformer/triangle.py similarity index 100% rename from evoformer/triangle.py rename to autochunk/evoformer/triangle.py diff --git a/openfold/checkpointing.py b/autochunk/openfold/checkpointing.py similarity index 100% rename from openfold/checkpointing.py rename to autochunk/openfold/checkpointing.py diff --git a/openfold/dropout.py b/autochunk/openfold/dropout.py similarity index 100% rename from openfold/dropout.py rename to autochunk/openfold/dropout.py diff --git a/openfold/evoformer.py b/autochunk/openfold/evoformer.py similarity index 100% rename from openfold/evoformer.py rename to autochunk/openfold/evoformer.py diff --git a/openfold/msa.py b/autochunk/openfold/msa.py similarity index 100% rename from openfold/msa.py rename to autochunk/openfold/msa.py diff --git a/openfold/outer_product_mean.py b/autochunk/openfold/outer_product_mean.py similarity index 100% rename from openfold/outer_product_mean.py rename to autochunk/openfold/outer_product_mean.py diff --git a/openfold/pair_transition.py b/autochunk/openfold/pair_transition.py similarity index 100% rename from openfold/pair_transition.py rename to autochunk/openfold/pair_transition.py diff --git a/openfold/primitives.py b/autochunk/openfold/primitives.py similarity index 100% rename from openfold/primitives.py rename to autochunk/openfold/primitives.py diff --git a/openfold/tensor_utils.py b/autochunk/openfold/tensor_utils.py similarity index 100% rename from openfold/tensor_utils.py rename to autochunk/openfold/tensor_utils.py diff --git a/openfold/triangular_attention.py b/autochunk/openfold/triangular_attention.py similarity index 100% rename from openfold/triangular_attention.py rename to autochunk/openfold/triangular_attention.py diff --git a/openfold/triangular_multiplicative_update.py b/autochunk/openfold/triangular_multiplicative_update.py similarity index 100% rename from openfold/triangular_multiplicative_update.py rename to autochunk/openfold/triangular_multiplicative_update.py diff --git a/autochunk_benchmark.py b/autochunk_benchmark.py index c938485ef..c34b5217e 100644 --- a/autochunk_benchmark.py +++ b/autochunk_benchmark.py @@ -3,13 +3,13 @@ import time import torch import torch.fx -from chunk_codegen import ChunkCodeGen +from autochunk.chunk_codegen import ChunkCodeGen from colossalai.fx import ColoTracer from colossalai.fx.graph_module import ColoGraphModule from colossalai.fx.passes.meta_info_prop import MetaInfoProp from colossalai.fx.profiler import MetaTensor -from evoformer.evoformer import evoformer_base -from openfold.evoformer import EvoformerBlock +from autochunk.evoformer.evoformer import evoformer_base +from autochunk.openfold.evoformer import EvoformerBlock def _benchmark_evoformer(model: torch.nn.Module, node, pair, title, chunk_size=None): @@ -94,23 +94,23 @@ def _build_openfold(): def benchmark_evoformer(): # init data and model msa_len = 256 - pair_len = 2048 + pair_len = 1024 node = torch.randn(1, msa_len, pair_len, 256).cuda() pair = torch.randn(1, pair_len, pair_len, 128).cuda() model = evoformer_base().cuda() # build autochunk model - max_memory = 10000 # MB fit memory mode - # max_memory = None # min memory mode + # max_memory = 10000 # MB fit memory mode + max_memory = None # min memory mode autochunk = _build_autochunk(evoformer_base().cuda(), max_memory, node, pair) # build openfold chunk_size = 64 - openfold = _build_openfold() + # openfold = _build_openfold() # benchmark - _benchmark_evoformer(model, node, pair, "base") - _benchmark_evoformer(openfold, node, pair, "openfold", chunk_size=chunk_size) + # _benchmark_evoformer(model, node, pair, "base") + # _benchmark_evoformer(openfold, node, pair, "openfold", chunk_size=chunk_size) _benchmark_evoformer(autochunk, node, pair, "autochunk") diff --git a/chunk_codegen_run.py b/autochunk_test.py similarity index 97% rename from chunk_codegen_run.py rename to autochunk_test.py index 3a3b3c599..63f393531 100644 --- a/chunk_codegen_run.py +++ b/autochunk_test.py @@ -12,8 +12,8 @@ from colossalai.core import global_context as gpc from colossalai.fx.graph_module import ColoGraphModule from colossalai.fx.passes.meta_info_prop import MetaInfoProp, TensorMetadata from colossalai.fx.profiler import MetaTensor -from evoformer.evoformer import evoformer_base -from chunk_codegen import ChunkCodeGen +from autochunk.evoformer.evoformer import evoformer_base +from autochunk.chunk_codegen import ChunkCodeGen with_codegen = True