Lightning-AI/pytorch-lightning · error · OSError

Multihost setups do not have a shared filesystem, so the che

Error message

Multihost setups do not have a shared filesystem, so the checkpoint shards cannot be consolidated into a single checkpoint after saving them. Please switch to `XLAFSDPStrategy(state_dict_type='sharded')`. TIP: You can consolidate them manually by getting them together into a single directory and running `python -m torch_xla.distributed.fsdp.consolidate_sharded_ckpts --ckpt_prefix {ckpt_prefix!r} --ckpt_suffix {ckpt_suffix!r} --save_path 'path/to/consolidated.ckpt'`.

What it means

Raised when saving a 'full' (consolidated) XLA FSDP checkpoint on a multihost TPU setup where len(parallel_devices) != world_size. Full checkpoints require gathering and consolidating shards on one host, which is impossible without a shared filesystem across hosts. The error suggests switching to sharded state dicts and consolidating manually afterwards.

Source

Thrown at src/lightning/fabric/strategies/xla_fsdp.py:463

        # ensure model parameters are updated
        xm.mark_step()

        parallel_devices = self.parallel_devices
        assert parallel_devices is not None
        if self._sequential_save:
            # each host runs this in parallel, but the ranks in the host run it sequentially
            for rank in range(len(parallel_devices)):
                if rank == self.local_rank:
                    self._save_checkpoint_shard(path, state, storage_options, filter)
                self.barrier(f"wait-for-{rank}-save")
        else:
            self._save_checkpoint_shard(path, state, storage_options, filter)

        if self._state_dict_type == "full":
            ckpt_prefix = str(path / "checkpoint")
            ckpt_suffix = "_rank-*-of-*.pth"
            if len(parallel_devices) != self.world_size:  # multihost
                raise OSError(
                    "Multihost setups do not have a shared filesystem, so the checkpoint shards cannot be consolidated"
                    " into a single checkpoint after saving them. Please switch to"
                    " `XLAFSDPStrategy(state_dict_type='sharded')`. TIP: You can consolidate them manually by getting"
                    " them together into a single directory and running `python -m"
                    f" torch_xla.distributed.fsdp.consolidate_sharded_ckpts --ckpt_prefix {ckpt_prefix!r} --ckpt_suffix"
                    f" {ckpt_suffix!r} --save_path 'path/to/consolidated.ckpt'`."
                )

            from torch_xla.distributed.fsdp import consolidate_sharded_model_checkpoints

            self.barrier("before_ckpt_consolidation")
            if self.is_global_zero:
                save_path = path.parent / "consolidated.ckpt"
                # save consolidated checkpoint separate to the shards
                consolidate_sharded_model_checkpoints(ckpt_prefix, ckpt_suffix, str(save_path))
                # remove the shards directory
                self.checkpoint_io.remove_checkpoint(path)
                # mv the consolidated checkpoint where the user would expect it

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Construct the strategy with sharded state dicts: XLAFSDPStrategy(state_dict_type='sharded'), then save per-rank shards
2. Alternatively save the shards, gather them into one directory, and run: python -m torch_xla.distributed.fsdp.consolidate_sharded_ckpts --ckpt_prefix <dir>/checkpoint --ckpt_suffix '_rank-*-of-*.pth' --save_path consolidated.ckpt
3. If you truly need full checkpoints on one host, restrict training to single-host setups (v8 TPU with all ranks visible to one process)

Example fix

# before
strategy = XLAFSDPStrategy(state_dict_type='full')

# after
strategy = XLAFSDPStrategy(state_dict_type='sharded')

strategy = XLAFSDPStrategy(state_dict_type='sharded' if len(devices) != world_size else 'full')

When it happens

Trigger: Using XLAFSDPStrategy(state_dict_type='full') (the default) with XLAFSDPStrategy.run(save) on a multihost TPU environment (e.g. TPU Pod slices via torchrun/tpu_pod_launch) where each host only sees its local ranks.

Common situations: Scaling single-host TPU training to a TPU Pod; CI or scripts that worked on v8-8 failing on v4-32/pod slices; default state_dict_type left as 'full' when moving to multihost.

Related errors

• Could not find a XLAFSDP model in the provided checkpoint st
• Found multiple XLAFSDP modules in the given state. Saving ch
• Got `XLAFSDPStrategy.load_checkpoint(..., state={state!r})`
• Loading a single module or optimizer object from a checkpoin
• The path {str(file)!r} does not point to valid sharded check

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