def test_validate_device_mesh_dimensions(num_nodes, devices, invalid_dp_size, invalid_tp_size):
    """Test passing sizes that don't multiply to the world size raises an error."""
    strategy = ModelParallelStrategy(
        parallelize_fn=(lambda m, _: m),
        data_parallel_size=invalid_dp_size,
        tensor_parallel_size=invalid_tp_size,
    )
    strategy._setup_distributed = Mock()
    strategy._accelerator = Mock()
    strategy.cluster_environment = Mock(
        world_size=Mock(return_value=(num_nodes * devices)), local_rank=Mock(return_value=1)
    )
    strategy.parallel_devices = [torch.device("cpu")] * devices
    strategy.num_nodes = num_nodes
    with pytest.raises(RuntimeError, match="multiplied should equal the world size"):
        strategy.setup_environment()
Test passing sizes that don't multiply to the world size raises an error.
tests/tests_fabric/strategies/test_model_parallel.py