Feature request: Support splitting model weights and training states into separate checkpoint files

[yilin404]

Description & Motivation

🚀 Feature Request: Support splitting model weights and training states into separate checkpoint files

Feature Request

Currently, PyTorch Lightning saves the entire training state (model weights, optimizer states, scheduler states, trainer state, etc.) into a single .ckpt file.

I would like to have an option to separate model weights (and config) from training states when saving checkpoints.

For example, the desired checkpoint structure could look like this:

checkpoints/
pretrained_model/
config.json # model configuration
model.safetensors # model weights only
training_states.pth # optimizer, LR scheduler, trainer states

Pitch

No response

Alternatives

No response

Additional context

No response

cc @lantiga @Borda

[SkafteNicki]

Hi @yilin404 thanks for proposing this new feature. Could you provide some more context for why this feature would be beneficial?

Regardless of this should be a build in feature of lightning or not, you can achieve it pretty simply by using the plugin system in lightning:
https://lightning.ai/docs/pytorch/stable/extensions/plugins.html

I written a SplitCheckpointIO plugin (expand the details part below) that can simply be used together with the trainer as:

    split_checkpoint_io = SplitCheckpointIO(use_safetensors=True)
    trainer = pl.Trainer(
        plugins=[split_checkpoint_io],  # 🎯 This is the key addition!
        ...
    )

Implementation of `SplitCheckpointIO`

# Copyright The Lightning AI team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import logging
import os
from pathlib import Path
from typing import Any, Callable, Optional

from typing_extensions import override

from lightning.fabric.plugins.io.torch_io import TorchCheckpointIO
from lightning.fabric.utilities.cloud_io import _atomic_save, get_filesystem
from lightning.fabric.utilities.cloud_io import _load as pl_load
from lightning.fabric.utilities.types import _PATH

log = logging.getLogger(__name__)


class SplitCheckpointIO(TorchCheckpointIO):
    """CheckpointIO that saves model weights and training states in separate files.
    
    Model weights are saved in safetensors format with a config.json file,
    while training states (optimizer, scheduler, etc.) are saved in a separate .pth file.
    
    Checkpoint structure:
    ```
    checkpoint_name/
    ├── config.json           # Model configuration and metadata
    ├── model.safetensors     # Model weights only  
    └── training_states.pth   # Optimizer, scheduler, trainer states
    ```
    """

    def __init__(self, use_safetensors: bool = True) -> None:
        """Initialize SplitCheckpointIO.
        
        Args:
            use_safetensors: Whether to use safetensors format for model weights.
                If False, will use PyTorch format (.pth).
        """
        super().__init__()
        self.use_safetensors = use_safetensors
        
        if use_safetensors:
            try:
                import safetensors.torch  # noqa: F401
            except ImportError:
                log.warning(
                    "safetensors is not installed, falling back to PyTorch format. "
                    "Install safetensors with: pip install safetensors"
                )
                self.use_safetensors = False

    @override
    def save_checkpoint(self, checkpoint: dict[str, Any], path: _PATH, storage_options: Optional[Any] = None) -> None:
        """Save model weights and training states in separate files.

        Args:
            checkpoint: dict containing model and trainer state
            path: write-target path (will be used as directory name)
            storage_options: not used in ``SplitCheckpointIO.save_checkpoint``

        Raises:
            TypeError: If ``storage_options`` arg is passed in
        """
        if storage_options is not None:
            raise TypeError(
                "`Trainer.save_checkpoint(..., storage_options=...)` with `storage_options` arg"
                f" is not supported for `{self.__class__.__name__}`. Please implement your custom `CheckpointIO`"
                " to define how you'd like to use `storage_options`."
            )

        # Convert path to directory
        checkpoint_dir = Path(str(path))
        if checkpoint_dir.suffix == '.ckpt':
            # Remove .ckpt extension to use as directory name
            checkpoint_dir = checkpoint_dir.with_suffix('')
        
        fs = get_filesystem(checkpoint_dir)
        fs.makedirs(checkpoint_dir, exist_ok=True)

        # Separate model weights from training states
        model_data = {}
        training_data = {}
        config_data = {}

        for key, value in checkpoint.items():
            if key == "state_dict":
                model_data[key] = value
            elif key in ["epoch", "global_step", "pytorch-lightning_version"]:
                # These go in config for easy access
                config_data[key] = value
            elif key in ["hyper_parameters", "CHECKPOINT_HYPER_PARAMS_NAME", 
                        "CHECKPOINT_HYPER_PARAMS_KEY", "CHECKPOINT_HYPER_PARAMS_TYPE"]:
                # Hyperparameters go in config
                config_data[key] = value
            else:
                # Everything else (optimizer, scheduler, callbacks, etc.) goes in training states
                training_data[key] = value

        # Save model weights
        model_path = checkpoint_dir / "model.safetensors" if self.use_safetensors else checkpoint_dir / "model.pth"
        if self.use_safetensors and "state_dict" in model_data:
            import safetensors.torch
            safetensors.torch.save_file(model_data["state_dict"], model_path)
        else:
            _atomic_save(model_data, model_path)

        # Save config as JSON
        config_path = checkpoint_dir / "config.json"
        with open(config_path, 'w') as f:
            # Convert non-serializable values to strings
            serializable_config = {}
            for key, value in config_data.items():
                try:
                    json.dumps(value)  # Test if serializable
                    serializable_config[key] = value
                except (TypeError, ValueError):
                    serializable_config[key] = str(value)
            json.dump(serializable_config, f, indent=2, default=str)

        # Save training states
        training_states_path = checkpoint_dir / "training_states.pth"
        if training_data:  # Only save if there's training data
            _atomic_save(training_data, training_states_path)

        log.debug(f"Saved split checkpoint to: {checkpoint_dir}")

    @override
    def load_checkpoint(
        self, path: _PATH, map_location: Optional[Callable] = lambda storage, loc: storage
    ) -> dict[str, Any]:
        """Load checkpoint from separate files.

        Args:
            path: Path to checkpoint directory or .ckpt file
            map_location: a function, :class:`torch.device`, string or a dict specifying how to remap storage
                locations.

        Returns: The loaded checkpoint as a single dictionary.

        Raises:
            FileNotFoundError: If checkpoint files are not found
        """
        # Handle both directory and .ckpt file paths
        checkpoint_dir = Path(str(path))
        if checkpoint_dir.suffix == '.ckpt':
            checkpoint_dir = checkpoint_dir.with_suffix('')

        fs = get_filesystem(checkpoint_dir)
        if not fs.exists(checkpoint_dir):
            raise FileNotFoundError(f"Checkpoint directory not found: {checkpoint_dir}")

        checkpoint = {}

        # Load config
        config_path = checkpoint_dir / "config.json"
        if fs.exists(config_path):
            with open(config_path, 'r') as f:
                config_data = json.load(f)
                checkpoint.update(config_data)

        # Load model weights
        model_safetensors_path = checkpoint_dir / "model.safetensors"
        model_pth_path = checkpoint_dir / "model.pth"
        
        if fs.exists(model_safetensors_path) and self.use_safetensors:
            import safetensors.torch
            # safetensors.load_file expects a device string, not a callable
            # We'll load to CPU first, then apply map_location if needed
            state_dict = safetensors.torch.load_file(model_safetensors_path, device="cpu")
            if map_location and callable(map_location):
                # Apply map_location to each tensor if it's a callable
                state_dict = {k: map_location(v, None) for k, v in state_dict.items()}
            elif map_location and not callable(map_location):
                # If map_location is a device, move tensors there
                import torch
                device = torch.device(map_location) if isinstance(map_location, str) else map_location
                state_dict = {k: v.to(device) for k, v in state_dict.items()}
            checkpoint["state_dict"] = state_dict
        elif fs.exists(model_pth_path):
            model_data = pl_load(model_pth_path, map_location=map_location)
            checkpoint.update(model_data)
        else:
            raise FileNotFoundError(f"Model weights file not found in: {checkpoint_dir}")

        # Load training states
        training_states_path = checkpoint_dir / "training_states.pth"
        if fs.exists(training_states_path):
            training_data = pl_load(training_states_path, map_location=map_location)
            checkpoint.update(training_data)

        return checkpoint

    @override
    def remove_checkpoint(self, path: _PATH) -> None:
        """Remove checkpoint directory and all its contents.

        Args:
            path: Path to checkpoint directory or .ckpt file
        """
        # Handle both directory and .ckpt file paths
        checkpoint_dir = Path(str(path))
        if checkpoint_dir.suffix == '.ckpt':
            checkpoint_dir = checkpoint_dir.with_suffix('')

        fs = get_filesystem(checkpoint_dir)
        if fs.exists(checkpoint_dir):
            fs.rm(checkpoint_dir, recursive=True)
            log.debug(f"Removed checkpoint directory: {checkpoint_dir}")

which creates a checkpoint structure like this:

checkpoint_name/
    ├── config.json           # Model configuration and metadata
    ├── model.safetensors     # Model weights only  
    └── training_states.pth   # Optimizer, scheduler, trainer states

a full example on mnist

MNIST example using `SplitCheckpointIO`

import os
import json
from pathlib import Path
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, random_split
from torchvision import transforms
from torchvision.datasets import MNIST

import lightning.pytorch as pl
from lightning.pytorch.callbacks import ModelCheckpoint, EarlyStopping
from lightning.pytorch.loggers import CSVLogger

class MNISTModel(pl.LightningModule):
    """Simple CNN model for MNIST classification."""
    
    def __init__(
        self,
        hidden_size: int = 128,
        learning_rate: float = 1e-3,
        dropout: float = 0.1,
    ):
        super().__init__()
        
        # Save hyperparameters
        self.save_hyperparameters()
        
        # Network layers
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.fc1 = nn.Linear(9216, hidden_size)
        self.fc2 = nn.Linear(hidden_size, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        
        # Log metrics
        self.log('train_loss', loss, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        preds = torch.argmax(logits, dim=1)
        acc = torch.sum(preds == y).float() / len(y)
        
        # Log metrics
        self.log('val_loss', loss, prog_bar=True)
        self.log('val_acc', acc, prog_bar=True)
        return loss

    def test_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        preds = torch.argmax(logits, dim=1)
        acc = torch.sum(preds == y).float() / len(y)
        
        # Log metrics
        self.log('test_loss', loss)
        self.log('test_acc', acc)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)
        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "interval": "epoch",
            },
        }


class MNISTDataModule(pl.LightningDataModule):
    """MNIST DataModule for PyTorch Lightning."""
    
    def __init__(self, data_dir: str = "./data", batch_size: int = 32, num_workers: int = 4):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.num_workers = num_workers
        
        # Define transforms
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])

    def prepare_data(self):
        # Download MNIST dataset
        MNIST(self.data_dir, train=True, download=True)
        MNIST(self.data_dir, train=False, download=True)

    def setup(self, stage: Optional[str] = None):
        # Assign train/val datasets for use in dataloaders
        if stage == "fit" or stage is None:
            mnist_full = MNIST(self.data_dir, train=True, transform=self.transform)
            self.mnist_train, self.mnist_val = random_split(mnist_full, [55000, 5000])

        # Assign test dataset for use in dataloader(s)
        if stage == "test" or stage is None:
            self.mnist_test = MNIST(self.data_dir, train=False, transform=self.transform)

    def train_dataloader(self):
        return DataLoader(
            self.mnist_train, 
            batch_size=self.batch_size, 
            num_workers=self.num_workers,
            shuffle=True
        )

    def val_dataloader(self):
        return DataLoader(
            self.mnist_val, 
            batch_size=self.batch_size, 
            num_workers=self.num_workers
        )

    def test_dataloader(self):
        return DataLoader(
            self.mnist_test, 
            batch_size=self.batch_size, 
            num_workers=self.num_workers
        )


def main():    
    # Set up directories
    checkpoint_dir = Path("./mnist_checkpoints")
    log_dir = Path("./mnist_logs")
    
    checkpoint_dir.mkdir(exist_ok=True)
    log_dir.mkdir(exist_ok=True)
    
    print(f"📁 Checkpoints will be saved to: {checkpoint_dir}")
    print(f"📊 Logs will be saved to: {log_dir}")
    
    # Create model and data module
    model = MNISTModel(hidden_size=128, learning_rate=1e-3, dropout=0.1)
    data_module = MNISTDataModule(batch_size=64, num_workers=0)  # num_workers=0 for compatibility
    
    # Create SplitCheckpointIO plugin
    split_checkpoint_io = SplitCheckpointIO(use_safetensors=True)
    print("🔧 Created SplitCheckpointIO plugin with safetensors=True")
    
    # Set up callbacks
    checkpoint_callback = ModelCheckpoint(
        dirpath=checkpoint_dir,
        filename="mnist-{epoch:02d}-{val_loss:.2f}",
        monitor="val_loss",
        mode="min",
        save_top_k=2,
        save_last=True,
    )
    
    early_stopping = EarlyStopping(
        monitor="val_loss",
        min_delta=0.00,
        patience=3,
        verbose=False,
        mode="min"
    )
    
    # Create logger
    logger = CSVLogger(log_dir, name="mnist_experiment")
    
    # Create trainer with SplitCheckpointIO plugin
    trainer = pl.Trainer(
        plugins=[split_checkpoint_io],  # 🎯 This is the key addition!
        callbacks=[checkpoint_callback, early_stopping],
        logger=logger,
        max_epochs=5,  # Short training for demo
        accelerator="auto",
        devices="auto",
        log_every_n_steps=50,
        enable_progress_bar=True,
    )

    # Train the model
    trainer.fit(model, data_module)
    
    # Test the model
    trainer.test(model, data_module)



if __name__ == "__main__":
    main()