Repository: mlflow/mlflow

Issue:
[BUG] mlflow.pytorch.load_checkpoint() fails when using LightningModule class with arguments.
### Issues Policy acknowledgement

- [x] I have read and agree to submit bug reports in accordance with the [issues policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md)

### Where did you encounter this bug?

Local machine

### Willingness to contribute

Yes. I can contribute a fix for this bug independently.

### MLflow version

- Client: 2.11.3
- Tracking server: 2.11.3


### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Ubuntu 18.04, MacOS
- **Python version**: 3.9.18
- **yarn version, if running the dev UI**:


### Describe the problem

When using mlflow.pytorch.autolog(), it uploads checkpoint automatically in artifacts.
And when the checkpoint is loaded, it needs to be with pl class parameters if it has arguments.
But mlflow official docs and source code(below) is called with only checkpoint file path.
It occur errors, so it needs **kwargs such as pytorch-lightning official docs.

- mlflow.pytorch.load_checkpoint() source code 
```python
...
with TempDir() as tmp_dir:
    downloaded_checkpoint_filepath = download_checkpoint_artifact(
        run_id=run_id, epoch=epoch, global_step=global_step, dst_path=tmp_dir.path()
    )
    return model_class.load_from_checkpoint(downloaded_checkpoint_filepath)
```

- pytorch-lightning load_from_checkpoint()    
**kwargs (`Any`) – Any extra keyword args needed to init the model.


[lightning-docs (load_from_checkpoint)](https://lightning.ai/docs/pytorch/stable/_modules/lightning/pytorch/core/module.html#LightningModule.load_from_checkpoint)
[mlflow-docs (mlflow.pytorch.load_model)](https://mlflow.org/docs/latest/python_api/mlflow.pytorch.html#mlflow.pytorch.load_model)

### Tracking information

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```shell
REPLACE_ME
```


### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```python
import os

import mlflow
import lightning as L

import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST

# set mlflow
mlflow.set_tracking_uri(os.environ['MLFLOW_TRACKING_URI'])
mlflow.set_experiment('test')


class CustomModel(L.LightningModule):
    def __init__(self, dim_in, dim_out):
        super().__init__()
        self.layer = nn.Linear(dim_in, dim_out)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        return

    def validation_step(self, batch, batch_idx):
        self.log('val_loss', 0.0, on_epoch=True)
        return

    def configure_optimizers(self):
        return


# Create a trainer
model = CustomModel(10, 1)
trainer = L.Trainer(max_epochs=1, accelerator='cpu',
                    limit_train_batches=0.1, limit_val_batches=0.1)

# Set dataloader
dataset = MNIST(os.getcwd(), download=True, transform=transforms.ToTensor())
loader = DataLoader(dataset, shuffle=True, batch_size=32)

# Train the model (autolog will log the checkpoints)
mlflow.pytorch.autolog()
with mlflow.start_run() as run:
    trainer.fit(model, loader, loader)

# Load model from MLFlow Server
model = mlflow.pytorch.load_checkpoint(CustomModel, run.info.run_id)
print(model)
print('Success to load model from MLFlow Server!')


```


### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
Traceback (most recent call last):
  File "/Users/jmjeon/mlflow_bug_reproducible.py", line 50, in <module>
    model = mlflow.pytorch.load_checkpoint(CustomModel, run.info.run_id)
  File "/Users/jmjeon/.pyenv/versions/venv/lib/python3.9/site-packages/mlflow/pytorch/__init__.py", line 1189, in load_checkpoint
    return model_class.load_from_checkpoint(downloaded_checkpoint_filepath)
  File "/Users/jmjeon/.pyenv/versions/venv/lib/python3.9/site-packages/lightning/pytorch/utilities/model_helpers.py", line 125, in wrapper
    return self.method(cls, *args, **kwargs)
  File "/Users/jmjeon/.pyenv/versions/venv/lib/python3.9/site-packages/lightning/pytorch/core/module.py", line 1581, in load_from_checkpoint
    loaded = _load_from_checkpoint(
  File "/Users/jmjeon/.pyenv/versions/venv/lib/python3.9/site-packages/lightning/pytorch/core/saving.py", line 91, in _load_from_checkpoint
    model = _load_state(cls, checkpoint, strict=strict, **kwargs)
  File "/Users/jmjeon/.pyenv/versions/venv/lib/python3.9/site-packages/lightning/pytorch/core/saving.py", line 158, in _load_state
    obj = cls(**_cls_kwargs)
TypeError: __init__() missing 2 required positional arguments: 'dim_in' and 'dim_out'
```


### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
REPLACE_ME
```


### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/deployments`: MLflow Deployments client APIs, server, and third-party Deployments integrations
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [X] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:
@jmjeon94 thanks for report this.


@WeichenXu123 @chenmoneygithub 
Thoughts on adding `kwargs` to `load_checkpoint`?
Agreed, we should add `**kwargs: Any` param. Thanks !

Resolving patch:
diff --git a/mlflow/pytorch/__init__.py b/mlflow/pytorch/__init__.py
--- a/mlflow/pytorch/__init__.py
+++ b/mlflow/pytorch/__init__.py
@@ -1133,7 +1133,7 @@ def print_auto_logged_info(r):
     )
 
 
-def load_checkpoint(model_class, run_id=None, epoch=None, global_step=None):
+def load_checkpoint(model_class, run_id=None, epoch=None, global_step=None, kwargs=None):
     """
     If you enable "checkpoint" in autologging, during pytorch-lightning model
     training execution, checkpointed models are logged as MLflow artifacts.
@@ -1156,6 +1156,7 @@ def load_checkpoint(model_class, run_id=None, epoch=None, global_step=None):
             "checkpoint_save_freq" to "epoch".
         global_step: The global step of the checkpoint to be loaded, if
             you set "checkpoint_save_freq" to an integer.
+        kwargs: Any extra kwargs needed to init the model.
 
     Returns:
         The instance of a pytorch-lightning model restored from the specified checkpoint.
@@ -1186,7 +1187,7 @@ def load_checkpoint(model_class, run_id=None, epoch=None, global_step=None):
         downloaded_checkpoint_filepath = download_checkpoint_artifact(
             run_id=run_id, epoch=epoch, global_step=global_step, dst_path=tmp_dir.path()
         )
-        return model_class.load_from_checkpoint(downloaded_checkpoint_filepath)
+        return model_class.load_from_checkpoint(downloaded_checkpoint_filepath, **(kwargs or {}))
 
 
 __all__ = [


Test patch: