use smaller pretrained models for ci
This commit is contained in:
@@ -9,7 +9,11 @@ from axolotl.train import train
|
||||
from axolotl.utils.config import normalize_config, validate_config
|
||||
from axolotl.utils.dict import DictDefault
|
||||
|
||||
from ..utils import check_model_output_exists, with_temp_dir
|
||||
from ..utils import (
|
||||
check_model_output_exists,
|
||||
check_tensorboard_loss_decreased,
|
||||
with_temp_dir,
|
||||
)
|
||||
|
||||
|
||||
class TestMixtral(unittest.TestCase):
|
||||
@@ -21,8 +25,7 @@ class TestMixtral(unittest.TestCase):
|
||||
def test_qlora(self, temp_dir):
|
||||
cfg = DictDefault(
|
||||
{
|
||||
"base_model": "hf-internal-testing/Mixtral-tiny",
|
||||
"tokenizer_config": "LoneStriker/Mixtral-8x7B-v0.1-HF",
|
||||
"base_model": "axolotl-ai-co/tiny-mixtral-30m",
|
||||
"flash_attention": True,
|
||||
"sample_packing": True,
|
||||
"sequence_len": 2048,
|
||||
@@ -41,17 +44,19 @@ class TestMixtral(unittest.TestCase):
|
||||
},
|
||||
],
|
||||
"num_epochs": 2,
|
||||
"micro_batch_size": 2,
|
||||
"micro_batch_size": 4,
|
||||
"gradient_accumulation_steps": 1,
|
||||
"output_dir": temp_dir,
|
||||
"learning_rate": 0.00001,
|
||||
"learning_rate": 2e-4,
|
||||
"optimizer": "adamw_bnb_8bit",
|
||||
"lr_scheduler": "cosine",
|
||||
"max_steps": 5,
|
||||
"save_steps": 3,
|
||||
"eval_steps": 4,
|
||||
"max_steps": 50,
|
||||
"logging_steps": 1,
|
||||
"save_steps": 50,
|
||||
"eval_steps": 50,
|
||||
"bf16": "auto",
|
||||
"save_first_step": False,
|
||||
"use_tensorboard": True,
|
||||
}
|
||||
)
|
||||
cfg = validate_config(cfg)
|
||||
@@ -60,13 +65,19 @@ class TestMixtral(unittest.TestCase):
|
||||
|
||||
train(cfg=cfg, dataset_meta=dataset_meta)
|
||||
check_model_output_exists(temp_dir, cfg)
|
||||
check_tensorboard_loss_decreased(
|
||||
temp_dir + "/runs",
|
||||
initial_window=5,
|
||||
final_window=5,
|
||||
max_initial=6.0,
|
||||
max_final=4.7,
|
||||
)
|
||||
|
||||
@with_temp_dir
|
||||
def test_ft(self, temp_dir):
|
||||
cfg = DictDefault(
|
||||
{
|
||||
"base_model": "hf-internal-testing/Mixtral-tiny",
|
||||
"tokenizer_config": "LoneStriker/Mixtral-8x7B-v0.1-HF",
|
||||
"base_model": "axolotl-ai-co/tiny-mixtral-30m",
|
||||
"flash_attention": True,
|
||||
"sample_packing": True,
|
||||
"sequence_len": 2048,
|
||||
@@ -79,17 +90,19 @@ class TestMixtral(unittest.TestCase):
|
||||
},
|
||||
],
|
||||
"num_epochs": 2,
|
||||
"micro_batch_size": 2,
|
||||
"micro_batch_size": 4,
|
||||
"gradient_accumulation_steps": 1,
|
||||
"output_dir": temp_dir,
|
||||
"learning_rate": 0.00001,
|
||||
"learning_rate": 2e-4,
|
||||
"optimizer": "adamw_bnb_8bit",
|
||||
"lr_scheduler": "cosine",
|
||||
"max_steps": 5,
|
||||
"save_steps": 3,
|
||||
"eval_steps": 4,
|
||||
"max_steps": 50,
|
||||
"logging_steps": 1,
|
||||
"save_steps": 50,
|
||||
"eval_steps": 50,
|
||||
"bf16": "auto",
|
||||
"save_first_step": False,
|
||||
"use_tensorboard": True,
|
||||
}
|
||||
)
|
||||
cfg = validate_config(cfg)
|
||||
@@ -98,3 +111,10 @@ class TestMixtral(unittest.TestCase):
|
||||
|
||||
train(cfg=cfg, dataset_meta=dataset_meta)
|
||||
check_model_output_exists(temp_dir, cfg)
|
||||
check_tensorboard_loss_decreased(
|
||||
temp_dir + "/runs",
|
||||
initial_window=5,
|
||||
final_window=5,
|
||||
max_initial=6.0,
|
||||
max_final=4.7,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user