def test_save_sharded_and_consolidate_and_load(tmp_path):
    """Test the consolidation of a FSDP-sharded checkpoint into a single file."""

    class CustomModel(BoringModel):
        def configure_optimizers(self):
            # Use Adam instead of SGD for this test because it has state
            # In PyTorch >= 2.4, saving an optimizer with empty state would result in a `KeyError: 'state'`
            # when loading the optimizer state-dict back.
            # TODO: To resolve this, switch to the new `torch.distributed.checkpoint` APIs in FSDPStrategy
            return torch.optim.Adam(self.parameters(), lr=0.1)

    model = CustomModel()
    trainer = Trainer(
        default_root_dir=tmp_path,
        accelerator="cuda",
        devices=2,
        strategy=FSDPStrategy(auto_wrap_policy=always_wrap_policy, state_dict_type="sharded"),
        max_steps=3,
    )
    trainer.fit(model)

    checkpoint_path_sharded = trainer.strategy.broadcast(str(trainer.checkpoint_callback.best_model_path))
    assert set(os.listdir(checkpoint_path_sharded)) == {"meta.pt", ".metadata", "__0_0.distcp", "__1_0.distcp"}

    # consolidate the checkpoint to a single file
    checkpoint_path_full = trainer.strategy.broadcast(str(tmp_path / "checkpoint_full.ckpt"))
    if trainer.global_rank == 0:
        checkpoint = _load_distributed_checkpoint(Path(checkpoint_path_sharded))
        checkpoint = _format_checkpoint(checkpoint)
        torch.save(checkpoint, checkpoint_path_full)
    trainer.strategy.barrier()

    model = CustomModel()
    trainer = Trainer(
        default_root_dir=tmp_path,
        accelerator="cuda",
        devices=2,
        strategy="ddp",
        max_steps=4,
    )
    trainer.fit(model, ckpt_path=checkpoint_path_full)
Test the consolidation of a FSDP-sharded checkpoint into a single file.
tests/tests_pytorch/strategies/test_fsdp.py