Is it okay to feed optimizer to configure_optimizers

[sophia1488]

Hi all, is it okay to feed the optimizer that's been initialized outside this code to pl.LightningModule?

def Model(pl.LightningModule):
    def __init__(optimizer):
        self.optimizer = optimizer
    def configure_optimizers(self) -> Any:                                                                                      
        optimizer = self.optimizer          # like this                                                                                    
        scheduler = {                                                                                                               
            'scheduler': LambdaLR(optimizer, self.lr_lambda),                                                                       
            'interval': 'step',                                                                                                     
            'frequency': 1,   
        }                                                                                                                                                                                                                   
    return [optimizer], [scheduler]

Thanks!

[rohitgr7]

yes, I think you can, but not a good practice, we recommend.

just curious, why are you feeding it like that?

[sophia1488]

Hi, thanks for your reply!!
I am using the diffQ quantizer on my model and it requires an optimizer when initializing the quantizer.
But I also need the quantizer for my callback functions, so I just put the quantizer at train.py, and then feed the quantizer to LightningModule and callback.
Cause I don't know how to initialize the quantizer in LightningModule and then pass it to callback...

[rohitgr7]

model = MyModel()
optim = ... # The optimizer must be created before the quantizer
quantizer = DiffQuantizer(model)
quantizer.setup_optimizer(optim)

from what I understand quantizer doesn't need an optimizer during init, but during setup_optimizer.

you can do something like:

class QuantizerCallback(Callback):
    def __init__(self, quantizer):
        self.quantizer = ...

   def on_fit_start(self, trainer, pl_module):
        opt = self.trainer.optimizers[0]  # this is a list
        self.quantizer.setup_optimizer(opt)


class LitModel(LightningModule):
    ...
    
    def configure_optimizers(self):
        opt = ...  # initialize optimizer here
        ...

[sophia1488]

Ohhh It looks nice, I'll try this. Thank you!

[sophia1488]

@rohitgr7
I know it's been a while but the problem still exists, as it turns out that the optimizer must be created before initializing the quantizer.

1. From what I understand, the configure_optimizers is called after __init__, so I'm not sure how to initialize the optimizer in configure_optimizers, before initializing the quantizer.
2. There are only 2 cases when I need the quantizer.
   a. save the quantized model (Either when valid loss is low, or simply save it at every epoch)
   torch.save(quantizer.get_quantized_state(), "some_file.th")
   b. add quantizer.model_size to training loss

class LitModel(LightningModule):
    def __init__():
     
    def training_step():
        loss = loss_function(predict, answer)
        loss += quantizer.model_size()               # quantizer is needed here!
        return loss 

   def configure_optimizers():

Do you have any thoughts on the code structure?

Thank you and hope to get your feedback!

[rohitgr7]

hey @sophia1488!

from the code it doesn't look like the quantizer needs an optimizer during initialization:

model = MyModel()
optim = ...  # The optimizer must be created before the quantizer
quantizer = DiffQuantizer(model)
quantizer.setup_optimizer(optim)

am I missing something here?

if you setup the quantizer inside on_fit_start, you can still access it inside training_step. The hook is present in both LightningModule and Callback. If you are using callback, you can access it using trainer.callbacks (this is a list), or initialize it inside LightningModule itself.

class LitModel(LightningModule):
    ...
    def __init__(self, ...)
        self.quantizer = DiffQuantizer(self)

   def on_fit_start(self):
        opt = self.trainer.optimizers[0]  # this is a list
        self.quantizer.setup_optimizer(opt)

    def training_step(...):
        ...
        loss += self.quantizer.model_size()   
    
    def configure_optimizers(self):
        opt = ...  # initialize optimizer here
        ...

    def on_save_checkpoint(self, checkpoint):
        checkpoint['state_dict'] = quantizer.get_quantized_state()

[sophia1488]

Hello! Thanks for your quick reply!
Yes, the quantizer doesn't need an optimizer during initialization, but during setup, it'd check the optimizer is created before the quantizer.
https://github.com/facebookresearch/diffq/blob/main/diffq/diffq.py#L176
RuntimeError("You should create the optimizer before the quantizer!")

I initialize the quantizer at on_fit_start, and it's working now~

class LitModel(LightningModule):
    ...
    def on_fit_start(self, trainer, pl_module):
        opt = self.trainer.optimizers[0]  # this is a list
        self.quantizer = DiffQuantizer(self.model)      # initialize here
        self.quantizer.setup_optimizer(opt)

It takes 6 hours to train a full epoch, so I'd wait to see if the quantized model is saved as expected!
Thanks a million !!

[rohitgr7]

you can also limit the batches with fast_dev_run to test things out rather than waiting for a complete epoch to finish.

[rohitgr7]

also I have updated the signature for on_fit_start in my example, you might need to update that too :)