Lightning-AI/pytorch-lightning · error · TypeError

The XLA strategy can only work with the `XLACheckpointIO` pl

Error message

The XLA strategy can only work with the `XLACheckpointIO` plugin, found {io}

What it means

Identical validation to single_xla: XLAStrategy's checkpoint_io setter only accepts XLACheckpointIO or _WrappingCheckpointIO; any other CheckpointIO is rejected with TypeError because XLA checkpoint save/load requires xm/xla-specific handling.

Source

Thrown at src/lightning/pytorch/strategies/xla.py:86

        )
        self.debug = debug
        self._launched = False
        self._sync_module_states = sync_module_states

    @property
    @override
    def checkpoint_io(self) -> Union[XLACheckpointIO, _WrappingCheckpointIO]:
        plugin = self._checkpoint_io
        if plugin is not None:
            assert isinstance(plugin, (XLACheckpointIO, _WrappingCheckpointIO))
            return plugin
        return XLACheckpointIO()

    @checkpoint_io.setter
    @override
    def checkpoint_io(self, io: Optional[CheckpointIO]) -> None:
        if io is not None and not isinstance(io, (XLACheckpointIO, _WrappingCheckpointIO)):
            raise TypeError(f"The XLA strategy can only work with the `XLACheckpointIO` plugin, found {io}")
        self._checkpoint_io = io

    @property
    @override
    def precision_plugin(self) -> XLAPrecision:
        plugin = self._precision_plugin
        if plugin is not None:
            assert isinstance(plugin, XLAPrecision)
            return plugin
        return XLAPrecision()

    @precision_plugin.setter
    @override
    def precision_plugin(self, precision_plugin: Optional[Precision]) -> None:
        if precision_plugin is not None and not isinstance(precision_plugin, XLAPrecision):
            raise TypeError(f"The XLA strategy can only work with the `XLAPrecision` plugin, found {precision_plugin}")
        self._precision_plugin = precision_plugin

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Use XLACheckpointIO or _WrappingCheckpointIO(your_io)
2. Subclass XLACheckpointIO for custom behavior
3. Pass storage credentials via environment instead of a custom IO

Example fix

# before
strategy = XLAStrategy(checkpoint_io=MyS3CheckpointIO())

# after
from lightning.pytorch.plugins.io.xla import XLACheckpointIO
from lightning.pytorch.strategies.xla import _WrappingCheckpointIO
strategy = XLAStrategy(checkpoint_io=_WrappingCheckpointIO(MyS3CheckpointIO()))

assert io is None or isinstance(io, (XLACheckpointIO, _WrappingCheckpointIO)), "wrap custom IO with _WrappingCheckpointIO"

def is_valid_xla_checkpoint_io(io) -> bool:
    return io is None or type(io).__name__ in {"XLACheckpointIO", "_WrappingCheckpointIO"} or isinstance(io, (XLACheckpointIO, _WrappingCheckpointIO))

When it happens

Trigger: strategy.checkpoint_io = <custom or torch CheckpointIO> or XLAStrategy(checkpoint_io=...) with an incompatible plugin on the multi-device XLA strategy.

Common situations: Reusing checkpoint plugin configuration from DDP/DeepSpeed setups; S3/GCS checkpointing plugins not wrapped for XLA.

Related errors

• The XLA strategy can only work with the `XLACheckpointIO` pl
• The XLA strategy can only work with the `XLAPrecision` plugi
• The XLA strategy can only work with the `XLAPrecision` plugi
• To spawn processes with the `{type(self.strategy).__name__}`
• The `{type(self._strategy).__name__}` requires the model and

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