github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

code style

pull/2364/head
oahzxl 2022-12-29 14:47:16 +08:00
parent 5a916c0adb
commit 7a23deb584
1 changed files with 13 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
autochunk_benchmark.py
View File

@ -34,15 +34,23 @@ def _benchmark_evoformer(model: torch.nn.Module, node, pair, title):
def benchmark_evoformer(): def benchmark_evoformer():
# data # init data and model
msa_len = 300 msa_len = 300
pair_len = 800 pair_len = 800
node = torch.randn(1, msa_len, pair_len, 256).cuda() node = torch.randn(1, msa_len, pair_len, 256).cuda()
pair = torch.randn(1, pair_len, pair_len, 128).cuda() pair = torch.randn(1, pair_len, pair_len, 128).cuda()
# build gm model
max_memory = 3000 # MB
model = evoformer_base().cuda() model = evoformer_base().cuda()
# build autochunk model
max_memory = 3000 # MB
autochunk = _build_autochunk(model, max_memory, node, pair)
# benchmark
_benchmark_evoformer(model, node, pair, "openfold")
_benchmark_evoformer(autochunk, node, pair, "autochunk")
def _build_autochunk(model, max_memory, node, pair):
# trace the module and replace codegen # trace the module and replace codegen
graph = ColoTracer().trace( graph = ColoTracer().trace(
model, model,
@ -70,9 +78,7 @@ def benchmark_evoformer():
# print # print
code = graph.python_code("self").src code = graph.python_code("self").src
print(code) print(code)
return gm
_benchmark_evoformer(gm, node, pair, "autochunk")
_benchmark_evoformer(model, node, pair, "openfold")
if __name__ == "__main__": if __name__ == "__main__":