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

The XLA strategies delegate checkpointing to XLACheckpointIO (or a _WrappingCheckpointIO). The checkpoint_io setter validates this: assigning any other CheckpointIO implementation raises TypeError.

Source

Thrown at src/lightning/pytorch/strategies/single_xla.py:72

            checkpoint_io=checkpoint_io,
            precision_plugin=precision_plugin,
        )
        self.debug = debug

    @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 wrap your custom IO with _WrappingCheckpointIO so the XLA handling stays in place
2. Pass parallel_devices/other options instead of replacing checkpoint_io
3. If you need custom logic, subclass XLACheckpointIO

Example fix

# before
strategy = SingleDeviceXLAStrategy(checkpoint_io=TorchCheckpointIO())

# after
from lightning.pytorch.plugins.io.xla import XLACheckpointIO
strategy = SingleDeviceXLAStrategy(checkpoint_io=XLACheckpointIO())

from lightning.pytorch.plugins.io.xla import XLACheckpointIO
assert io is None or isinstance(io, (XLACheckpointIO, _WrappingCheckpointIO))

def is_valid_xla_checkpoint_io(io) -> bool:
    from lightning.pytorch.plugins.io.xla import XLACheckpointIO
    return io is None or isinstance(io, (XLACheckpointIO, _WrappingCheckpointIO))

When it happens

Trigger: Assigning strategy.checkpoint_io = TorchCheckpointIO() (or any custom CheckpointIO that is not XLACheckpointIO/_WrappingCheckpointIO) on SingleDeviceXLAStrategy/XLAStrategy; passing checkpoint_io=... to the strategy constructor with an incompatible plugin.

Common situations: Copy-pasting strategy configs from GPU setups that set a custom or cloud CheckpointIO; swapping plugins at runtime.

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/ff7766da554b1c47. Report an issue: GitHub.