Gist on how to use Schedule-Free in Lightning

[CompRhys]

Recent interest has emerged on twitter regarding https://github.com/facebookresearch/schedule_free given promising results without the need to configure a learning rate schedule.

However the optimizer requires .eval() and .train() calls to ensure that the correct parameter buffer is used. Where is the right place to add these additional calls to the methods of a lightning module?

[CompRhys]

class MyNet(pl.LightningModule):

...

    def train(self, mode: bool = True) -> None:
        """Set the model to training mode"""
        self.model.train(mode)
        self.optimizers().train()

    def eval(self) -> None:
        """Set the model to evaluation mode"""
        self.model.eval()
        self.optimizers().eval()

Seemingly the following doesn't cause issues, need to verify whether the results are good. This only would work if the module has a single optimizer

[elyall]

That doesn't seem to work for me after calling trainer.fit(model, train_dataloader):

----> 9) model, trainer = train(model, train_dataloader, valid_dataloader)
...
File /lib/python3.10/site-packages/schedulefree/adamw_schedulefree.py:133, in AdamWScheduleFree.step(self, closure)
...
Exception: Not in train mode!

And when trying to run model.train() beforehand I get a different error:

--> 361) self.optimizers().train()
...
File /lib/python3.10/site-packages/lightning/pytorch/core/module.py:212, in LightningModule.trainer(self)
...
RuntimeError: LitUNet is not attached to a `Trainer`.

Versions:

• torch = 2.2.2
• lightning = 2.2.3
• schedulefree = 1.2.5

[elyall]

Using the on_fit_start and on_predict_start hooks seems to work for a single optimizer (or at least doesn't throw an error):

class MyNet(pl.LightningModule):

...

    def on_fit_start(self) -> None:
        self.optimizers().train()

    def on_predict_start(self) -> None:
        self.optimizers().eval()

Versions

• torch = 2.2.2
• lightning = 2.2.3
• schedulefree = 1.2.5

UPDATE: Looking more at the Hooks pseudocode I believe you also want to set some others as well:

    def on_validation_model_eval(self) -> None:
        self.model.eval()
        self.optimizers().eval()

    def on_validation_model_train(self) -> None:
        self.model.train()
        self.optimizers().train()

    def on_test_model_eval(self) -> None:
        self.model.eval()
        self.optimizers().eval()

    def on_test_model_train(self) -> None:
        self.model.train()
        self.optimizers().train()

    def on_predict_model_eval(self) -> None:  # redundant with on_predict_start()
        self.model.eval()
        self.optimizers().eval()