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

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

Browse Source 
* 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
Download Patch File Download Diff File Expand all files Collapse all files

8
tests/e2e/patched/lora_kernels/test_lora_kernel_patching.py
View File

@@ -32,12 +32,12 @@ from axolotl.utils.dict import DictDefault
MODEL_CONFIGS = [
{
"name": "trl-internal-testing/tiny-MistralForCausalLM-0.2",
"name": "axolotl-ai-co/tiny-mistral-25m",
"expected_activation": apply_lora_mlp_swiglu,
"dtype": torch.float16,
},
{
"name": "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
"name": "axolotl-ai-co/tiny-qwen2-129m",
"expected_activation": apply_lora_mlp_swiglu,
"dtype": torch.float16,
},
@@ -47,7 +47,7 @@ MODEL_CONFIGS = [
"dtype": torch.float32,
},
{
"name": "trl-internal-testing/tiny-Gemma2ForCausalLM",
"name": "axolotl-ai-co/tiny-gemma2-137m",
"expected_activation": apply_lora_mlp_geglu,
"dtype": torch.float16,
},
@@ -159,7 +159,7 @@ def test_swiglu_mlp_integration(small_llama_model):
def test_geglu_model_integration():
"""Test GeGLU activation with Gemma model."""
model = AutoModelForCausalLM.from_pretrained(
"trl-internal-testing/tiny-Gemma2ForCausalLM",
"axolotl-ai-co/tiny-gemma2-137m",
dtype=torch.float16,
device_map="cuda:0",
)

58
tests/e2e/patched/test_falcon_samplepack.py
View File

@@ -4,14 +4,16 @@ E2E tests for falcon
import unittest
import pytest
from axolotl.common.datasets import load_datasets
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 TestFalconPatched(unittest.TestCase):
@@ -19,13 +21,12 @@ class TestFalconPatched(unittest.TestCase):
Test case for Falcon models
"""
@pytest.mark.skip(reason="no tiny models for testing with safetensors")
@with_temp_dir
def test_qlora(self, temp_dir):
cfg = DictDefault(
{
"base_model": "illuin/tiny-random-FalconForCausalLM",
"flash_attention": True,
"base_model": "axolotl-ai-co/tiny-falcon-42m",
"flash_attention": False,
"sample_packing": True,
"sequence_len": 2048,
"load_in_4bit": True,
@@ -47,17 +48,20 @@ class TestFalconPatched(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": 20,
"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,
"seed": 42,
}
)
cfg = validate_config(cfg)
@@ -66,14 +70,20 @@ class TestFalconPatched(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,
)
@pytest.mark.skip(reason="no tiny models for testing with safetensors")
@with_temp_dir
def test_ft(self, temp_dir):
cfg = DictDefault(
{
"base_model": "illuin/tiny-random-FalconForCausalLM",
"flash_attention": True,
"base_model": "axolotl-ai-co/tiny-falcon-42m",
"flash_attention": False,
"sample_packing": True,
"sequence_len": 2048,
"val_set_size": 0.05,
@@ -88,17 +98,20 @@ class TestFalconPatched(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": 20,
"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,
"seed": 42,
}
)
cfg = validate_config(cfg)
@@ -107,3 +120,10 @@ class TestFalconPatched(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,
)

51
tests/e2e/patched/test_mistral_samplepack.py
View File

@@ -9,7 +9,12 @@ 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, require_torch_2_6_0, with_temp_dir
from ..utils import (
check_model_output_exists,
check_tensorboard_loss_decreased,
require_torch_2_6_0,
with_temp_dir,
)
class TestMistral(unittest.TestCase):
@@ -22,7 +27,7 @@ class TestMistral(unittest.TestCase):
def test_lora_packing(self, temp_dir):
cfg = DictDefault(
{
"base_model": "trl-internal-testing/tiny-MistralForCausalLM-0.2",
"base_model": "axolotl-ai-co/tiny-mistral-25m",
"flash_attention": True,
"sample_packing": True,
"sequence_len": 1024,
@@ -45,17 +50,20 @@ class TestMistral(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_torch_fused",
"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,
"seed": 42,
}
)
cfg = validate_config(cfg)
@@ -64,12 +72,19 @@ class TestMistral(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.5,
max_final=4.3,
)
@with_temp_dir
def test_ft_packing(self, temp_dir):
cfg = DictDefault(
{
"base_model": "trl-internal-testing/tiny-MistralForCausalLM-0.2",
"base_model": "axolotl-ai-co/tiny-mistral-25m",
"flash_attention": True,
"sample_packing": True,
"sequence_len": 1024,
@@ -86,17 +101,20 @@ class TestMistral(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_torch_fused",
"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,
"seed": 42,
}
)
cfg = validate_config(cfg)
@@ -105,3 +123,10 @@ class TestMistral(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.5,
max_final=4.3,
)

58
tests/e2e/patched/test_mixtral_samplepack.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 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,
@@ -30,7 +33,7 @@ class TestMixtral(unittest.TestCase):
"adapter": "qlora",
"lora_r": 16,
"lora_alpha": 32,
"lora_dropout": 0.1,
"lora_dropout": 0.0,
"lora_target_linear": True,
"val_set_size": 0.05,
"special_tokens": {},
@@ -41,17 +44,21 @@ 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": 3e-3,
"optimizer": "adamw_bnb_8bit",
"lr_scheduler": "cosine",
"max_steps": 5,
"save_steps": 3,
"eval_steps": 4,
"max_steps": 80,
"warmup_steps": 5,
"logging_steps": 1,
"save_steps": 80,
"eval_steps": 80,
"bf16": "auto",
"save_first_step": False,
"use_tensorboard": True,
"seed": 42,
}
)
cfg = validate_config(cfg)
@@ -60,13 +67,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=10,
final_window=10,
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 +92,21 @@ 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,
"optimizer": "adamw_bnb_8bit",
"learning_rate": 5e-4,
"optimizer": "adamw_torch_fused",
"lr_scheduler": "cosine",
"max_steps": 5,
"save_steps": 3,
"eval_steps": 4,
"max_steps": 80,
"warmup_steps": 5,
"logging_steps": 1,
"save_steps": 80,
"eval_steps": 80,
"bf16": "auto",
"save_first_step": False,
"use_tensorboard": True,
"seed": 42,
}
)
cfg = validate_config(cfg)
@@ -98,3 +115,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,
)

5
tests/e2e/patched/test_model_patches.py
View File

@@ -22,8 +22,7 @@ class TestModelPatches(unittest.TestCase):
def test_mixtral_multipack(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,
@@ -57,7 +56,7 @@ class TestModelPatches(unittest.TestCase):
def test_mistral_multipack(self, temp_dir):
cfg = DictDefault(
{
"base_model": "trl-internal-testing/tiny-MistralForCausalLM-0.2",
"base_model": "axolotl-ai-co/tiny-mistral-25m",
"flash_attention": True,
"sample_packing": True,
"sequence_len": 2048,

52
tests/e2e/patched/test_phi_multipack.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 TestPhiMultipack(unittest.TestCase):
@@ -21,7 +25,7 @@ class TestPhiMultipack(unittest.TestCase):
def test_ft_packed(self, temp_dir):
cfg = DictDefault(
{
"base_model": "microsoft/phi-1_5",
"base_model": "axolotl-ai-co/tiny-phi-64m",
"model_type": "PhiForCausalLM",
"tokenizer_type": "AutoTokenizer",
"sequence_len": 1024,
@@ -43,17 +47,20 @@ class TestPhiMultipack(unittest.TestCase):
"dataset_shard_num": 10,
"dataset_shard_idx": 0,
"num_epochs": 1,
"micro_batch_size": 1,
"micro_batch_size": 2,
"gradient_accumulation_steps": 1,
"output_dir": temp_dir,
"learning_rate": 0.00001,
"optimizer": "adamw_bnb_8bit",
"learning_rate": 2e-4,
"optimizer": "adamw_torch_fused",
"lr_scheduler": "cosine",
"max_steps": 5,
"eval_steps": 3,
"save_steps": 4,
"max_steps": 50,
"logging_steps": 1,
"eval_steps": 50,
"save_steps": 50,
"bf16": "auto",
"save_first_step": False,
"use_tensorboard": True,
"seed": 42,
}
)
@@ -63,12 +70,19 @@ class TestPhiMultipack(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_qlora_packed(self, temp_dir):
cfg = DictDefault(
{
"base_model": "microsoft/phi-1_5",
"base_model": "axolotl-ai-co/tiny-phi-64m",
"model_type": "PhiForCausalLM",
"tokenizer_type": "AutoTokenizer",
"sequence_len": 1024,
@@ -94,17 +108,20 @@ class TestPhiMultipack(unittest.TestCase):
"dataset_shard_num": 10,
"dataset_shard_idx": 0,
"num_epochs": 1,
"micro_batch_size": 1,
"micro_batch_size": 2,
"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,
"eval_steps": 3,
"save_steps": 4,
"max_steps": 50,
"logging_steps": 1,
"eval_steps": 50,
"save_steps": 50,
"bf16": "auto",
"save_first_step": False,
"use_tensorboard": True,
"seed": 42,
}
)
@@ -114,3 +131,10 @@ class TestPhiMultipack(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,
)