Lightning-AI/pytorch-lightning · error · TypeError

The XLA strategy can only work with the `XLAPrecision` plugi

Error message

The XLA strategy can only work with the `XLAPrecision` plugin, found {precision_plugin}

What it means

XLAStrategy's precision_plugin setter only accepts XLAPrecision, since XLA mixed precision and optimizer handling are TPU-specific. Assigning any other Precision plugin raises TypeError.

Source

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

    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

    @property
    @override
    def root_device(self) -> torch.device:
        if not self._launched:
            raise RuntimeError("Accessing the XLA device before processes have spawned is not allowed.")
        import torch_xla.core.xla_model as xm

        return xm.xla_device()

    @property
    @override
    def global_rank(self) -> int:
        return super().global_rank if self._launched else 0

    @property
    @override

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Let the strategy default to XLAPrecision; express precision via Trainer(precision=...)
2. Subclass XLAPrecision if you need custom behavior
3. Drop precision_plugin from shared config files for TPU runs

Example fix

# before
strategy = XLAStrategy(precision_plugin=MixedPrecision())

# after
strategy = XLAStrategy()
trainer = L.Trainer(strategy=strategy, precision="bf16-true")

assert precision_plugin is None or isinstance(precision_plugin, XLAPrecision), "XLA strategies require XLAPrecision"

def is_valid_xla_precision(p) -> bool:
    return p is None or type(p).__name__ == "XLAPrecision"

When it happens

Trigger: XLAStrategy(precision_plugin=MixedPrecision(...)) or strategy.precision_plugin = Precision() with the multi-device XLA strategy.

Common situations: Porting GPU training configs (MixedPrecision with native amp) to TPU; plugin lists in YAML/Sweeps that include precision plugins.

Related errors

• The XLA strategy can only work with the `XLAPrecision` plugi
• The XLA strategy can only work with the `XLACheckpointIO` pl
• The XLA strategy can only work with the `XLACheckpointIO` pl
• 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/362dedfe1897417a. Report an issue: GitHub.