github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-07 20:10:17 +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)