Implementation:Bentoml BentoML Framework PyTorch Lightning
| Knowledge Sources | |
|---|---|
| Domains | ML Framework, PyTorch Lightning, Deep Learning, Model Serialization |
| Last Updated | 2026-02-13 15:00 GMT |
Overview
Provides the BentoML framework integration for PyTorch Lightning, enabling saving, loading, and serving pl.LightningModule models through the BentoML model store by converting them to TorchScript.
Description
This module implements the BentoML framework adapter for PyTorch Lightning. It works by converting pl.LightningModule instances to torch.ScriptModule via model.to_torchscript() and then delegating to the TorchScript framework module for actual serialization.
save_model() validates the input is a pl.LightningModule, calls to_torchscript() to produce a torch.ScriptModule, and delegates to torchscript.save_model() with _framework_name="pytorch_lightning" and _module_name=MODULE_NAME to properly tag the model. Saving a dict of modules (multi-module export) is explicitly not supported and raises an assertion error.
load_model() retrieves the model from the store and loads it using torch.jit.load() with an optional device_id parameter (defaults to "cpu"), returning a torch.ScriptModule. It uses the same MODEL_FILENAME as the TorchScript module.
get() retrieves the BentoML Model metadata for a given tag, validating it was saved with the PyTorch Lightning module.
get_runnable() creates a PytorchModelRunnable using the shared PyTorch runnable utilities (partial_class, make_pytorch_runnable_method), mapping each model signature to a runnable method.
Usage
Use this module to save PyTorch Lightning models to the BentoML model store and serve them. The models are automatically converted to TorchScript format for optimized inference. Access via bentoml.pytorch_lightning.
Code Reference
Source Location
- Repository: Bentoml_BentoML
- File: src/bentoml/_internal/frameworks/pytorch_lightning.py
- Lines: 1-204
Signature
def get(tag_like: str | Tag) -> Model: ...
def load_model(
bentoml_model: str | Tag | Model,
device_id: t.Optional[str] = "cpu",
) -> torch.ScriptModule: ...
def save_model(
name: Tag | str,
model: pl.LightningModule,
*,
signatures: ModelSignaturesType | None = None,
labels: t.Dict[str, str] | None = None,
custom_objects: t.Dict[str, t.Any] | None = None,
external_modules: t.List[ModuleType] | None = None,
metadata: t.Dict[str, t.Any] | None = None,
) -> bentoml.Model: ...
def get_runnable(bento_model: Model) -> type[PytorchModelRunnable]: ...
Import
import bentoml
# Via the public API
bento_model = bentoml.pytorch_lightning.save_model("lit_classifier", lit_model)
loaded = bentoml.pytorch_lightning.load_model("lit_classifier:latest", device_id="cuda:0")
# Direct import
from bentoml._internal.frameworks.pytorch_lightning import save_model, load_model, get, get_runnable
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| name | Tag or str | Yes (save_model) | Name or tag for the model in the store |
| model | pl.LightningModule | Yes (save_model) | The PyTorch Lightning module to save; converted to TorchScript internally |
| tag_like | str or Tag | Yes (get) | Tag to retrieve from the model store |
| bentoml_model | str, Tag, or Model | Yes (load_model) | Model identifier or object to load |
| device_id | str or None | No (default "cpu") | Device to load the model onto (e.g., "cpu", "cuda:0") |
| signatures | ModelSignaturesType or None | No | Method signatures; defaults inherited from torchscript module |
| labels | dict[str, str] or None | No | User-defined labels for model management |
| custom_objects | dict[str, Any] or None | No | Additional objects to save with the model |
| external_modules | List[ModuleType] or None | No | Additional Python modules to bundle |
| metadata | dict[str, Any] or None | No | Custom metadata for the model |
Outputs
| Name | Type | Description |
|---|---|---|
| bentoml.Model (save_model) | bentoml.Model | The saved model reference in the BentoML store (stored as TorchScript) |
| torch.ScriptModule (load_model) | torch.ScriptModule | The deserialized TorchScript module on the specified device |
| Model (get) | Model | BentoML Model metadata object |
| PytorchModelRunnable (get_runnable) | type[Runnable] | A Runnable class using shared PyTorch utilities |
Usage Examples
import bentoml
import torch
import pytorch_lightning as pl
class LitClassifier(pl.LightningModule):
def __init__(self, hidden_dim: int = 128, learning_rate: float = 0.0001):
super().__init__()
self.save_hyperparameters()
self.l1 = torch.nn.Linear(28 * 28, self.hparams.hidden_dim)
self.l2 = torch.nn.Linear(self.hparams.hidden_dim, 10)
def forward(self, x):
x = x.view(x.size(0), -1)
x = torch.relu(self.l1(x))
x = torch.relu(self.l2(x))
return x
# Save the Lightning module (auto-converted to TorchScript)
tag = bentoml.pytorch_lightning.save_model("lit_classifier", LitClassifier())
print(f"Saved: {tag}")
# Load the model (returns a torch.ScriptModule)
model = bentoml.pytorch_lightning.load_model("lit_classifier:latest", device_id="cuda:0")
# Run inference
input_tensor = torch.randn(1, 28 * 28)
output = model(input_tensor)