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

add test in import

pull/2364/head
oahzxl 2023-01-10 11:20:28 +08:00
parent fd818cf144
commit c1492e5013
2 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tests/test_autochunk/test_autochunk_codegen.py
View File

@ -6,7 +6,6 @@ import torch.fx
import torch.multiprocessing as mp
import colossalai
from colossalai.autochunk.autochunk_codegen import AutoChunkCodeGen
from colossalai.core import global_context as gpc
from colossalai.fx import ColoTracer
from colossalai.fx.codegen.activation_checkpoint_codegen import CODEGEN_AVAILABLE
@ -16,6 +15,9 @@ from colossalai.fx.profiler import MetaTensor
from colossalai.utils import free_port
from tests.test_autochunk.evoformer.evoformer import evoformer_base
if CODEGEN_AVAILABLE:
from colossalai.autochunk.autochunk_codegen import AutoChunkCodeGen
def _test_fwd(model: torch.nn.Module, gm: ColoGraphModule, node, pair):
# for memory test

4
tests/test_autochunk/test_autochunk_search.py
View File

@ -6,7 +6,6 @@ import torch.fx
import torch.multiprocessing as mp
import colossalai
from colossalai.autochunk.autochunk_codegen import AutoChunkCodeGen
from colossalai.core import global_context as gpc
from colossalai.fx.codegen.activation_checkpoint_codegen import CODEGEN_AVAILABLE
from colossalai.fx.passes.meta_info_prop import MetaInfoProp
@ -14,6 +13,9 @@ from colossalai.fx.profiler import MetaTensor
from colossalai.utils import free_port
from tests.test_autochunk.evoformer.evoformer import evoformer_base
if CODEGEN_AVAILABLE:
from colossalai.autochunk.autochunk_codegen import AutoChunkCodeGen
def assert_chunk_infos(chunk_infos, max_memory, msa_len, pair_len):
found_regions = [i["region"] for i in chunk_infos]