various tests fixes for flakey tests (#2110)

* add mhenrichsen/alpaca_2k_test with revision dataset download fixture for flaky tests

* log slowest tests

* pin pynvml==11.5.3

* fix load local hub path

* optimize for speed w smaller models and val_set_size

* replace pynvml

* make the resume from checkpoint e2e faster

* make tests smaller
This commit is contained in:
Wing Lian
2024-12-02 17:28:58 -05:00
committed by bursteratom
parent b0fbd4d11d
commit c0c53eb62f
13 changed files with 78 additions and 44 deletions

View File

@@ -94,6 +94,7 @@ class TestCustomOptimizers(unittest.TestCase):
},
],
"num_epochs": 1,
"max_steps": 5,
"micro_batch_size": 8,
"gradient_accumulation_steps": 1,
"output_dir": temp_dir,
@@ -115,7 +116,7 @@ class TestCustomOptimizers(unittest.TestCase):
{
"base_model": "HuggingFaceTB/SmolLM2-135M",
"sequence_len": 1024,
"val_set_size": 0.1,
"val_set_size": 0.01,
"special_tokens": {
"pad_token": "<|endoftext|>",
},
@@ -126,13 +127,14 @@ class TestCustomOptimizers(unittest.TestCase):
},
],
"num_epochs": 1,
"micro_batch_size": 4,
"micro_batch_size": 2,
"gradient_accumulation_steps": 2,
"output_dir": temp_dir,
"learning_rate": 0.00001,
"optimizer": "schedule_free_adamw",
"lr_scheduler": "constant",
"save_safetensors": True,
"max_steps": 10,
}
)
# pylint: disable=duplicate-code