tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

use smaller pretrained models for ci

Browse Source
This commit is contained in:
Wing Lian
2026-04-23 13:51:01 +00:00
parent 1bf65c500e
commit 431888c1de
24 changed files with 614 additions and 205 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
tests/e2e/test_phi.py
View File

@@ -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 TestPhi(unittest.TestCase):
@@ -21,7 +25,7 @@ class TestPhi(unittest.TestCase):
def test_phi_ft(self, temp_dir):
cfg = DictDefault(
{
"base_model": "microsoft/phi-1_5",
"base_model": "axolotl-ai-co/tiny-phi-64m",
"model_type": "AutoModelForCausalLM",
"tokenizer_type": "AutoTokenizer",
"sequence_len": 2048,
@@ -41,18 +45,20 @@ class TestPhi(unittest.TestCase):
"dataset_shard_num": 10,
"dataset_shard_idx": 0,
"num_epochs": 1,
"micro_batch_size": 1,
"micro_batch_size": 4,
"gradient_accumulation_steps": 1,
"output_dir": temp_dir,
"learning_rate": 0.00001,
"optimizer": "paged_adamw_8bit",
"learning_rate": 2e-4,
"optimizer": "adamw_torch_fused",
"lr_scheduler": "cosine",
"flash_attention": True,
"max_steps": 10,
"save_steps": 10,
"eval_steps": 10,
"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)
@@ -61,12 +67,19 @@ class TestPhi(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=5.0,
max_final=4.7,
)
@with_temp_dir
def test_phi_qlora(self, temp_dir):
cfg = DictDefault(
{
"base_model": "microsoft/phi-1_5",
"base_model": "axolotl-ai-co/tiny-phi-64m",
"model_type": "AutoModelForCausalLM",
"tokenizer_type": "AutoTokenizer",
"sequence_len": 2048,
@@ -90,18 +103,20 @@ class TestPhi(unittest.TestCase):
"dataset_shard_num": 10,
"dataset_shard_idx": 0,
"num_epochs": 1,
"micro_batch_size": 1,
"micro_batch_size": 4,
"gradient_accumulation_steps": 1,
"output_dir": temp_dir,
"learning_rate": 0.00001,
"learning_rate": 2e-4,
"optimizer": "paged_adamw_8bit",
"lr_scheduler": "cosine",
"flash_attention": True,
"max_steps": 10,
"save_steps": 10,
"eval_steps": 10,
"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)
@@ -110,3 +125,10 @@ class TestPhi(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=5.0,
max_final=4.7,
)