def test_save_checkpoint_one_dist_module_required(tmp_path):
    """Test that the ModelParallelStrategy strategy can only save one distributed model per checkpoint."""
    strategy = ModelParallelStrategy(parallelize_fn=(lambda m, _: m))

    # missing DTensor model
    with pytest.raises(ValueError, match="Could not find a distributed model in the provided checkpoint state."):
        strategy.save_checkpoint(path=tmp_path, state={})
    with pytest.raises(ValueError, match="Could not find a distributed model in the provided checkpoint state."):
        strategy.save_checkpoint(path=tmp_path, state={"model": torch.nn.Linear(3, 3)})

    # multiple DTensor models
    with mock.patch("lightning.fabric.strategies.model_parallel._has_dtensor_modules", return_value=True):
        model1 = Mock(spec=nn.Module)
        model1.modules.return_value = [model1]
        model2 = Mock(spec=nn.Module)
        model2.modules.return_value = [model2]
        with pytest.raises(ValueError, match="Found multiple distributed models in the given state."):
            strategy.save_checkpoint(path=tmp_path, state={"model1": model1, "model2": model2})
Test that the ModelParallelStrategy strategy can only save one distributed model per checkpoint.
tests/tests_fabric/strategies/test_model_parallel.py