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()

[faroit]

is there a way to do this without overriding hooks (e.g. modifying the trainer) as this can lead to a lot of additional code when this has to be done for multiple systems?

[elyall]

Not as far as I can tell currently. I would just create a new LightningModule subclass with these changes and then update your code so that your models subclass it instead, which shouldn't be that much additional code. However if those models all have unique changes to these hooks already then I could see how that would result in a lot of code.

[hendriks73]

This is interesting!
What did you end up using? eval/train calls on all 7 hooks?

[abcamiletto]

This solution looks a bit clumsy. Is there any suggestion on doing this cleanly? Maybe a wrapper that hides this complexity in lightning_bolts?

[maxfreu]

The schedulerfree docs state that the scheduler should be put into eval mode when saving checkpoints. How is this best achieved?

[elyall]

Looks like there's another hook that gets called before it saves that you can instantiate: on_save_checkpoint. e.g.:

    def on_save_checkpoint(self) -> None:
        self.model.eval()  # probably not necessary
        self.optimizers().eval()

[maxfreu]

That was also my first guess - but unfortunately this would lead to a silent bug. I looked into the saving code and on_save_checkpoint gets called after the optimizer state was read.

Here the optmizer state is read:

pytorch-lightning/src/lightning/pytorch/trainer/connectors/checkpoint_connector.py

Lines 444 to 450 in 3627c5b

	optimizer_states = []
	for i, optimizer in enumerate(trainer.optimizers):
	# Rely on accelerator to dump optimizer state
	optimizer_state = trainer.strategy.optimizer_state(optimizer)
	optimizer_states.append(optimizer_state)
	checkpoint["optimizer_states"] = optimizer_states

And a few lines further down the callback is called:

pytorch-lightning/src/lightning/pytorch/trainer/connectors/checkpoint_connector.py

Lines 487 to 492 in 3627c5b

	if not weights_only:
	# if state is returned from callback's on_save_checkpoint
	# it overrides the returned state from callback's state_dict
	# support for returning state in on_save_checkpoint
	# will be removed in v1.8
	call._call_callbacks_on_save_checkpoint(trainer, checkpoint)

So one could no go two routes: 1) Copy the code for setting the optimizer state into the callback and set the opimizer to eval beforehand or 2) subclass the ModelCheckpoint callback and overwrite the _save_checkpoint method. I went the second way:

from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint
from schedulefree import AdamWScheduleFree


class SchedulerFreeModelCheckpoint(ModelCheckpoint):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def _save_checkpoint(self, trainer: Trainer, filepath: str) -> None:
        opts = trainer.optimizers
        if not isinstance(opts, list):
            opts = [opts]

        # save current train/eval state so that we can restore them
        states = []
        for opt in opts:
            if isinstance(opt, AdamWScheduleFree):
                # probe optimizer state by checking the train state of the first parameter group
                is_train_mode = opt.param_groups[0]["train_mode"]
                if is_train_mode:
                    opt.eval()
            else:
                is_train_mode = None
            states.append(is_train_mode)

        super()._save_checkpoint(trainer, filepath)

        # restore previous state
        for (opt, wasintrainmode) in zip(opts, states):
            if isinstance(opt, AdamWScheduleFree):
                if wasintrainmode:
                    opt.train()

In the model I set things via:

def set_optimizer_state(self, state: str):
    opts = self.optimizers()
    if not isinstance(opts, list):
        opts = [opts]

    for opt in opts:
        if isinstance(opt, AdamWScheduleFree):
            if state == "train":
                opt.train()
            elif state == "eval":
                opt.eval()
            else:
                raise ValueError(f"Unknown train state {state}")

def on_train_epoch_start(self) -> None:
    self.set_optimizer_state("train")

def on_validation_epoch_start(self) -> None:
    self.set_optimizer_state("eval")