Lightning-AI/pytorch-lightning · critical · ModuleNotFoundError

raise ModuleNotFoundError(str(_XLA_AVAILABLE))

Error message

raise ModuleNotFoundError(str(_XLA_AVAILABLE))

What it means

XLAStrategy requires torch_xla. Its __init__ raises ModuleNotFoundError with Lightning's _XLA_AVAILABLE explanatory message when the package is absent, before setting up the XLA cluster environment and defaults.

Source

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

class XLAStrategy(DDPStrategy):
    """Strategy for training multiple TPU devices using the :func:`torch_xla.distributed.xla_multiprocessing.spawn`
    method."""

    strategy_name = "xla"

    def __init__(
        self,
        accelerator: Optional["pl.accelerators.Accelerator"] = None,
        parallel_devices: Optional[list[torch.device]] = None,
        checkpoint_io: Optional[Union[XLACheckpointIO, _WrappingCheckpointIO]] = None,
        precision_plugin: Optional[XLAPrecision] = None,
        debug: bool = False,
        sync_module_states: bool = True,
        **_: Any,
    ) -> None:
        if not _XLA_AVAILABLE:
            raise ModuleNotFoundError(str(_XLA_AVAILABLE))
        super().__init__(
            accelerator=accelerator,
            parallel_devices=parallel_devices,
            cluster_environment=XLAEnvironment(),
            checkpoint_io=checkpoint_io,
            precision_plugin=precision_plugin,
            start_method="fork",
        )
        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))

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Install torch_xla matching your torch version, or pip install lightning[xla]
  2. Confirm python -c "import torch_xla" works in the same interpreter/environment the script uses
  3. If not targeting TPU, use DDPStrategy/SingleDeviceStrategy

Example fix

# before
strategy = XLAStrategy()  # ModuleNotFoundError

# after
# pip install lightning[xla]
strategy = XLAStrategy()
Defensive patterns

Strategy: validation

Validate before calling

from lightning.fabric.utilities.imports import _XLA_AVAILABLE
assert _XLA_AVAILABLE, "pip install lightning[xla] before using XLAStrategy"

Prevention

When it happens

Trigger: Constructing XLAStrategy(...) in an environment without torch_xla (local GPU/CPU box, CI, or a container missing the XLA wheel).

Common situations: Running TPU-targeted scripts locally for debugging; requirements.txt missing torch_xla; version mismatch after upgrading torch.

Understand the failure class

Background: "X is not installed. Please install it with pip install Y": missing optional dependency errors — ImportError/ValueError raised when a library's optional extra was never installed — this error's family across 22 libraries.

Related errors


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