github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-16 14:41:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

[autochunk] support complete benchmark (#3121)

Browse Source 
* refact memory code

* dont log free var memory

* add memory align

* update chunk target

* update setting for new memory

* finish test

* update tracer

* update typo

* update test

* add unet test

* add bench

* update bench

* update bench

* init

* support vit

* move to cpu

* add cpu benchmark
This commit is contained in:
Xuanlei Zhao
2023-03-13 17:42:37 +08:00
committed by GitHub
parent 68577fbc43
commit 30dd13c450
3 changed files with 8 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/test_autochunk/test_autochunk_alphafold/benchmark_autochunk_alphafold.py
View File

@@ -23,7 +23,7 @@ def _benchmark_evoformer_stack_gm(
get_data: Any,
) -> None:
# build model and input
model = get_model()
model = get_model().cpu().eval()
meta_args, concrete_args = get_data(*data_args)
if concrete_args is None:
concrete_args = []
@@ -35,7 +35,7 @@ def _benchmark_evoformer_stack_gm(
concrete_args={k: v for k, v in concrete_args},
)
interp = MetaInfoProp(meta_graph)
meta_tensors = [MetaTensor(i[1], fake_device="cuda:0") for i in meta_args] + [i[1] for i in concrete_args]
meta_tensors = [MetaTensor(i[1], fake_device="cpu") for i in meta_args] + [i[1] for i in concrete_args]
interp.propagate(*meta_tensors)
codegen = AutoChunkCodeGen(
meta_graph,