Adding defaults to hparams in base class LightningModule causes incorrect _hparams_name in checkpoints

Bug description

Scenario

Consider a LightningModule that you want configurable via some hparams dictionary, adding some defaults for reproducibility and usability:

class MyBaseModel(LightningModule):
    def __init__(self, hparams: dict):
        super().__init__()
        hparams = some_defaults | hparams
        self.save_hyperparameters(hparams)

If users want to create variations of this model, they would write something like

class MyDerivedModel(MyBaseModel):
    def __init__(self, hparams):
        super().__init__(hparams)

Problem

However, this model cannot be loaded from a checkpoint:

>>> MyDerivedModel.load_from_checkpoint(checkpoint_path)
TypeError: MyDerivedModel.__init__() missing 1 required positional argument: 'hparams'

This cryptic message is actually caused by the fact that the *.ckpt files do not contain the proper _hparams_name if we set the default values in the MyBaseModel (it will be None instead of "hparams").

Why does this happen?

In save_hyperparameters, we assign the _hparams_name according to the name used in the stack frame above, which is saved in init_args. However, there are two problems with this

1. init_args is created not only from one frame level up, but recursively from all levels above.
2. _hparams_name is only assigned if the argument passed to save_hyperparameters is equal (==) to the argument in init_args

1 leads to a name clash in init_args due to the inheritance, where the hparams with the filled defaults is overwritten with the one without any defaults (since both classes named the parameter the same). Thus, 2 does not set the _hparams_name appropriately.

How to fix?

My proposed fix would be to retrieve the _hparams_name regardless of the equality, i.e. change the following lines in save_hyperparameters:

From

hp = args[isx_non_str[0]]
cand_names = [k for k, v in init_args.items() if v == hp]
obj._hparams_name = cand_names[0] if cand_names else None

To

hp = args[isx_non_str[0]]
obj._hparams_name = list(init_args.keys())[isx_non_str[0]]

I think it would also be useful to fix the name clash in init_args, unless this is somehow desired behaviour. This could be done e.g. by replacing

for local_args in collect_init_args(frame, [], classes=(HyperparametersMixin,)):
    init_args.update(local_args)

with

local_args = collect_init_args(frame, [], classes=(HyperparametersMixin,))[0]
init_args.update(local_args)

Is there a workaround?

Yes. Simply set the _hparams_name manually in the base model:

class MyBaseModel(LightningModule):
    def __init__(self, hparams: dict):
        super().__init__()
        hparams = some_defaults | hparams
        self.save_hyperparameters(hparams)
        self._hparams_name = "hparams"

What version are you seeing the problem on?

v2.0

How to reproduce the bug

""" This is a complete script to reproduce the issue """

import torch
from torch.utils.data import TensorDataset, DataLoader

import lightning
import lightning.pytorch.callbacks as callbacks


class MyBaseModel(lightning.LightningModule):
    base_defaults = {"some_string": "hello"}

    def __init__(self, hparams: dict):
        super().__init__()
        hparams = hparams | self.base_defaults
        self.save_hyperparameters(hparams)
        self.param = torch.nn.Parameter(torch.randn(8, 1))

    def training_step(self, batch, batch_idx):
        return torch.tensor(0.0, requires_grad=True)

    def validation_step(self, batch, batch_idx):
        return torch.tensor(0.0, requires_grad=True)

    def configure_callbacks(self):
        return [
            callbacks.ModelCheckpoint(save_last=True),
        ]

    def configure_optimizers(self):
        return None

    def train_dataloader(self):
        train_data = TensorDataset(torch.randn(128, 1))
        return DataLoader(train_data, batch_size=8)

    def val_dataloader(self):
        val_data = TensorDataset(torch.randn(128, 1))
        return DataLoader(val_data, batch_size=8)


class MyDerivedModel(MyBaseModel):
    def __init__(self, hparams):
        super().__init__(hparams)


hparams = dict(
    some_int=42
)

model = MyDerivedModel(hparams)

trainer = lightning.Trainer(max_epochs=1)

trainer.fit(model)

checkpoint = trainer.checkpoint_callback.best_model_path

# this line fails
model = MyDerivedModel.load_from_checkpoint(checkpoint)

Error messages and logs

TypeError: MyDerivedModel.__init__() missing 1 required positional argument: 'hparams'

Environment

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):

More info

No response

Start at save_hyperparameters and collect_init_args, then use the provided MyBaseModel/MyDerivedModel reproduction to inspect how init_args and _hparams_name are recorded. Run the checkpoint round trip through MyDerivedModel.load_from_checkpoint; done means the reproduced checkpoint loads without the missing hparams argument error and the intended hyperparameter name is retained.