pytorch-lightning: Can't load_from_checkpoint with QuantizationAwareTraining callback during training

The error in load_from_checkpoint is caused be model quantization preparation is done by checkpoint hook, which is not triggered by PL module’s load_from_checkpoint API. At the checkpoint loading time, Lightning module is not prepared for quantization but the checkpoint has quantization states.

Have we considered to delegate quantization preparation and conversion to PL module itself? In this way, we could call preparation and conversion in module’s on_load_checkpoint hook. This is the design in my mind:

class QuantizationMixin(ABC):
    def prepare_for_quant(self):
        # insert quant and dequant
        # prepare_qat

    def convert_for_quant(self):
        # convert

    def setup_for_quant_checkpoint(self):
        # if model in checkpoint has been prepared
        # self.prepare_for_quant()
        # if model in checkpoint has been convert
        # self.convert_for_quant()

class MyModule(pl.LightningModule, QuantizationMixin):
    def on_load_checkpoint(self, checkpoint):
        self.setup_for_quant_checkpoint()

class QuantizationAwareTraining(Callback):
    def on_fit_start(self, trainer, pl_module):
        pl_module.prepare_for_quant()

    def on_fit_end(self, trainer, pl_module):
        pl_module.convert_for_quant()

This design gives several benefits:

1. The checkpoint could be loaded outside of training;
2. Module writer can override prepare_for_quant and convert_for_quant to define their own customization;
3. Simplifies quantization callback initiation. The same callback could be used across different modules.

hi, So how should we use quantitative weights? Calculate quantization weights based on weights and statistical parameters?

I see two options:

• you would need to save the model after training 🚀
• adding extra arguments which would save as checkpoint already quantized model 👍

@Borda by “integrate it with the main LightningMdule”, do you mean adding model preparation and conversion function to the core PyTorch LightningModule? Well it does simplify the requirements for user to use the quant callback, but I’m not sure if it’s the right thing to do - adding a function to core LightningModule just for a specific callback.

Is there any temporary fix to apply? QAT is basically unusable without checkpoint loading