use smaller pretrained models for ci (#3620) [skip ci]

* use smaller pretrained models for ci

* more steps for loss check

* fix tests

* more train steps

* fix losses
This commit is contained in:
Wing Lian
2026-04-27 13:22:56 -04:00
committed by GitHub
parent 798c8fba89
commit ac77da96da
24 changed files with 716 additions and 288 deletions

View File

@@ -13,6 +13,7 @@ from axolotl.utils.dict import DictDefault
from .utils import (
check_model_output_exists,
check_tensorboard_loss_decreased,
require_torch_2_5_1,
require_torch_2_6_0,
require_torch_2_7_0,
@@ -243,20 +244,18 @@ class TestCustomOptimizers(unittest.TestCase):
def test_came_pytorch(self, temp_dir):
cfg = DictDefault(
{
"base_model": "JackFram/llama-68m",
"tokenizer_type": "LlamaTokenizer",
"base_model": "axolotl-ai-co/tiny-llama-50m",
"tokenizer_type": "AutoTokenizer",
"sequence_len": 1024,
"load_in_8bit": True,
"adapter": "lora",
"lora_r": 8,
"lora_alpha": 16,
"lora_dropout": 0.05,
"lora_dropout": 0.0,
"lora_target_linear": True,
"val_set_size": 0.1,
"special_tokens": {
"unk_token": "<unk>",
"bos_token": "<s>",
"eos_token": "</s>",
"pad_token": "<|endoftext|>",
},
"datasets": [
{
@@ -265,16 +264,22 @@ class TestCustomOptimizers(unittest.TestCase):
},
],
"num_epochs": 1,
"sample_packing": True,
"pad_to_sequence_len": True,
"micro_batch_size": 8,
"gradient_accumulation_steps": 1,
"output_dir": temp_dir,
"learning_rate": 0.00001,
"learning_rate": 1e-4,
"optimizer": "came_pytorch",
"adam_beta3": 0.9999,
"adam_epsilon2": 1e-16,
"max_steps": 5,
"max_steps": 80,
"warmup_steps": 5,
"logging_steps": 1,
"lr_scheduler": "cosine",
"save_first_step": False,
"use_tensorboard": True,
"seed": 42,
}
)
@@ -284,6 +289,13 @@ class TestCustomOptimizers(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=10,
final_window=10,
max_initial=4.0,
max_final=3.0,
)
@require_torch_2_7_0