Lightning-AI/pytorch-lightning · error · TypeError

`Trainer.save_checkpoint(..., storage_options=...)` with `st

Error message

`Trainer.save_checkpoint(..., storage_options=...)` with `storage_options` arg is not supported for `{self.__class__.__name__}`. Please implement your custom `CheckpointIO` to define how you'd like to use `storage_options`.

What it means

XLACheckpointIO.save_checkpoint raises TypeError when a storage_options argument is passed, because the XLA checkpoint implementation has no way to forward arbitrary storage options to the xm.save/xla filesystem layer.

Source

Thrown at src/lightning/fabric/plugins/io/xla.py:58

            raise ModuleNotFoundError(str(_XLA_AVAILABLE))
        super().__init__(*args, **kwargs)

    @override
    def save_checkpoint(self, checkpoint: dict[str, Any], path: _PATH, storage_options: Optional[Any] = None) -> None:
        """Save model/training states as a checkpoint file through state-dump and file-write.

        Args:
            checkpoint: dict containing model and trainer state
            path: write-target path
            storage_options: not used in ``XLACheckpointIO.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`."
            )
        fs = get_filesystem(path)
        fs.makedirs(os.path.dirname(path), exist_ok=True)
        if RequirementCache("omegaconf"):
            # workaround for https://github.com/pytorch/xla/issues/2773
            from omegaconf import DictConfig, ListConfig, OmegaConf

            checkpoint = apply_to_collection(checkpoint, (DictConfig, ListConfig), OmegaConf.to_container)
        import torch_xla.core.xla_model as xm

        cpu_data = xm._maybe_convert_to_cpu(checkpoint, convert=True)
        log.debug(f"Saving checkpoint: {path}")
        torch.save(cpu_data, path)

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Remove the storage_options argument when saving checkpoints on XLA
2. If you need storage_options, subclass XLACheckpointIO (or implement CheckpointIO) and handle storage_options yourself
3. Save locally then upload to remote storage manually after xm.save completes

Example fix

# before
fabric.save_checkpoint(path, storage_options={"rkwargs": {"aws_access_key_id": ...}})

# after
fabric.save_checkpoint(path)

from lightning.fabric.plugins.io.xla import XLACheckpointIO
from lightning.fabric.plugins.io import CheckpointIO

def supports_storage_options(io) -> bool:
    import inspect
    sig = inspect.signature(io.save_checkpoint)
    return "storage_options" in sig.parameters and type(io) is not XLACheckpointIO

def is_xla_checkpoint_io(io) -> bool:
    "TypeGuard[XLACheckpointIO]"
    from lightning.fabric.plugins.io.xla import XLACheckpointIO
    return isinstance(io, XLACheckpointIO)

try:
    fabric.save_checkpoint(path, storage_options=opts)
except TypeError as e:
    if "storage_options" in str(e):
        fabric.save_checkpoint(path)  # fallback: save without options
    else:
        raise

When it happens

Trigger: Calling fabric.save_checkpoint(...) / trainer.save_checkpoint(...) with storage_options=... while the XLACheckpointIO plugin is active, or calling XLACheckpointIO.save_checkpoint(path, storage_options={...}) directly.

Common situations: Copying storage_options usage (e.g. for fsspec/S3 options) from a TorchCheckpointIO setup onto a TPU/XLA run; passing remote-fs credentials options that work with other plugins.

Related errors

• str(_XLA_AVAILABLE)
• str(_XLA_AVAILABLE)
• `precision={precision!r})` is not supported in XLA. `precisi
• `FSDPStrategy.save_checkpoint(..., storage_options=...)` is
• Found multiple FSDP models in the given state. Saving checkp

AI-assisted analysis of Lightning-AI/pytorch-lightning@9fed5c27d2 (2026-08-28). Data as JSON: /api/errors/d87fa5cbd6a67365. Report an issue: GitHub.