Add a `dtype` option for `load_from_checkpoint`

[arijit-hub]

Description & Motivation

Hi,
It would be nice to have a dtype argument for load_from_checkpoint along with the very cool map_location. It would allow the user to set the dtype automatically via load_from_checkpoint without the manual need of doing .to(dtype).

Pitch

The fix is should be pretty straightforward. Currently we have:

def load_from_checkpoint(
        cls,
        checkpoint_path: Union[_PATH, IO],
        map_location: _MAP_LOCATION_TYPE = None,
        hparams_file: Optional[_PATH] = None,
        strict: Optional[bool] = None,
        **kwargs: Any,
    ) -> Self:
...

It would just get one more argument:

def load_from_checkpoint(
        cls,
        checkpoint_path: Union[_PATH, IO],
        map_location: _MAP_LOCATION_TYPE = None,
        type:torch.dtype=None,
        hparams_file: Optional[_PATH] = None,
        strict: Optional[bool] = None,
        **kwargs: Any,
    ) -> Self:
...

Given this actually points to _load_from_checkpoint

pytorch-lightning/src/lightning/pytorch/core/saving.py

Line 53 in 3d39824

def _load_from_checkpoint(

We can add our extra dtype argument here too and can easily change this line

pytorch-lightning/src/lightning/pytorch/core/saving.py

Line 99 in 3d39824

return model.to(device)

to

return model.to(dtype).to(device)

Hope this is taken into consideration :)

@Borda @awaelchli @lantiga

Alternatives

No response

Additional context

No response

[matsumotosan]

@arijit-hub I could be misunderstanding the motivation for this, but if it's a simple mismatch between the model parameters and data, would a conversion of the data's type suffice?

If you'd like to convert the model type in the context of post-training quantization, there is an example that might be helpful.

[arijit-hub]

Hi @matsumotosan ,

Thanks for the reference :)

For me, the motivation comes from generative models, where you might want to sample faster using float16 for a model trained in bfloat16 or float32.

I personally felt its much more simpler to do just .to() to typecast into different datatypes. Do you feel that the post-quantization examples that you shared are more practical? If that's the case, I will close the issue :D

Best,
Arijit

[matsumotosan]

@arijit-hub Just to preface, I'm not well-informed on quantization for generative models, so please take my word with a grain of salt.

It is definitely simpler to use .to() to cast the model parameters, but the loss in precision can affect the model's performance. If you have a specific metric that needs to be maintained, then using something like post-training quantization can help you mitigate some of that decrease in performance while getting the boost in inference speed and memory efficiency.