Lightning-AI/pytorch-lightning · warning

When saving the DeepSpeed Stage 3 checkpoint, each worker wi

Error message

When saving the DeepSpeed Stage 3 checkpoint, each worker will save a shard of the checkpoint within a directory. If a single file is required after training, see https://lightning.ai/docs/pytorch/stable/advanced/model_parallel.html#deepspeed-zero-stage-3-single-file for instructions.

What it means

With DeepSpeed ZeRO Stage 3 and multiple devices, saving a checkpoint writes one sharded directory per worker, not a single consolidated .ckpt file. Lightning warns and links docs for consolidating to a single file.

Source

Thrown at src/lightning/pytorch/strategies/deepspeed.py:658

            filepath: write-target file's path
            storage_options: not used for ``DeepSpeedStrategy`` as ``CheckpointIO`` is not used

        Raises:
            TypeError:
                If ``storage_options`` arg is passed in

        """
        # broadcast the filepath from rank 0 to ensure all the states are saved in a common filepath
        filepath = self.broadcast(filepath)

        if storage_options is not None:
            raise TypeError(
                "`Trainer.save_checkpoint(..., storage_options=...)` with `storage_options` arg"
                f" is not supported for `{self.__class__.__name__}` as `CheckpointIO` is not used."
            )

        if self.zero_stage_3 and self._multi_device and self.is_global_zero:
            warning_cache.warn(
                "When saving the DeepSpeed Stage 3 checkpoint, "
                "each worker will save a shard of the checkpoint within a directory. "
                "If a single file is required after training, "
                "see https://lightning.ai/docs/pytorch/stable/advanced/model_parallel.html#"
                "deepspeed-zero-stage-3-single-file for instructions."
            )
        # Use deepspeed's internal checkpointing function to handle partitioned weights across processes
        # dump states as a checkpoint dictionary object
        _exclude_keys = ["state_dict", "optimizer_states"]
        checkpoint = {k: v for k, v in checkpoint.items() if k not in _exclude_keys}
        self.deepspeed_engine.save_checkpoint(
            filepath,
            client_state=checkpoint,
            tag="checkpoint",
            exclude_frozen_parameters=self.exclude_frozen_parameters,
        )

    @override

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Follow the linked docs: load the sharded checkpoint with DeepSpeed's conversion, or use zero_to_fp32.py to consolidate weights into a single file
2. Resume training directly from the sharded directory (supported)
3. If single-file checkpoints are essential, use stage 2 or zero_offload variants without sharded state

Example fix

# before
trainer.save_checkpoint('model.ckpt')  # produces sharded dir with stage3 multi-device
# after
# after training, consolidate:
# python zero_to_f32.py path/to/sharded_ckpt model.ckpt

assert not (strategy.zero_stage_3 and trainer.num_devices > 1) or allow_shards, 'stage-3 multi-device saves sharded dir'

When it happens

Trigger: trainer.save_checkpoint(...) or ModelCheckpoint with Trainer(strategy=deepspeed_stage_3) on multi-GPU/multi-node.

Common situations: Expecting a portable single-file checkpoint for inference/serving after stage-3 training; resuming on different hardware.

Related errors

• DeepSpeed was unable to load the checkpoint. Ensure you pass
• {default_message}. It looks like you passed the path to a su
• {default_message}. It looks like you passed the path to a fi
• The provided path is not a valid DeepSpeed checkpoint: {path
• SWA does not currently support sharded models.

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