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

[yilin404]

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

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

[SkafteNicki]

@yilin404 thanks for raising this issue, I am going to do my best to give you an answer. The best way I could come up with to integrate huggingfaces save_pretrained functionality into lightning is using a custom Callback. In this case the two important methods that we need to implement is on_save_checkpoint and on_load_checkpoint which lightning will automatically call when ModelCheckpoint is activated. This is a solution:

class HuggingFaceIntegrationCallback(Callback):
    def __init__(self, save_dir="hf_checkpoints"):
        super().__init__()
        self.save_dir = save_dir

    def on_save_checkpoint(self, trainer, pl_module, checkpoint):
        save_path = os.path.join(self.save_dir, f"epoch={trainer.current_epoch}-step={trainer.global_step}")
        os.makedirs(save_path, exist_ok=True)

        saved_modules = []
        for name, module in pl_module.named_modules():
            if hasattr(module, "_save_pretrained") and module != pl_module:
                module_save_path = os.path.join(save_path, name.replace(".", "_"))
                os.makedirs(module_save_path, exist_ok=True)
                module._save_pretrained(module_save_path)
                saved_modules.append(name)

        if saved_modules:
            print(f"Saved HuggingFace modules to {save_path}: {saved_modules}")

    def on_load_checkpoint(self, trainer, pl_module, checkpoint):
        if "hf_module_paths" in checkpoint:
            for name, path in checkpoint["hf_module_paths"].items():
                if os.path.exists(path):
                    module = dict(pl_module.named_modules())[name]
                    if hasattr(module.__class__, "_from_pretrained"):
                        loaded_module = module.__class__._from_pretrained(path)
                        module.load_state_dict(loaded_module.state_dict())
                        print(f"Loaded HuggingFace module {name} from {path}")

Essentially, what is does it that whenever one of the methods is called. it will automatically looks for modules and submodules of your model that has the _save_pretrained and _from_pretrained methods/attributes defined and call them appropriately. Here is a full example:

import json
import os

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

from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.callbacks import Callback, ModelCheckpoint


class BasePolicy(nn.Module):
    def __init__(self, input_dim=10, hidden_dim=64, output_dim=10):
        super().__init__()
        self.config = {
            "input_dim": input_dim,
            "hidden_dim": hidden_dim,
            "output_dim": output_dim,
        }
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim),
        )

    def forward(self, x):
        logits = self.net(x)
        loss = torch.nn.functional.mse_loss(logits, torch.zeros_like(logits))
        loss_dict = {"mse_loss": loss.item()}
        return loss, loss_dict

    def _save_pretrained(self, save_directory):
        torch.save(self.state_dict(), f"{save_directory}/model.safetensors")
        with open(f"{save_directory}/config.json", "w") as f:
            json.dump(self.config, f)

    @classmethod
    def _from_pretrained(cls, model_id, **kwargs):
        with open(f"{model_id}/config.json") as f:
            config = json.load(f)
        model = cls(**config)
        model.load_state_dict(torch.load(f"{model_id}/model.safetensors", weights_only=True))
        return model


class DataProcessorPipeline(nn.Module):
    def __init__(self, transform_type="normalize"):
        super().__init__()
        self.config = {"transform_type": transform_type}
        self.transform_type = transform_type

    def forward(self, batch):
        if self.transform_type == "normalize":
            return batch / 2.0
        return batch

    def _save_pretrained(self, save_directory):
        with open(f"{save_directory}/config.json", "w") as f:
            json.dump(self.config, f)

    @classmethod
    def _from_pretrained(cls, model_id, **kwargs):
        with open(f"{model_id}/config.json") as f:
            config = json.load(f)
        return cls(**config)


class HuggingFaceIntegrationCallback(Callback):
    def __init__(self, save_dir="hf_checkpoints"):
        super().__init__()
        self.save_dir = save_dir

    def on_save_checkpoint(self, trainer, pl_module, checkpoint):
        save_path = os.path.join(self.save_dir, f"epoch={trainer.current_epoch}-step={trainer.global_step}")
        os.makedirs(save_path, exist_ok=True)

        saved_modules = []
        for name, module in pl_module.named_modules():
            if hasattr(module, "_save_pretrained") and module != pl_module:
                module_save_path = os.path.join(save_path, name.replace(".", "_"))
                os.makedirs(module_save_path, exist_ok=True)
                module._save_pretrained(module_save_path)
                saved_modules.append(name)

        if saved_modules:
            print(f"Saved HuggingFace modules to {save_path}: {saved_modules}")

    def on_load_checkpoint(self, trainer, pl_module, checkpoint):
        if "hf_module_paths" in checkpoint:
            for name, path in checkpoint["hf_module_paths"].items():
                if os.path.exists(path):
                    module = dict(pl_module.named_modules())[name]
                    if hasattr(module.__class__, "_from_pretrained"):
                        loaded_module = module.__class__._from_pretrained(path)
                        module.load_state_dict(loaded_module.state_dict())
                        print(f"Loaded HuggingFace module {name} from {path}")


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 = batch[0]
        batch = self.preprocessor(batch)
        loss, loss_dict = self.policy(batch)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)


def main():
    policy = BasePolicy(input_dim=10, hidden_dim=64, output_dim=10)
    preprocessor = DataProcessorPipeline(transform_type="normalize")
    postprocessor = DataProcessorPipeline(transform_type="normalize")

    model = PolicyModule(policy=policy, preprocessor=preprocessor, postprocessor=postprocessor)

    x = torch.randn(100, 10)
    dataset = TensorDataset(x)
    dataloader = DataLoader(dataset, batch_size=10)

    checkpoint_callback = ModelCheckpoint(
        dirpath="checkpoints",
        filename="epoch={epoch}-step={step}",
        save_top_k=3,
        monitor="train/_loss",
    )

    hf_callback = HuggingFaceIntegrationCallback(save_dir="hf_checkpoints")
    trainer = Trainer(max_epochs=2, callbacks=[checkpoint_callback, hf_callback])
    trainer.fit(model, dataloader)


if __name__ == "__main__":
    main()

which when run should produce a checkpoiint structure like this:

hf_checkpoints/
    epoch=0-step=10/
        policy/
            config.json
            model.safetensor
        postprocessor/
            config.json
            model.safetensor
        preprocessor/
            config.json
            model.safetensor
    epoch=1-step=20/
...

The exact logic can be changed of how files are organized, but hope this give you an idea on how to integrate the two frameworks.