github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-12 12:47:21 +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_torchrec_model/test_deepfm_model.py
View File

@@ -2,6 +2,7 @@ import pytest
import torch
from colossalai._analyzer.fx import symbolic_trace
from colossalai.testing import clear_cache_before_run
from tests.kit.model_zoo import model_zoo
BATCH = 2
@@ -47,6 +48,7 @@ def trace_and_compare(model_cls, data, output_transform_fn, meta_args=None):
), f'{model.__class__.__name__} has inconsistent outputs, {fx_out} vs {non_fx_out}'
@clear_cache_before_run()
def test_torchrec_deepfm_models():
deepfm_models = model_zoo.get_sub_registry('deepfm')
torch.backends.cudnn.deterministic = True

2
tests/test_fx/test_tracer/test_torchrec_model/test_dlrm_model.py
View File

@@ -2,6 +2,7 @@ import pytest
import torch
from colossalai._analyzer.fx import symbolic_trace
from colossalai.testing import clear_cache_before_run
from tests.kit.model_zoo import model_zoo
BATCH = 2
@@ -47,6 +48,7 @@ def trace_and_compare(model_cls, data, output_transform_fn, meta_args=None):
), f'{model.__class__.__name__} has inconsistent outputs, {fx_out} vs {non_fx_out}'
@clear_cache_before_run()
def test_torchrec_dlrm_models():
torch.backends.cudnn.deterministic = True
dlrm_models = model_zoo.get_sub_registry('dlrm')