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
2026-04-27 13:22:56 -04:00
parent 798c8fba89
commit ac77da96da
24 changed files with 716 additions and 288 deletions
--- a/tests/e2e/test_quantization.py
+++ b/tests/e2e/test_quantization.py
@@ -45,7 +45,7 @@ def _get_fake_quant_config_dtype(config):
@pytest.fixture()
 def model():
    dummy_model = AutoModelForCausalLM.from_pretrained(
-        "Qwen/Qwen2-0.5B",
+        "axolotl-ai-co/tiny-qwen2-129m",
        device_map="auto",
        dtype=torch.bfloat16,
    )