def test_strategy_state_dict(distributed, tmp_path, save_distributed_checkpoint):
    """Test that the strategy returns the correct state dict of the LightningModule."""
    model = FSDP2Model()
    correct_state_dict = model.state_dict()  # State dict before wrapping

    strategy = ModelParallelStrategy(save_distributed_checkpoint=save_distributed_checkpoint)
    trainer = Trainer(
        default_root_dir=tmp_path,
        accelerator="cuda",
        devices=2,
        strategy=strategy,
        max_epochs=1,
        barebones=True,
    )
    trainer.fit(model)

    state_dict = trainer.strategy.lightning_module_state_dict()

    if save_distributed_checkpoint:
        # All ranks return a state dict
        assert len(state_dict) > 0
        # State dict should contain same keys as non-distributed state dict
        assert list(state_dict.keys()) == list(correct_state_dict.keys())
    else:
        if trainer.global_rank != 0:
            # The full state-dict is only returned on rank 0
            assert len(state_dict) == 0
            return
        # State dict should contain same keys as non-distributed state dict
        assert list(state_dict.keys()) == list(correct_state_dict.keys())
Test that the strategy returns the correct state dict of the LightningModule.
tests/tests_pytorch/strategies/test_model_parallel_integration.py