diff --git a/tests/kit/model_zoo/transformers/gpt.py b/tests/kit/model_zoo/transformers/gpt.py index ab5d97420..f71776b6b 100644 --- a/tests/kit/model_zoo/transformers/gpt.py +++ b/tests/kit/model_zoo/transformers/gpt.py @@ -18,23 +18,8 @@ def data_gen(): # tokenized_input = tokenizer(input, return_tensors='pt') # input_ids = tokenized_input['input_ids'] # attention_mask = tokenized_input['attention_mask'] - # input_ids = torch.tensor([[15496, 11, 616, 3290, 318, 13779, 318, 13779]], dtype=torch.int64) - # attention_mask = torch.tensor([[1, 1, 1, 1, 1, 1, 1, 1]], dtype=torch.int64) - input_ids = torch.tensor( - [ - [15496, 11, 616, 3290, 318, 13779, 318, 13779, 15496, 11, 616, 3290, 318, 13779, 318, 13779], - [15496, 11, 616, 3290, 318, 13779, 318, 13779, 15496, 11, 616, 3290, 318, 13779, 318, 13779], - ], - dtype=torch.int64, - ) - attention_mask = torch.tensor( - [ - [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], - [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], - ], - dtype=torch.int64, - ) - + input_ids = torch.tensor([[22, 11, 616, 4, 5, 13, 318, 345]], dtype=torch.int64) + attention_mask = torch.tensor([[1, 1, 1, 1, 1, 1, 1, 1]], dtype=torch.int64) return dict(input_ids=input_ids, attention_mask=attention_mask) @@ -50,9 +35,9 @@ def data_gen_for_question_answering(): # question answering data gen # `labels` is the type not the token id for token classification, 0 or 1 data = data_gen() - start_positions = torch.tensor([[0], [0]], dtype=torch.int64) + start_positions = torch.tensor([0], dtype=torch.int64) data["start_positions"] = start_positions - end_positions = torch.tensor([[1], [1]], dtype=torch.int64) + end_positions = torch.tensor([1], dtype=torch.int64) data["end_positions"] = end_positions return data @@ -61,20 +46,14 @@ def data_gen_for_token_classification(): # token classification data gen # `labels` is the type not the token id for token classification, 0 or 1 data = data_gen() - data["labels"] = torch.tensor( - [ - [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1], - [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1], - ], - dtype=torch.int64, - ) + data["labels"] = torch.tensor([[0, 0, 0, 0, 0, 0, 0, 1]], dtype=torch.int64) return data def data_gen_for_sequence_classification(): # sequence classification data gen data = data_gen() - data["labels"] = torch.tensor([[1], [1]], dtype=torch.int64) + data["labels"] = torch.tensor([1], dtype=torch.int64) return data @@ -82,18 +61,12 @@ def date_gen_for_double_heads(): num_choices = 2 batch_size = 2 input_ids = torch.tensor( - [ - [15496, 11, 616, 3290, 318, 13779, 318, 13779, 15496, 11, 616, 3290, 318, 13779, 318, 13779], - [15496, 11, 616, 3290, 318, 13779, 318, 13779, 15496, 11, 616, 3290, 318, 13779, 318, 13779], - ], + [[46, 11, 616, 432, 318, 19, 318, 555], [777, 11, 235, 333, 318, 231, 468, 136]], dtype=torch.int64, ) - attention_mask = torch.tensor( - [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], - dtype=torch.int64, - ) - + attention_mask = torch.tensor([[1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1]], dtype=torch.int64) mc_labels = torch.zeros(input_ids.shape[0], dtype=torch.int64) + mc_token_ids = torch.arange(0, num_choices, dtype=torch.int64) mc_token_ids = mc_token_ids.expand((batch_size, num_choices)) multiple_choice_inputs_ids = input_ids.unsqueeze(1).expand(-1, num_choices, -1).contiguous() @@ -122,14 +95,14 @@ config = transformers.GPT2Config( n_layer=2, n_head=4, n_embd=128, - vocab_size=50258, + vocab_size=1024, attn_pdrop=0, embd_pdrop=0, resid_pdrop=0, summary_first_dropout=0, hidden_dropout=0, problem_type="single_label_classification", - pad_token_id=50256, + pad_token_id=1022, tie_word_embeddings=True, ) diff --git a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py index 0e941f4b9..fd13ce0bf 100644 --- a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py +++ b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py @@ -21,14 +21,10 @@ from colossalai.testing import ( from tests.kit.model_zoo import model_zoo MODEL_PLACEMENT_CONFIGS = [ - {"placement_policy": "static", "shard_param_frac": 0.0}, # zero2 - {"placement_policy": "static", "shard_param_frac": 1.0}, # zero3 - {"placement_policy": "static", "shard_param_frac": 0.5}, # zero3-half + {"placement_policy": "static", "shard_param_frac": 0.5}, ] OPTIM_PLACEMENT_CONFIGS = [ - {"placement_policy": "static", "shard_param_frac": 0.0, "offload_optim_frac": 0.0}, # zero2 - {"placement_policy": "static", "shard_param_frac": 0.0, "offload_optim_frac": 1.0}, # zero2-offload {"placement_policy": "static", "shard_param_frac": 0.0, "offload_optim_frac": 0.5}, # zero2-offload-half ] diff --git a/tests/test_zero/test_gemini/test_grad_accum.py b/tests/test_zero/test_gemini/test_grad_accum.py index 002741389..3299cf631 100644 --- a/tests/test_zero/test_gemini/test_grad_accum.py +++ b/tests/test_zero/test_gemini/test_grad_accum.py @@ -15,9 +15,7 @@ from colossalai.zero.gemini.chunk import search_chunk_configuration from tests.kit.model_zoo import model_zoo, run_fwd PLACEMENT_CONFIGS = [ - {"placement_policy": "static", "shard_param_frac": 0.0}, # zero2 - {"placement_policy": "static", "shard_param_frac": 1.0}, # zero3 - {"placement_policy": "static", "shard_param_frac": 0.5}, # zero3-half + {"placement_policy": "static", "shard_param_frac": 0.75}, {"placement_policy": "auto"}, ] @@ -109,7 +107,7 @@ def exam_gemini_grad_acc( torch_model = DDP(torch_model, device_ids=[rank]) set_seed(rank) - accum_iter = 4 + accum_iter = 2 train_dataloader = DummyDataloader(data_gen_fn) for i, data in enumerate(train_dataloader): delay_unscale = False if (i + 1) % accum_iter == 0 else True diff --git a/tests/test_zero/test_gemini/test_optim.py b/tests/test_zero/test_gemini/test_optim.py index c610259b2..39cf348d9 100644 --- a/tests/test_zero/test_gemini/test_optim.py +++ b/tests/test_zero/test_gemini/test_optim.py @@ -15,17 +15,7 @@ from colossalai.zero.gemini.chunk import search_chunk_configuration from tests.kit.model_zoo import model_zoo, run_fwd_bwd PLACEMENT_CONFIGS = [ - {"placement_policy": "static", "shard_param_frac": 0.0, "offload_optim_frac": 0.0}, # zero2 - {"placement_policy": "static", "shard_param_frac": 0.0, "offload_optim_frac": 1.0}, # zero2-offload - {"placement_policy": "static", "shard_param_frac": 0.0, "offload_optim_frac": 0.5}, # zero2-offload-half - {"placement_policy": "static", "shard_param_frac": 1.0}, # zero3 - {"placement_policy": "static", "shard_param_frac": 0.5}, # zero3-half - { - "placement_policy": "static", - "shard_param_frac": 1.0, - "offload_optim_frac": 1.0, - "offload_param_frac": 1.0, - }, # zero3-offload-all + {"placement_policy": "static", "shard_param_frac": 0.3, "offload_param_frac": 0.3, "offload_optim_frac": 0.3}, {"placement_policy": "auto"}, ] @@ -73,7 +63,7 @@ def check_param(model: GeminiDDP, torch_model: torch.nn.Module, dtype: torch.dty @parameterize("model_name", TEST_MODELS) @parameterize("mixed_precision", [torch.half, torch.bfloat16]) @parameterize("master_weights", [True, False]) -@parameterize("enable_async_reduce", [False, True]) +@parameterize("enable_async_reduce", [True]) def exam_model_step( placement_config, model_name: str, mixed_precision: torch.dtype, master_weights: bool, enable_async_reduce=True ): @@ -136,7 +126,7 @@ def exam_model_step( check_param(model, torch_model, mixed_precision) -@parameterize("placement_config", [PLACEMENT_CONFIGS[3]]) +@parameterize("placement_config", [{"placement_policy": "static", "shard_param_frac": 1.0}]) @parameterize("model_name", EXAMPLE_MODELS) @parameterize("mixed_precision", [torch.half]) def exam_tiny_example(placement_config, model_name: str, mixed_precision: torch.dtype): @@ -197,7 +187,7 @@ def run_dist(rank, world_size, port): @pytest.mark.dist -@pytest.mark.parametrize("world_size", [1, 4]) +@pytest.mark.parametrize("world_size", [4]) @rerun_if_address_is_in_use() def test_optim(world_size): spawn(run_dist, world_size) diff --git a/tests/test_zero/test_gemini/test_search.py b/tests/test_zero/test_gemini/test_search.py index 9c8c497f3..58e585474 100644 --- a/tests/test_zero/test_gemini/test_search.py +++ b/tests/test_zero/test_gemini/test_search.py @@ -1,18 +1,31 @@ import pytest import torch +import transformers import colossalai from colossalai.accelerator import get_accelerator from colossalai.testing import rerun_if_address_is_in_use, spawn from colossalai.zero.gemini.chunk import init_chunk_manager, search_chunk_configuration -from tests.kit.model_zoo import model_zoo + +CONFIG = transformers.GPT2Config( + n_layer=2, + n_head=4, + n_embd=128, + vocab_size=50258, + attn_pdrop=0, + embd_pdrop=0, + resid_pdrop=0, + summary_first_dropout=0, + hidden_dropout=0, + problem_type="single_label_classification", + pad_token_id=50256, + tie_word_embeddings=True, +) + +model_builder = lambda: transformers.GPT2LMHeadModel(CONFIG) def exam_search_chunk_size(): - model_builder, data_gen_fn, output_transform_fn, *_ = next( - iter(model_zoo.get_sub_registry("transformers_gpt_lm").values()) - ) - # make sure torch_model and model has the same parameter values model = model_builder() config_dict, *_ = search_chunk_configuration( @@ -27,10 +40,6 @@ def exam_search_chunk_size(): def exam_chunk_manager(): world_size = torch.distributed.get_world_size() - model_builder, data_gen_fn, output_transform_fn, *_ = next( - iter(model_zoo.get_sub_registry("transformers_gpt_lm").values()) - ) - sharded_ddp_model = model_builder() chunk_manager = init_chunk_manager( sharded_ddp_model, diff --git a/tests/test_zero/test_gemini/test_zeroddp_state_dict.py b/tests/test_zero/test_gemini/test_zeroddp_state_dict.py index 23e2d8083..00d28f1c0 100644 --- a/tests/test_zero/test_gemini/test_zeroddp_state_dict.py +++ b/tests/test_zero/test_gemini/test_zeroddp_state_dict.py @@ -10,9 +10,7 @@ from colossalai.zero.gemini.chunk import search_chunk_configuration from tests.kit.model_zoo import model_zoo PLACEMENT_CONFIGS = [ - {"placement_policy": "static", "shard_param_frac": 0.0}, # zero2 - {"placement_policy": "static", "shard_param_frac": 1.0}, # zero3 - {"placement_policy": "static", "shard_param_frac": 0.5}, # zero3-half + {"placement_policy": "static", "shard_param_frac": 0.75}, {"placement_policy": "auto"}, ] @@ -26,8 +24,8 @@ def ignore_the_first_parameter(model: torch.nn.Module): @parameterize("placement_config", PLACEMENT_CONFIGS) @parameterize("keep_gathered", [True, False]) -@parameterize("model_name", ["transformers_gpt_lm", "transformers_bert_for_sequence_classification"]) -@parameterize("master_weights", [False, True]) +@parameterize("model_name", ["transformers_gpt_lm"]) +@parameterize("master_weights", [True, False]) def exam_state_dict(placement_config, keep_gathered, model_name: str, master_weights: bool): set_seed(431) model_builder, data_gen_fn, output_transform_fn, *_ = next(iter(model_zoo.get_sub_registry(model_name).values())) @@ -81,7 +79,7 @@ def run_dist(rank, world_size, port): @pytest.mark.dist -@pytest.mark.parametrize("world_size", [1, 4]) +@pytest.mark.parametrize("world_size", [4]) @rerun_if_address_is_in_use() def test_zero_ddp(world_size): spawn(run_dist, world_size)