Lightning-AI / Lightning-AI/pytorch-lightning

How to integrate `huggingface.PyTorchModelHubMixin.save_pretrained()` with Lightning `Trainer` (checkpointing & loading)

Open
#21,366 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

3rd party checkpointing question
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Description & Motivation

## 📝 Summary

I am trying to train a custom policy module using **PyTorch Lightning**, where my model components (`policy`, `preprocessor`, `postprocessor`) all inherit from `huggingface_hub.PyTorchModelHubMixin`.

This mixin provides:
- `_save_pretrained()`
- `_from_pretrained()`
which work similarly to HuggingFace’s `save_pretrained()` and `from_pretrained()`, and are very convenient for packaging model weights + config.

However, I am not sure how to properly integrate these HF-style save/load utilities into Lightning's standard training flow — especially Lightning's `ModelCheckpoint` callback.

## 📦 Minimal example

```python
class PolicyModule(LightningModule):
def __init__(
self,
policy: BasePolicy,
preprocessor: DataProcessorPipeline,
postprocessor: DataProcessorPipeline,
**kwargs,
):
super().__init__()

self.save_hyperparameters(
logger=False, ignore=["policy", "preprocessor", "postprocessor"]
)

self.policy = policy
self.preprocessor = preprocessor
self.postprocessor = postprocessor

def training_step(self, batch, batch_idx):
batch = self.preprocessor(batch)
loss, loss_dict = self.policy(batch)

self.log_dict(
{"train/_loss": loss},
on_step=True, on_epoch=True, prog_bar=True, sync_dist=True,
)
self.log_dict(
{f"train/{k}": v for k, v in loss_dict.items()},
on_step=True, on_epoch=True, sync_dist=True,
)
return loss
```

## ❗The problem
### How to make Lightning call `save_pretrained()` during checkpointing?

## 🙏 Additional context
PyTorch Lightning provides many excellent callback tools—especially `ModelCheckpoint`—which greatly simplify training workflows. Ideally, I would like to remain fully within the standard Lightning Trainer + callbacks framework, while still taking advantage of HuggingFace-style `save_pretrained()` / `from_pretrained()` for model components.

Thanks in advance for any guidance or best practices!

### Pitch

_No response_

### Alternatives

_No response_

### Additional context

_No response_

cc @lantiga

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reviewing Lightning's Trainer and ModelCheckpoint callback flow alongside the PyTorchModelHubMixin save_pretrained() and from_pretrained() methods described in the issue. Determine whether the expected outcome is a documented integration pattern or a framework change, and define done as a clear checkpointing and loading workflow for the listed model components.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
20/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.