Repository: Lightning-AI/pytorch-lightning
Pull request: Quant as optional step 
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
 Please also include relevant motivation and context.
 List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)
-->
Fixes #7427

### Does your PR introduce any breaking changes ? If yes, please list them.

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

## Before submitting
- [x] Was this discussed/approved via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your PR does only one thing, instead of bundling different changes together?
- [x] Did you make sure to update the documentation with your changes? (if necessary)
- [ ] Did you write any new necessary tests? (not for typos and docs)
- [ ] Did you verify new and existing tests pass locally with your changes?
- [ ] Did you update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)? (not for typos, docs, test updates, or internal minor changes/refactorings)
- [ ] Did you list all the breaking changes introduced by this pull request?

<!-- For CHANGELOG separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review
Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

 - [ ] Is this pull request ready for review? (if not, please submit in draft mode)
 - [x] Check that all items from **Before submitting** are resolved
 - [x] Make sure the title is self-explanatory and the description concisely explains the PR
 - [x] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?
Make sure you had fun coding 🙃


Problem statement:
Allow calling fit multiple times with Quantization (QAT)
## 🐛 Bug

When using the QuantizationAwareTraining() callback, one cannot call Trainer.fit twice.

### To Reproduce

Call Trainer.fit twice on any model gives an exception (which is pretty qpaque, I might add)

### Expected behavior

Resumes training

### Additional context

I imagine here is a tradeoff between two goals:
1. Have fit return a model that's ready for inference,
2. don't catapult yourself irreversibly out of the option to continue training.

With Quantization Aware Training's conversion step at the end (moving from fake quantization for QAT to quantized layers), we have to choose one. Currently, the QAT hook converts, so it has 1 but not 2.

To my mind, the conversion step is "preparing for export / inference" rather than part of training, so I would suggest to drop it from the fitting part.

Hints from issue discussion:
To avoid the confusion, calling fit twice means something like
```python
dm = BoringDataModule()
model = BoringModel()
t = pl.Trainer()
trainer = pl.Trainer(gpus=1, max_epochs=1, check_val_every_n_epoch=1, callbacks=[QuantizationAwareTraining()])
trainer.fit(model, datamodule=dm)
trainer.fit(model, datamodule=dm)
```

With BoringModel, you I get
```
GPU available: True, used: False
TPU available: False, using: 0 TPU cores
GPU available: True, used: True
TPU available: False, using: 0 TPU cores
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

  | Name    | Type        | Params
----------------------------------------
0 | layer   | Linear      | 66    
1 | quant   | QuantStub   | 0     
2 | dequant | DeQuantStub | 0     
----------------------------------------
66        Trainable params
0         Non-trainable params
66        Total params
0.000     Total estimated model params size (MB)

Epoch 0:  50%|█████     | 64/128 [00:00<00:00, 389.96it/s, loss=0.197, v_num=40]
Validating: 0it [00:00, ?it/s]
Epoch 0: 100%|██████████| 128/128 [00:00<00:00, 546.09it/s, loss=0.197, v_num=40]
Epoch 0: 100%|██████████| 128/128 [00:00<00:00, 540.57it/s, loss=0.197, v_num=40]

LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]


---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-47-27d928f7fc20> in <module>
    152 trainer = pl.Trainer(gpus=1, max_epochs=1, check_val_every_n_epoch=1, callbacks=[QuantizationAwareTraining()])
    153 trainer.fit(model, datamodule=dm)
--> 154 trainer.fit(model, datamodule=dm)

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloader, val_dataloaders, datamodule)
    456         )
    457 
--> 458         self._run(model)
    459 
    460         assert self.state.stopped

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/trainer/trainer.py in _run(self, model)
    713         self.call_setup_hook(model)  # allow user to setup lightning_module in accelerator environment
    714         self.call_configure_sharded_model(model)  # allow user to setup in model sharded environment
--> 715         self.accelerator.setup(self, model)  # note: this sets up self.lightning_module
    716 
    717         # ----------------------------

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/accelerators/gpu.py in setup(self, trainer, model)
     39         self.set_nvidia_flags()
     40         torch.cuda.set_device(self.root_device)
---> 41         return super().setup(trainer, model)
     42 
     43     def on_train_start(self) -> None:

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/accelerators/accelerator.py in setup(self, trainer, model)
     90         self.setup_training_type_plugin(self.training_type_plugin, model)
     91         if not self.training_type_plugin.setup_optimizers_in_pre_dispatch:
---> 92             self.setup_optimizers(trainer)
     93         self.setup_precision_plugin(self.precision_plugin)
     94 

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/accelerators/accelerator.py in setup_optimizers(self, trainer)
    372         if trainer.state.fn not in (TrainerFn.FITTING, TrainerFn.TUNING):
    373             return
--> 374         optimizers, lr_schedulers, optimizer_frequencies = self.training_type_plugin.init_optimizers(
    375             trainer=trainer, model=self.lightning_module
    376         )

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py in init_optimizers(self, trainer, model)
    188 
    189     def init_optimizers(self, trainer: 'pl.Trainer', model: 'pl.LightningModule'):
--> 190         return trainer.init_optimizers(model)
    191 
    192     def optimizer_step(self, optimizer: torch.optim.Optimizer, lambda_closure: Callable, **kwargs):

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/trainer/optimizers.py in init_optimizers(self, model)
     32     def init_optimizers(self, model: LightningModule) -> Tuple[List, List, List]:
     33         self._lightning_optimizers = None
---> 34         optim_conf = model.configure_optimizers()
     35         if optim_conf is None:
     36             rank_zero_warn(

<ipython-input-47-27d928f7fc20> in configure_optimizers(self)
     98 
     99     def configure_optimizers(self):
--> 100         optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.1)
    101         lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
    102         return [optimizer], [lr_scheduler]

/usr/local/lib/python3.9/dist-packages/torch/optim/sgd.py in __init__(self, params, lr, momentum, dampening, weight_decay, nesterov)
     67         if nesterov and (momentum <= 0 or dampening != 0):
     68             raise ValueError("Nesterov momentum requires a momentum and zero dampening")
---> 69         super(SGD, self).__init__(params, defaults)
     70 
     71     def __setstate__(self, state):

/usr/local/lib/python3.9/dist-packages/torch/optim/optimizer.py in __init__(self, params, defaults)
     47         param_groups = list(params)
     48         if len(param_groups) == 0:
---> 49             raise ValueError("optimizer got an empty parameter list")
     50         if not isinstance(param_groups[0], dict):
     51             param_groups = [{'params': param_groups}]

ValueError: optimizer got an empty parameter list
```

If you use a more interesting model (M5 from the PyTorch Audio Keyword detection) that still has parameters after quantization, you get something like

```
NotImplementedError                       Traceback (most recent call last)
<ipython-input-62-a5d40beafe36> in <module>
      2 trainer = pl.Trainer(gpus=1, max_epochs=1, check_val_every_n_epoch=1, callbacks=[QuantizationAwareTraining()])
      3 trainer.fit(model, datamodule=dm)
----> 4 trainer.fit(model, datamodule=dm)

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloader, val_dataloaders, datamodule)
    456         )
    457 
--> 458         self._run(model)
    459 
    460         assert self.state.stopped

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/trainer/trainer.py in _run(self, model)
    754 
    755         # dispatch `start_training` or `start_evaluating` or `start_predicting`
--> 756         self.dispatch()
    757 
    758         # plugin will finalized fitting (e.g. ddp_spawn will load trained model)

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/trainer/trainer.py in dispatch(self)
    795             self.accelerator.start_predicting(self)
    796         else:
--> 797             self.accelerator.start_training(self)
    798 
    799     def run_stage(self):

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/accelerators/accelerator.py in start_training(self, trainer)
     94 
     95     def start_training(self, trainer: 'pl.Trainer') -> None:
---> 96         self.training_type_plugin.start_training(trainer)
     97 
     98     def start_evaluating(self, trainer: 'pl.Trainer') -> None:

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py in start_training(self, trainer)
    142     def start_training(self, trainer: 'pl.Trainer') -> None:
    143         # double dispatch to initiate the training loop
--> 144         self._results = trainer.run_stage()
    145 
    146     def start_evaluating(self, trainer: 'pl.Trainer') -> None:

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/trainer/trainer.py in run_stage(self)
    805         if self.predicting:
    806             return self.run_predict()
--> 807         return self.run_train()
    808 
    809     def _pre_training_routine(self):

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/trainer/trainer.py in run_train(self)
    840             self.progress_bar_callback.disable()
    841 
--> 842         self.run_sanity_check(self.lightning_module)
    843 
    844         self.checkpoint_connector.has_trained = False

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/trainer/trainer.py in run_sanity_check(self, ref_model)
   1105 
   1106             # run eval step
-> 1107             self.run_evaluation()
   1108 
   1109             self.on_sanity_check_end()

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/trainer/trainer.py in run_evaluation(self, on_epoch)
    960                 # lightning module methods
    961                 with self.profiler.profile("evaluation_step_and_end"):
--> 962                     output = self.evaluation_loop.evaluation_step(batch, batch_idx, dataloader_idx)
    963                     output = self.evaluation_loop.evaluation_step_end(output)
    964 

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/trainer/evaluation_loop.py in evaluation_step(self, batch, batch_idx, dataloader_idx)
    172             model_ref._current_fx_name = "validation_step"
    173             with self.trainer.profiler.profile("validation_step"):
--> 174                 output = self.trainer.accelerator.validation_step(args)
    175 
    176         # capture any logged information

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/accelerators/accelerator.py in validation_step(self, args)
    224 
    225         with self.precision_plugin.val_step_context(), self.training_type_plugin.val_step_context():
--> 226             return self.training_type_plugin.validation_step(*args)
    227 
    228     def test_step(self, args: List[Union[Any, int]]) -> Optional[STEP_OUTPUT]:

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py in validation_step(self, *args, **kwargs)
    159 
    160     def validation_step(self, *args, **kwargs):
--> 161         return self.lightning_module.validation_step(*args, **kwargs)
    162 
    163     def test_step(self, *args, **kwargs):

<ipython-input-61-234454f3c7b4> in validation_step(self, batch, batch_idx)
     57     def validation_step(self, batch, batch_idx):
     58         inp, label = batch
---> 59         pred = self(inp)
     60         loss = torch.nn.functional.cross_entropy(pred, label)
     61         acc = self.val_accuracy(pred.softmax(dim=-1), label)

/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1013         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1014                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1015             return forward_call(*input, **kwargs)
   1016         # Do not call functions when jit is used
   1017         full_backward_hooks, non_full_backward_hooks = [], []

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/callbacks/quantization.py in wrapper(data)
     49             quant_cb._forward_calls += 1
     50             data = model.quant(data)
---> 51         data = func(data)
     52         # apply custom trigger
     53         if _quant_run:

/usr/local/lib/python3.9/dist-packages/pytorch_lightning/callbacks/quantization.py in wrapper(data)
     66     def wrapper(data) -> Any:
     67         data = model.quant(data)
---> 68         data = func(data)
     69         data = model.dequant(data)
     70         return data

<ipython-input-61-234454f3c7b4> in forward(self, x)
     45 
     46     def forward(self, x):
---> 47         return self.model(x)
     48 
     49     def training_step(self, batch, batch_idx):

/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1013         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1014                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1015             return forward_call(*input, **kwargs)
   1016         # Do not call functions when jit is used
   1017         full_backward_hooks, non_full_backward_hooks = [], []

<ipython-input-61-234454f3c7b4> in forward(self, x)
     17 
     18     def forward(self, x):
---> 19         x = self.conv1(x)
     20         x = F.relu(self.bn1(x))
     21         x = self.pool1(x)

/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1013         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1014                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1015             return forward_call(*input, **kwargs)
   1016         # Do not call functions when jit is used
   1017         full_backward_hooks, non_full_backward_hooks = [], []

/usr/local/lib/python3.9/dist-packages/torch/nn/quantized/modules/conv.py in forward(self, input)
    328             input = F.pad(input, _reversed_padding_repeated_twice,
    329                           mode=self.padding_mode)
--> 330         return ops.quantized.conv1d(input, self._packed_params, self.scale, self.zero_point)
    331 
    332     @classmethod

NotImplementedError: Could not run 'quantized::conv1d' with arguments from the 'CUDA' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'quantized::conv1d' is only available for these backends: [QuantizedCPU, BackendSelect, Named, InplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, UNKNOWN_TENSOR_TYPE_ID, AutogradMLC, Tracer, Autocast, Batched, VmapMode].

QuantizedCPU: registered at ../aten/src/ATen/native/quantized/cpu/qconv.cpp:876 [kernel]
BackendSelect: fallthrough registered at ../aten/src/ATen/core/BackendSelectFallbackKernel.cpp:3 [backend fallback]
Named: registered at ../aten/src/ATen/core/NamedRegistrations.cpp:7 [backend fallback]
InplaceOrView: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:60 [backend fallback]
AutogradOther: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:35 [backend fallback]
AutogradCPU: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:39 [backend fallback]
AutogradCUDA: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:47 [backend fallback]
AutogradXLA: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:51 [backend fallback]
UNKNOWN_TENSOR_TYPE_ID: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:43 [backend fallback]
AutogradMLC: fallthrough registered at ../aten/src/ATen/core/VariableFallbackKernel.cpp:55 [backend fallback]
Tracer: fallthrough registered at ../torch/csrc/jit/frontend/tracer.cpp:1027 [backend fallback]
Autocast: fallthrough registered at ../aten/src/ATen/autocast_mode.cpp:250 [backend fallback]
Batched: registered at ../aten/src/ATen/BatchingRegistrations.cpp:1016 [backend fallback]
VmapMode: fallthrough registered at ../aten/src/ATen/VmapModeRegistrations.cpp:33 [backend fallback]

```
@Borda mind taking a look?
@t-vi how about adding an extra argument to the callback which would skip this quantization in `on_fit_end` and in such case user needs to call it separately and anytime he/she wants...
So there are two practically irreversible transitions: Model -> QAT-Model -> Quantized model.
Yes, it would make sense to have something to skip the second one and offer it as an extra step. It's really a bit sad that you cannot unquantize a model, but it seems that the PyTorch quantization didn't really appreciate the philosophy of PyTorch as much as I would have wished.

> it seems that the PyTorch quantization didn't really appreciate the philosophy of PyTorch as much as I would have wished.

kind of way around could be replacing the quantized modules with the initial ones and will value from the quantized ones...
what do you think, does it worse to make this kind of reverse engineering?
I think this reversal could be made (you'd have the rounding errors, but hey), but the problem of doing it outside PyTorch that I see is operator coverage. 
> the problem of doing it outside PyTorch that I see is operator coverage

what do you mean by operator coverage?
I mean that if there are quantized operators we don't handle (correctly) then we can't unquantize the model. I'd have to think a bit about how much of a problem I'd expect it to be, it might be not as bad as I thought it was. (Maybe I just wrote too much about operator coverage in those quantization tutorials... :slightly_smiling_face:  ).

Commit message: Apply suggestions from code review

Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com>
Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com>

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2be735ba3f855..4848ec105270c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -177,12 +177,16 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Enabled traditional/manual launching of DDP processes through `LOCAL_RANK` and `NODE_RANK` environment variable assignments ([#7480](https://github.com/PyTorchLightning/pytorch-lightning/pull/7480))
 
 
+- Added `convert_on_fit_end` argument to `QuantizationAwareTraining` ([#8464](https://github.com/PyTorchLightning/pytorch-lightning/pull/8464))
+
+
 - Added experimental support for loop specialization ([#8226](https://github.com/PyTorchLightning/pytorch-lightning/pull/8226))
 
 
 - Added support for `devices` flag to Trainer ([#8440](https://github.com/PyTorchLightning/pytorch-lightning/pull/8440))
 
 
+
 ### Changed
 
 
diff --git a/pytorch_lightning/callbacks/quantization.py b/pytorch_lightning/callbacks/quantization.py
index 9153fbacbf764..1b2dec0a933da 100644
--- a/pytorch_lightning/callbacks/quantization.py
+++ b/pytorch_lightning/callbacks/quantization.py
@@ -82,60 +82,65 @@ def _recursive_hasattr(obj: Any, attribs: str, state: bool = True) -> bool:
 
 
 class QuantizationAwareTraining(Callback):
-    OBSERVER_TYPES = ('histogram', 'average')
+    """
+    Quantization allows speeding up inference and decreasing memory requirements
+    by performing computations and storing tensors at lower bitwidths
+    (such as INT8 or FLOAT16) than floating point precision.
+    We use native PyTorch API so for more information
+    see `Quantization <https://pytorch.org/docs/stable/quantization.html#quantization-aware-training>`_.
+
+    .. warning:: ``QuantizationAwareTraining`` is in beta and subject to change.
 
-    def __init__(
-        self,
-        qconfig: Union[str, QConfig] = 'fbgemm',
-        observer_type: str = "average",
-        collect_quantization: Optional[Union[int, Callable]] = None,
-        modules_to_fuse: Optional[Sequence] = None,
-        input_compatible: bool = True,
-    ) -> None:
-        """
-        Quantization allows speeding up inference and decreasing memory requirements
-        by performing computations and storing tensors at lower bitwidths
-        (such as INT8 or FLOAT16) than floating point precision.
-        We use native PyTorch API so for more information
-        see `Quantization <https://pytorch.org/docs/stable/quantization.html#quantization-aware-training>`_.
 
-        .. warning:: ``QuantizationAwareTraining`` is in beta and subject to change.
+    Args:
 
+        qconfig: quantization configuration:
 
-        Args:
+            - 'fbgemm' for server inference.
+            - 'qnnpack' for mobile inference.
+            -  a custom `torch.quantization.QConfig
+                <https://pytorch.org/docs/stable/torch.quantization.html#torch.quantization.QConfig>`_.
 
-            qconfig: quantization configuration:
+        observer_type: allows switching between ``MovingAverageMinMaxObserver`` as "average" (default)
+            and ``HistogramObserver`` as "histogram" which is more computationally expensive.
 
-                - 'fbgemm' for server inference.
-                - 'qnnpack' for mobile inference.
-                -  a custom `torch.quantization.QConfig <https://pytorch.org/docs/stable/torch.quantization.html#torch.quantization.QConfig>`_.
+        collect_quantization: count or custom function to collect quantization statistics:
 
-            observer_type: allows switching between ``MovingAverageMinMaxObserver`` as "average" (default)
-                and ``HistogramObserver`` as "histogram" which is more computationally expensive.
+            - ``None`` (deafult). The quantization observer is called in each module forward
+                (useful for collecting extended statistic when useing image/data augmentation).
+            - ``int``. Use to set a fixed number of calls, starting from the beginning.
+            - ``Callable``. Custom function with single trainer argument.
+                See this example to trigger only the last epoch:
 
-            collect_quantization: count or custom function to collect quantization statistics:
+                .. code-block:: python
 
-                - ``None`` (deafult). The quantization observer is called in each module forward
-                    (useful for collecting extended statistic when useing image/data augmentation).
-                - ``int``. Use to set a fixed number of calls, starting from the beginning.
-                - ``Callable``. Custom function with single trainer argument.
-                    See this example to trigger only the last epoch:
+                    def custom_trigger_last(trainer):
+                        return trainer.current_epoch == (trainer.max_epochs - 1)
 
-                    .. code-block:: python
+                    QuantizationAwareTraining(collect_quantization=custom_trigger_last)
 
-                        def custom_trigger_last(trainer):
-                            return trainer.current_epoch == (trainer.max_epochs - 1)
+        modules_to_fuse: allows you fuse a few layers together as shown in
+            `diagram <https://pytorch.org/docs/stable/quantization.html#quantization-aware-training>`_
+            to find which layer types can be fused, check https://github.com/pytorch/pytorch/pull/43286.
 
-                        QuantizationAwareTraining(collect_quantization=custom_trigger_last)
+        input_compatible: preserve quant/dequant layers. This allows to feat any input as to the original model,
+            but break compatibility to torchscript and export with ``torch.save``.
 
-            modules_to_fuse: allows you fuse a few layers together as shown in
-                `diagram <https://pytorch.org/docs/stable/quantization.html#quantization-aware-training>`_
-                to find which layer types can be fused, check https://github.com/pytorch/pytorch/pull/43286.
+        convert_on_fit_end: perform the quantization in `on_fit_end`.
+            Note that once converted, the model cannot be put in training mode again.
 
-            input_compatible: preserve quant/dequant layers. This allows to feat any input as to the original model,
-                but break compatibility to torchscript.
+    """
+    OBSERVER_TYPES = ('histogram', 'average')
 
-        """  # noqa: E501
+    def __init__(
+        self,
+        qconfig: Union[str, QConfig] = 'fbgemm',
+        observer_type: str = "average",
+        collect_quantization: Optional[Union[int, Callable]] = None,
+        modules_to_fuse: Optional[Sequence] = None,
+        input_compatible: bool = True,
+        convert_on_fit_end: bool = True,
+    ) -> None:
         _valid_qconf_str = isinstance(qconfig, str) and qconfig in torch.backends.quantized.supported_engines
         if not isinstance(qconfig, QConfig) and not _valid_qconf_str:
             raise MisconfigurationException(
@@ -157,6 +162,7 @@ def custom_trigger_last(trainer):
 
         self.modules_to_fuse = modules_to_fuse
         self._input_compatible = input_compatible
+        self._convert_on_fit_end = convert_on_fit_end
         self._forward_calls = 0
 
     def _check_feasible_fuse(self, model):
@@ -199,6 +205,9 @@ def on_fit_start(self, trainer, pl_module):
         torch.quantization.prepare_qat(pl_module, inplace=True)
 
     def on_fit_end(self, trainer, pl_module):
+        if not self._convert_on_fit_end:
+            pl_module.forward = self.__module_forward
+            return
         pl_module.eval()
         # Convert the observed model to a quantized model. This does several things:
         # quantizes the weights, computes and stores the scale and bias value to be
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
index f6342b5e8e458..4741bdf76dfa5 100644
--- a/pytorch_lightning/trainer/trainer.py
+++ b/pytorch_lightning/trainer/trainer.py
@@ -1227,8 +1227,10 @@ def call_hook(self, hook_name: str, *args, **kwargs) -> Any:
         return output
 
     def _parse_devices(
-        self, gpus: Optional[Union[List[int], str, int]], auto_select_gpus: bool, tpu_cores: Optional[Union[List[int],
-                                                                                                            str, int]]
+        self,
+        gpus: Optional[Union[List[int], str, int]],
+        auto_select_gpus: bool,
+        tpu_cores: Optional[Union[List[int], str, int]],
     ) -> Tuple[Optional[List[int]], Optional[Union[List[int], int]]]:
         if auto_select_gpus and isinstance(gpus, int):
             gpus = pick_multiple_gpus(gpus)
diff --git a/tests/callbacks/test_quantization.py b/tests/callbacks/test_quantization.py
index 8627acd23f6be..a102346852625 100644
--- a/tests/callbacks/test_quantization.py
+++ b/tests/callbacks/test_quantization.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import copy
+import os
 from typing import Callable, Union
 
 import pytest
@@ -28,15 +29,16 @@
 
 @pytest.mark.parametrize("observe", ['average', 'histogram'])
 @pytest.mark.parametrize("fuse", [True, False])
+@pytest.mark.parametrize("convert", [True, False])
 @RunIf(quantization=True)
-def test_quantization(tmpdir, observe: str, fuse: bool):
+def test_quantization(tmpdir, observe: str, fuse: bool, convert: bool):
     """Parity test for quant model"""
     seed_everything(42)
     dm = RegressDataModule()
     trainer_args = dict(
         default_root_dir=tmpdir,
-        max_epochs=10,
-        gpus=1 if torch.cuda.is_available() else None,
+        max_epochs=7,
+        gpus=torch.cuda.device_count(),
     )
     model = RegressionModel()
     qmodel = copy.deepcopy(model)
@@ -47,20 +49,39 @@ def test_quantization(tmpdir, observe: str, fuse: bool):
     org_score = torch.mean(torch.tensor([mean_relative_error(model(x), y) for x, y in dm.test_dataloader()]))
 
     fusing_layers = [(f'layer_{i}', f'layer_{i}a') for i in range(3)] if fuse else None
-    qcb = QuantizationAwareTraining(observer_type=observe, modules_to_fuse=fusing_layers)
+    qcb = QuantizationAwareTraining(observer_type=observe, modules_to_fuse=fusing_layers, convert_on_fit_end=convert)
     trainer = Trainer(callbacks=[qcb], **trainer_args)
     trainer.fit(qmodel, datamodule=dm)
 
     quant_calls = qcb._forward_calls
     assert quant_calls == qcb._forward_calls
+    quant_score = torch.mean(torch.tensor([mean_relative_error(qmodel(x), y) for x, y in dm.test_dataloader()]))
+    # test that the test score is almost the same as with pure training
+    assert torch.allclose(org_score, quant_score, atol=0.45)
+    model_path = trainer.checkpoint_callback.best_model_path
+
+    trainer_args.update(dict(max_epochs=1, checkpoint_callback=False))
+    if not convert:
+        trainer = Trainer(callbacks=[QuantizationAwareTraining()], **trainer_args)
+        trainer.fit(qmodel, datamodule=dm)
+        qmodel.eval()
+        torch.quantization.convert(qmodel, inplace=True)
 
     quant_size = qmodel.model_size
-    quant_score = torch.mean(torch.tensor([mean_relative_error(qmodel(x), y) for x, y in dm.test_dataloader()]))
     # test that the trained model is smaller then initial
     size_ratio = quant_size / org_size
     assert size_ratio < 0.65
-    # test that the test score is almost the same as with pure training
-    assert torch.allclose(org_score, quant_score, atol=0.45)
+
+    # todo: make it work also with strict loading
+    qmodel2 = RegressionModel.load_from_checkpoint(model_path, strict=False)
+    quant2_score = torch.mean(torch.tensor([mean_relative_error(qmodel2(x), y) for x, y in dm.test_dataloader()]))
+    assert torch.allclose(org_score, quant2_score, atol=0.45)
+
+    trainer = Trainer(**trainer_args)
+    trainer.fit(qmodel2, datamodule=dm)
+
+    trainer = Trainer(callbacks=[QuantizationAwareTraining()], **trainer_args)
+    trainer.fit(qmodel2, datamodule=dm)
 
 
 @RunIf(quantization=True)


Code review comments:
diff_hunk: @@ -82,60 +82,65 @@ def _recursive_hasattr(obj: Any, attribs: str, state: bool = True) -> bool:
 
 
 class QuantizationAwareTraining(Callback):
-    OBSERVER_TYPES = ('histogram', 'average')
+    """
+    Quantization allows speeding up inference and decreasing memory requirements
+    by performing computations and storing tensors at lower bitwidths
+    (such as INT8 or FLOAT16) than floating point precision.
+    We use native PyTorch API so for more information
+    see `Quantization <https://pytorch.org/docs/stable/quantization.html#quantization-aware-training>`_.
+
+    .. warning:: ``QuantizationAwareTraining`` is in beta and subject to change.
 
-    def __init__(
-        self,
-        qconfig: Union[str, QConfig] = 'fbgemm',
-        observer_type: str = "average",
-        collect_quantization: Optional[Union[int, Callable]] = None,
-        modules_to_fuse: Optional[Sequence] = None,
-        input_compatible: bool = True,
-    ) -> None:
-        """
-        Quantization allows speeding up inference and decreasing memory requirements
-        by performing computations and storing tensors at lower bitwidths
-        (such as INT8 or FLOAT16) than floating point precision.
-        We use native PyTorch API so for more information
-        see `Quantization <https://pytorch.org/docs/stable/quantization.html#quantization-aware-training>`_.
 
-        .. warning:: ``QuantizationAwareTraining`` is in beta and subject to change.
+    Args:
 
+        qconfig: quantization configuration:
 
-        Args:
+            - 'fbgemm' for server inference.
+            - 'qnnpack' for mobile inference.
+            -  a custom `torch.quantization.QConfig
+                <https://pytorch.org/docs/stable/torch.quantization.html#torch.quantization.QConfig>`_.
 
-            qconfig: quantization configuration:
+        observer_type: allows switching between ``MovingAverageMinMaxObserver`` as "average" (default)
+            and ``HistogramObserver`` as "histogram" which is more computationally expensive.
 
-                - 'fbgemm' for server inference.
-                - 'qnnpack' for mobile inference.
-                -  a custom `torch.quantization.QConfig <https://pytorch.org/docs/stable/torch.quantization.html#torch.quantization.QConfig>`_.
+        collect_quantization: count or custom function to collect quantization statistics:
 
-            observer_type: allows switching between ``MovingAverageMinMaxObserver`` as "average" (default)
-                and ``HistogramObserver`` as "histogram" which is more computationally expensive.
+            - ``None`` (deafult). The quantization observer is called in each module forward
+                (useful for collecting extended statistic when useing image/data augmentation).
+            - ``int``. Use to set a fixed number of calls, starting from the beginning.
+            - ``Callable``. Custom function with single trainer argument.
+                See this example to trigger only the last epoch:
 
-            collect_quantization: count or custom function to collect quantization statistics:
+                .. code-block:: python
 
-                - ``None`` (deafult). The quantization observer is called in each module forward
-                    (useful for collecting extended statistic when useing image/data augmentation).
-                - ``int``. Use to set a fixed number of calls, starting from the beginning.
-                - ``Callable``. Custom function with single trainer argument.
-                    See this example to trigger only the last epoch:
+                    def custom_trigger_last(trainer):
+                        return trainer.current_epoch == (trainer.max_epochs - 1)
 
-                    .. code-block:: python
+                    QuantizationAwareTraining(collect_quantization=custom_trigger_last)
 
-                        def custom_trigger_last(trainer):
-                            return trainer.current_epoch == (trainer.max_epochs - 1)
+        modules_to_fuse: allows you fuse a few layers together as shown in
+            `diagram <https://pytorch.org/docs/stable/quantization.html#quantization-aware-training>`_
+            to find which layer types can be fused, check https://github.com/pytorch/pytorch/pull/43286.
 
-                        QuantizationAwareTraining(collect_quantization=custom_trigger_last)
+        input_compatible: preserve quant/dequant layers. This allows to feat any input as to the original model,
+            but break compatibility to torchscript and export with ``torch.save``.
 
-            modules_to_fuse: allows you fuse a few layers together as shown in
-                `diagram <https://pytorch.org/docs/stable/quantization.html#quantization-aware-training>`_
-                to find which layer types can be fused, check https://github.com/pytorch/pytorch/pull/43286.
+        convert_on_fit_end: perform the quantization in `on_fit_end`.
+            Note that once converted, the model cannot be put in training mode again.
 
-            input_compatible: preserve quant/dequant layers. This allows to feat any input as to the original model,
-                but break compatibility to torchscript.
+    """
+    OBSERVER_TYPES = ('histogram', 'average')
 
-        """  # noqa: E501
+    def __init__(
+        self,
+        qconfig: Union[str, QConfig] = 'fbgemm',
+        observer_type: str = "average",
+        collect_quantization: Optional[Union[int, Callable]] = None,
+        modules_to_fuse: Optional[Sequence] = None,
+        input_compatible: bool = True,
+        convert_on_fit_end: bool = True,
path: pytorch_lightning/callbacks/quantization.py
text: @user1:
Maybe rename to `quantize_on_fit_end`?

diff_hunk: @@ -47,20 +49,39 @@ def test_quantization(tmpdir, observe: str, fuse: bool):
     org_score = torch.mean(torch.tensor([mean_relative_error(model(x), y) for x, y in dm.test_dataloader()]))
 
     fusing_layers = [(f'layer_{i}', f'layer_{i}a') for i in range(3)] if fuse else None
-    qcb = QuantizationAwareTraining(observer_type=observe, modules_to_fuse=fusing_layers)
+    qcb = QuantizationAwareTraining(observer_type=observe, modules_to_fuse=fusing_layers, convert_on_fit_end=convert)
     trainer = Trainer(callbacks=[qcb], **trainer_args)
     trainer.fit(qmodel, datamodule=dm)
 
     quant_calls = qcb._forward_calls
     assert quant_calls == qcb._forward_calls
+    quant_score = torch.mean(torch.tensor([mean_relative_error(qmodel(x), y) for x, y in dm.test_dataloader()]))
+    # test that the test score is almost the same as with pure training
+    assert torch.allclose(org_score, quant_score, atol=0.45)
+    model_path = trainer.checkpoint_callback.best_model_path
+
+    trainer_args.update(dict(max_epochs=1, checkpoint_callback=False))
+    if not convert:
+        trainer = Trainer(callbacks=[QuantizationAwareTraining()], **trainer_args)
+        trainer.fit(qmodel, datamodule=dm)
+        qmodel.eval()
+        torch.quantization.convert(qmodel, inplace=True)
 
     quant_size = qmodel.model_size
-    quant_score = torch.mean(torch.tensor([mean_relative_error(qmodel(x), y) for x, y in dm.test_dataloader()]))
     # test that the trained model is smaller then initial
     size_ratio = quant_size / org_size
     assert size_ratio < 0.65
-    # test that the test score is almost the same as with pure training
-    assert torch.allclose(org_score, quant_score, atol=0.45)
+
+    # todo: make it work also with strict loading
+    qmodel2 = RegressionModel.load_from_checkpoint(model_path, strict=False)
+    quant2_score = torch.mean(torch.tensor([mean_relative_error(qmodel2(x), y) for x, y in dm.test_dataloader()]))
+    assert torch.allclose(org_score, quant2_score, atol=0.45)
+
+    trainer = Trainer(**trainer_args)
+    trainer.fit(qmodel2, datamodule=dm)
+
+    trainer = Trainer(callbacks=[QuantizationAwareTraining()], **trainer_args)
+    trainer.fit(qmodel2, datamodule=dm)
path: tests/callbacks/test_quantization.py
text: @user1:
What are these testing?
```suggestion
```

No assertions?

@author:
that was requested by @user2-vi that the double training did not work... #7427

@user1:
But you already have 2 `fit` calls for `qmodel` if `not convert`. So the reported issue is already tested

Do you expect that the behaviour for `qmodel2` (loaded for checkpoint) would be different?