def test_ddp_module_state_dict():
    """Test that the module state dict can be retrieved and loaded without the prefixed wrapper keys from DDP."""

    class DistributedDataParallelMock(MagicMock):
        def __instancecheck__(self, instance):
            # to make the strategy's `isinstance(model, DistributedDataParallel)` pass with a mock as class
            return True

    strategy = DDPStrategy(parallel_devices=[torch.device("cpu"), torch.device("cpu")])

    # Without DDP applied (no setup call)
    original_module = torch.nn.Linear(2, 3)
    original_state_dict = deepcopy(original_module.state_dict())
    retrieved_state_dict = strategy.get_module_state_dict(original_module)
    assert retrieved_state_dict.keys() == original_state_dict.keys()
    strategy.load_module_state_dict(original_module, retrieved_state_dict)

    # With DDP applied (setup called)
    with mock.patch("lightning.fabric.strategies.ddp.DistributedDataParallel", DistributedDataParallelMock):
        wrapped_module = strategy.setup_module(original_module)
        retrieved_state_dict = strategy.get_module_state_dict(wrapped_module)
    assert retrieved_state_dict.keys() == original_state_dict.keys()
    strategy.load_module_state_dict(wrapped_module, retrieved_state_dict)
    strategy.load_module_state_dict(wrapped_module, original_state_dict)
Test that the module state dict can be retrieved and loaded without the prefixed wrapper keys from DDP.
tests/tests_fabric/strategies/test_ddp.py