github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-04 18:40:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

[ci] fixed booster test (#5251)

Browse Source 
* [ci] fixed booster test

* [ci] fixed booster test

* [ci] fixed booster test
This commit is contained in:
Frank Lee
2024-01-11 16:04:45 +08:00
committed by GitHub
parent edf94a35c3
commit d5eeeb1416
5 changed files with 12 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/test_booster/test_plugin/test_3d_plugin.py
View File

@@ -10,10 +10,11 @@ from colossalai.booster.plugin import HybridParallelPlugin
from colossalai.fx import is_compatible_with_meta
from colossalai.lazy.lazy_init import LazyInitContext
from colossalai.nn.optimizer import HybridAdam
from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
from colossalai.testing import clear_cache_before_run, parameterize, rerun_if_address_is_in_use, spawn
from tests.kit.model_zoo import model_zoo
@clear_cache_before_run()
def run_fn(init_method, model_fn, data_gen_fn, output_transform_fn) -> Optional[str]:
try:
if init_method == "lazy":
@@ -69,7 +70,6 @@ def check_3d_plugin(init_method: str = "none", early_stop: bool = True):
"transformers_llama_for_casual_lm"
).items():
err = run_fn(init_method, model_fn, data_gen_fn, output_transform_fn)
torch.cuda.empty_cache()
if err is None:
passed_models.append(name)