Lightning-AI/pytorch-lightning · error · RuntimeError

Accessing the device mesh before processes have initialized

Error message

Accessing the device mesh before processes have initialized is not allowed.

What it means

ModelParallelStrategy exposes a device_mesh property that returns the torch.distributed.device_mesh.DeviceMesh created during setup. Accessing it before the strategy has initialized (mesh not yet built) raises RuntimeError, because the mesh depends on process group initialization that only happens at setup time.

Source

Thrown at src/lightning/pytorch/strategies/model_parallel.py:106

        data_parallel_size: Union[Literal["auto"], int] = "auto",
        tensor_parallel_size: Union[Literal["auto"], int] = "auto",
        save_distributed_checkpoint: bool = True,
        process_group_backend: Optional[str] = None,
        timeout: Optional[timedelta] = default_pg_timeout,
    ) -> None:
        super().__init__()
        self._data_parallel_size = data_parallel_size
        self._tensor_parallel_size = tensor_parallel_size
        self._save_distributed_checkpoint = save_distributed_checkpoint
        self._process_group_backend: Optional[str] = process_group_backend
        self._timeout: Optional[timedelta] = timeout
        self._device_mesh: Optional[DeviceMesh] = None
        self.num_nodes = 1

    @property
    def device_mesh(self) -> "DeviceMesh":
        if self._device_mesh is None:
            raise RuntimeError("Accessing the device mesh before processes have initialized is not allowed.")
        return self._device_mesh

    @property
    @override
    def root_device(self) -> torch.device:
        assert self.parallel_devices is not None
        return self.parallel_devices[self.local_rank]

    @property
    def num_processes(self) -> int:
        return len(self.parallel_devices) if self.parallel_devices is not None else 0

    @property
    @override
    def distributed_sampler_kwargs(self) -> dict[str, Any]:
        assert self.device_mesh is not None
        data_parallel_mesh = self.device_mesh["data_parallel"]
        return {"num_replicas": data_parallel_mesh.size(), "rank": data_parallel_mesh.get_local_rank()}

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Move mesh access into/after setup: use setup() hooks, on_fit_start callbacks, or access it inside forward/training_step where setup already ran
  2. Create your own DeviceMesh and pass parallelize_to_plan / use the strategy's API for pre-setup mesh needs
  3. Guard with an attribute check: getattr(strategy, '_device_mesh', None) before use

Example fix

# before
class LitModel(L.LightningModule):
    def __init__(self):
        mesh = self.trainer.strategy.device_mesh  # RuntimeError

# after
class LitModel(L.LightningModule):
    def setup(self, stage=None):
        mesh = self.trainer.strategy.device_mesh  # OK: after strategy setup
Defensive patterns

Strategy: validation

Validate before calling

mesh = getattr(strategy, "_device_mesh", None)
if mesh is None:
    # too early: defer access to setup()/hooks
    ...

Prevention

When it happens

Trigger: Reading strategy.device_mesh in LightningModule.__init__, configure_model (before setup), or any code executed before trainer setup (e.g. in module-level code or a pre-run callback) with ModelParallelStrategy.

Common situations: Trying to inspect or pass the mesh to torch tensor creation in __init__ or configure_model; logging mesh info before training starts.

Related errors


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