Lightning-AI / Lightning-AI/pytorch-lightning
How to integrate `huggingface.PyTorchModelHubMixin.save_pretrained()` with Lightning `Trainer` (checkpointing & loading)
Nobody has claimed this yet.
- 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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