Lightning-AI/pytorch-lightning · error · RuntimeError

Accessing the XLA device before processes have spawned is no

Error message

Accessing the XLA device before processes have spawned is not allowed.

What it means

XLAStrategy.root_device returns xm.xla_device(), but the XLA runtime is only initialized after the launcher has spawned processes. Accessing root_device before launch (self._launched is False) raises RuntimeError to prevent initializing XLA outside the multiprocess context.

Source

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

    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
    def local_rank(self) -> int:
        return super().local_rank if self._launched else 0

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

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Access root_device only after launch: inside setup(), on_fit_start, or training_step hooks
2. For device-needing setup in LightningModule, use configure_model()/setup() which run post-launch
3. Guard with strategy._launched if you must probe early

Example fix

# before
class Lit(L.LightningModule):
    def __init__(self):
        self.x = torch.zeros(3, device=self.trainer.strategy.root_device)  # RuntimeError

# after
class Lit(L.LightningModule):
    def setup(self, stage=None):
        self.x = torch.zeros(3, device=self.trainer.strategy.root_device)  # after launch

if not getattr(strategy, "_launched", False):
    # defer device access to post-launch hooks
    ...

When it happens

Trigger: Reading strategy.root_device (directly or via trainer logic) before trainer fit/launch — e.g. in LightningModule.__init__, configure_model, or module-level code — while using XLAStrategy.

Common situations: Moving device setup out of hooks; logging devices pre-run; utility code that queries strategy.root_device at import time.

Related errors

• To spawn processes with the `{type(self.strategy).__name__}`
• The `{type(self._strategy).__name__}` requires the model and
• str(_XLA_AVAILABLE)
• `Trainer.save_checkpoint(..., storage_options=...)` with `st
• str(_XLA_AVAILABLE)

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