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

Plugins create_lr_scheduler support (#2584)

Browse Source 
* lr_scheduler support

* fix

* Update scheduler.py

* Update scheduler.py

* cfg handling

* black

* remove debug

* remove adding the axolotl cfg to the scheduler mixin

---------

Co-authored-by: Wing Lian <wing@axolotl.ai>
This commit is contained in:
Aleksandr Dremov
2025-04-29 23:08:30 +02:00
committed by GitHub
parent ecac731922
commit 41a1ec0c95
4 changed files with 40 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/e2e/integrations/test_hooks.py
View File

@@ -72,7 +72,7 @@ class LogHooksPlugin(BasePlugin):
f.write("get_trainer_cls\n")
def create_lr_scheduler(
self, cfg, trainer, optimizer
self, cfg, trainer, optimizer, num_training_steps
): # pylint: disable=unused-argument
with open(
self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
@@ -172,7 +172,7 @@ class TestPluginHooks:
assert "post_model_load" in file_contents
# assert "create_optimizer" in file_contents # not implemented yet
assert "get_trainer_cls" in file_contents
# assert "create_lr_scheduler" in file_contents # not implemented yet
assert "create_lr_scheduler" in file_contents
assert "add_callbacks_pre_trainer" in file_contents
assert "add_callbacks_post_trainer" in file_contents
assert "post_train" in file_contents