def test_save_load_sharded_state_dict(distributed, tmp_path):
    """Test saving and loading with the distributed state dict format."""

    class CheckpointModel(FSDP2Model):
        def __init__(self, params_to_compare=None):
            super().__init__()
            self.params_to_compare = params_to_compare

        def on_train_start(self):
            if self.params_to_compare is None:
                return
            for p0, p1 in zip(self.params_to_compare, self.trainer.model.parameters()):
                assert torch.equal(p0, p1.full_tensor())

    seed_everything(0)

    strategy = ModelParallelStrategy(save_distributed_checkpoint=True)
    trainer_kwargs = {
        "default_root_dir": tmp_path,
        "accelerator": "cuda",
        "devices": 2,
        "max_epochs": 1,
        "enable_progress_bar": False,
        "enable_model_summary": False,
        "logger": False,
    }

    # Initial training
    model = CheckpointModel()
    trainer = Trainer(**trainer_kwargs, strategy=strategy)
    trainer.fit(model)
    params_before = [p.full_tensor() for p in trainer.model.parameters()]

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

    metadata = torch.load(checkpoint_path / "meta.pt", weights_only=True)
    assert "pytorch-lightning_version" in metadata
    assert len(metadata["callbacks"]) == 1  # model checkpoint callback
    assert "state_dict" not in metadata
    assert "optimizer_states" not in metadata

    # Load checkpoint and continue training
    trainer_kwargs.update(max_epochs=2)
    model = CheckpointModel(params_to_compare=params_before)
    strategy = ModelParallelStrategy(save_distributed_checkpoint=True)
    trainer = Trainer(**trainer_kwargs, strategy=strategy)
    trainer.fit(model, ckpt_path=checkpoint_path)
Test saving and loading with the distributed state dict format.
tests/tests_pytorch/strategies/test_model_parallel_integration.py