github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-15 22:19:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

[test] refactor tests with spawn (#3452)

Browse Source 
* [test] added spawn decorator

* polish code

* polish code

* polish code

* polish code

* polish code

* polish code
This commit is contained in:
Frank Lee
2023-04-06 14:51:35 +08:00
committed by GitHub
parent 62f4e2eb07
commit 80eba05b0a
240 changed files with 1723 additions and 2342 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_fx/test_tracer/test_hf_model/test_hf_albert.py
View File

@@ -3,6 +3,7 @@ import torch
from hf_tracer_utils import trace_model_and_compare_output
from packaging import version
from colossalai.testing import clear_cache_before_run
from tests.kit.model_zoo import model_zoo
BATCH_SIZE = 2
@@ -10,6 +11,7 @@ SEQ_LENGTH = 16
@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12')
@clear_cache_before_run()
def test_albert():
sub_registry = model_zoo.get_sub_registry('transformers_albert')

2
tests/test_fx/test_tracer/test_hf_model/test_hf_bert.py
View File

@@ -3,10 +3,12 @@ import torch
from hf_tracer_utils import trace_model_and_compare_output
from packaging import version
from colossalai.testing import clear_cache_before_run
from tests.kit.model_zoo import model_zoo
@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12')
@clear_cache_before_run()
def test_bert():
sub_registry = model_zoo.get_sub_registry('transformers_bert')

3
tests/test_fx/test_tracer/test_hf_model/test_hf_diffuser.py
View File

@@ -2,6 +2,7 @@ import pytest
import torch
from colossalai.fx import symbolic_trace
from colossalai.testing import clear_cache_before_run
from colossalai.testing.random import seed_all
from tests.kit.model_zoo import model_zoo
@@ -40,6 +41,7 @@ def trace_and_compare(model_cls, data, output_fn):
@pytest.mark.skip(reason='cannot pass this test yet')
@clear_cache_before_run()
def test_diffusers():
seed_all(9091, cuda_deterministic=True)
@@ -52,6 +54,7 @@ def test_diffusers():
print(f"{name:40s} √")
@clear_cache_before_run()
def test_torch_diffusers():
seed_all(65535, cuda_deterministic=True)

2
tests/test_fx/test_tracer/test_hf_model/test_hf_gpt.py
View File

@@ -3,10 +3,12 @@ import torch
from hf_tracer_utils import trace_model_and_compare_output
from packaging import version
from colossalai.testing import clear_cache_before_run
from tests.kit.model_zoo import model_zoo
@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12')
@clear_cache_before_run()
def test_gpt():
sub_registry = model_zoo.get_sub_registry('transformers_gpt')

2
tests/test_fx/test_tracer/test_hf_model/test_hf_opt.py
View File

@@ -3,10 +3,12 @@ import torch
from hf_tracer_utils import trace_model_and_compare_output
from packaging import version
from colossalai.testing import clear_cache_before_run
from tests.kit.model_zoo import model_zoo
@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12')
@clear_cache_before_run()
def test_opt():
sub_registry = model_zoo.get_sub_registry('transformers_opt')

2
tests/test_fx/test_tracer/test_hf_model/test_hf_t5.py
View File

@@ -3,10 +3,12 @@ import torch
from hf_tracer_utils import trace_model_and_compare_output
from packaging import version
from colossalai.testing import clear_cache_before_run
from tests.kit.model_zoo import model_zoo
@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12')
@clear_cache_before_run()
def test_t5():
sub_registry = model_zoo.get_sub_registry('transformers_t5')