From 9c5dbe1286483a20f96c97f99903ac7a3e40e222 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Fri, 15 Nov 2024 15:03:37 -0500 Subject: [PATCH 1/4] support for schedule free and e2e ci smoke test --- requirements.txt | 1 + tests/e2e/test_optimizers.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/requirements.txt b/requirements.txt index 76c3273fcf..8474e04255 100644 --- a/requirements.txt +++ b/requirements.txt @@ -53,3 +53,4 @@ immutabledict==4.2.0 antlr4-python3-runtime==4.13.2 torchao==0.5.0 +schedulefree==1.3.0 diff --git a/tests/e2e/test_optimizers.py b/tests/e2e/test_optimizers.py index 3b68ec5ad5..b321916a71 100644 --- a/tests/e2e/test_optimizers.py +++ b/tests/e2e/test_optimizers.py @@ -108,3 +108,35 @@ def test_adopt_adamw(self, temp_dir): train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta) assert (Path(temp_dir) / "adapter_model.bin").exists() + + @with_temp_dir + def test_fft_schedule_free_adamw(self, temp_dir): + cfg = DictDefault( + { + "base_model": "HuggingFaceTB/SmolLM-135M", + "sequence_len": 1024, + "val_set_size": 0.1, + "special_tokens": { + "pad_token": "<|endoftext|>", + }, + "datasets": [ + { + "path": "mhenrichsen/alpaca_2k_test", + "type": "alpaca", + }, + ], + "num_epochs": 1, + "micro_batch_size": 4, + "gradient_accumulation_steps": 2, + "output_dir": temp_dir, + "learning_rate": 0.00001, + "optimizer": "schedule_free_adamw", + "save_savetensors": True, + } + ) + normalize_config(cfg) + cli_args = TrainerCliArgs() + dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args) + + train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta) + assert (Path(temp_dir) / "model.safetensors").exists() From 2a3d0898394de25caab8b8ca8e4369557444246b Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Fri, 15 Nov 2024 15:05:06 -0500 Subject: [PATCH 2/4] set default lr scheduler to constant in test --- tests/e2e/test_optimizers.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/e2e/test_optimizers.py b/tests/e2e/test_optimizers.py index b321916a71..2b54cdd823 100644 --- a/tests/e2e/test_optimizers.py +++ b/tests/e2e/test_optimizers.py @@ -131,6 +131,7 @@ def test_fft_schedule_free_adamw(self, temp_dir): "output_dir": temp_dir, "learning_rate": 0.00001, "optimizer": "schedule_free_adamw", + "lr_scheduler": "constant", "save_savetensors": True, } ) From fcee5422c47b96eaaf10e1bd4416eba56c20bb4e Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Fri, 15 Nov 2024 15:10:24 -0500 Subject: [PATCH 3/4] ignore duplicate code --- tests/e2e/test_optimizers.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/e2e/test_optimizers.py b/tests/e2e/test_optimizers.py index 2b54cdd823..e410f22d75 100644 --- a/tests/e2e/test_optimizers.py +++ b/tests/e2e/test_optimizers.py @@ -135,6 +135,7 @@ def test_fft_schedule_free_adamw(self, temp_dir): "save_savetensors": True, } ) + # pylint: disable=duplicate-code normalize_config(cfg) cli_args = TrainerCliArgs() dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args) From 56f4865e81f56b8699b3cb628e4af72c4685f531 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Fri, 15 Nov 2024 15:56:02 -0500 Subject: [PATCH 4/4] fix quotes for config/dict --- tests/e2e/test_optimizers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/test_optimizers.py b/tests/e2e/test_optimizers.py index e410f22d75..b9fa368f6f 100644 --- a/tests/e2e/test_optimizers.py +++ b/tests/e2e/test_optimizers.py @@ -132,7 +132,7 @@ def test_fft_schedule_free_adamw(self, temp_dir): "learning_rate": 0.00001, "optimizer": "schedule_free_adamw", "lr_scheduler": "constant", - "save_savetensors": True, + "save_safetensors": True, } ) # pylint: disable=duplicate-code