def test_train_save_load(distributed, precision, tmp_path):
    """Test 2D-parallel training, saving and loading precision settings."""
    strategy = ModelParallelStrategy(
        _parallelize_feed_forward_fsdp2_tp,
        data_parallel_size=2,
        tensor_parallel_size=2,
    )
    fabric = Fabric(accelerator="cuda", devices=4, strategy=strategy, precision=precision)
    fabric.launch()
    model, optimizer = _train(fabric)

    checkpoint_path = fabric.broadcast(str(tmp_path / "dist-checkpoint"))

    params_before = [p.full_tensor().clone() for p in model.parameters()]
    state = {"model": model, "optimizer": optimizer, "steps": 1}
    fabric.save(checkpoint_path, state)
    assert set(os.listdir(checkpoint_path)) == {
        ".metadata",
        "__0_0.distcp",
        "__1_0.distcp",
        "__2_0.distcp",
        "__3_0.distcp",
        "meta.pt",
    }

    # re-init all objects and resume
    strategy = ModelParallelStrategy(
        _parallelize_feed_forward_fsdp2_tp,
        data_parallel_size=2,
        tensor_parallel_size=2,
    )
    fabric = Fabric(accelerator="cuda", devices=4, strategy=strategy, precision=precision)
    fabric.launch()
    model, optimizer = _train(fabric)

    # check correctness with loaded state
    state = {"model": model, "optimizer": optimizer, "steps": 0}
    metadata = fabric.load(checkpoint_path, state)
    for p0, p1 in zip(params_before, (p.full_tensor() for p in model.parameters())):
        torch.testing.assert_close(p0, p1, atol=0, rtol=0, equal_nan=True)

    # check user data in state reloaded
    assert state["steps"] == 1
    assert not metadata

    # attempt to load a key not in the metadata checkpoint
    state = {"model": model, "coconut": 11}
    with pytest.raises(KeyError, match="The requested state contains a key 'coconut' that does not exist"):
        fabric.load(checkpoint_path, state)

    # `strict=False` ignores the missing key
    state = {"model": model, "coconut": 11}
    fabric.load(checkpoint_path, state, strict=False)
    assert state["coconut"] == 11
Test 2D-parallel training, saving and loading precision settings.
tests/tests_fabric/strategies/test_model_parallel_integration.py