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

[workflow] fixed build CI (#5240)

Browse Source 
* [workflow] fixed build CI

* polish

* polish

* polish

* polish

* polish
This commit is contained in:
Frank Lee
2024-01-10 22:34:16 +08:00
committed by GitHub
parent 41e52c1c6e
commit edf94a35c3
14 changed files with 101 additions and 156 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
tests/test_checkpoint_io/test_gemini_checkpoint_io.py
View File

@@ -7,6 +7,7 @@ from transformers import LlamaForCausalLM
from utils import shared_tempdir
import colossalai
from colossalai.testing import skip_if_not_enough_gpus
from colossalai.booster import Booster
from colossalai.booster.plugin import GeminiPlugin
from colossalai.lazy import LazyInitContext
@@ -68,7 +69,7 @@ def exam_state_dict_with_origin(placement_config, model_name, use_safetensors: b
@clear_cache_before_run()
@parameterize("placement_config", OPTIM_PLACEMENT_CONFIGS)
@parameterize("shard", [True, False])
@parameterize("model_name", ["transformers_gpt"])
@parameterize("model_name", ["transformers_llama_for_casual_lm"])
@parameterize("size_per_shard", [32])
@parameterize("tp_size", [1, 2])
@parameterize("zero_size", [2])
@@ -156,13 +157,12 @@ def run_dist(rank, world_size, port):
@pytest.mark.dist
@pytest.mark.parametrize("world_size", [4])
@rerun_if_address_is_in_use()
def test_gemini_ckpIO(world_size):
spawn(run_dist, world_size)
def test_gemini_ckpIO():
spawn(run_dist, 4)
@pytest.mark.largedist
@pytest.mark.parametrize("world_size", [8])
@skip_if_not_enough_gpus(min_gpus=8)
@rerun_if_address_is_in_use()
def test_gemini_ckpIO_3d(world_size):
spawn(run_dist, world_size)
def test_gemini_ckpIO_3d():
spawn(run_dist, 8)

2
tests/test_checkpoint_io/test_gemini_torch_compability.py
View File

@@ -20,7 +20,7 @@ from tests.kit.model_zoo import model_zoo
@clear_cache_before_run()
@parameterize("shard", [False, True])
@parameterize("model_name", ["transformers_gpt"])
@parameterize("model_name", ["transformers_llama_for_casual_lm"])
def exam_torch_load_from_gemini(shard: bool, model_name: str):
(model_fn, data_gen_fn, output_transform_fn, _, _) = next(iter(model_zoo.get_sub_registry(model_name).values()))
criterion = lambda x: x.mean()

2
tests/test_checkpoint_io/test_hybrid_parallel_plugin_checkpoint_io.py
View File

@@ -40,7 +40,7 @@ else:
@clear_cache_before_run()
@parameterize("shard", [True, False])
@parameterize("model_name", ["transformers_gpt"])
@parameterize("model_name", ["transformers_llama_for_casual_lm"])
@parameterize("size_per_shard", [32])
@parameterize("test_config", TEST_CONFIGS)
def exam_state_dict(shard: bool, model_name: str, size_per_shard: int, test_config: dict):

2
tests/test_checkpoint_io/test_plugins_huggingface_compatibility.py
View File

@@ -18,7 +18,7 @@ from tests.kit.model_zoo import model_zoo
@clear_cache_before_run()
@parameterize("model_name", ["transformers_gpt"])
@parameterize("model_name", ["transformers_llama_for_casual_lm"])
@parameterize("plugin_type", ["ddp", "zero", "gemini"])
def exam_from_pretrained(plugin_type: str, model_name: str, shard=True, size_per_shard=32):
(model_fn, data_gen_fn, output_transform_fn, loss_fn, _) = next(