def test_save_full_state_dict(distributed, tmp_path):
    """Test that ModelParallelStrategy saves the full state into a single file with
    `save_distributed_checkpoint=False`."""
    from torch.distributed.checkpoint.state_dict import get_optimizer_state_dict

    strategy = ModelParallelStrategy(
        _parallelize_feed_forward_fsdp2,
        data_parallel_size=2,
        tensor_parallel_size=1,
        save_distributed_checkpoint=False,
    )
    fabric = Fabric(accelerator="cuda", strategy=strategy, devices=2)
    fabric.launch()
    model, optimizer = _train(fabric)

    checkpoint_path = Path(fabric.broadcast(str(tmp_path / "fsdp-checkpoint.pt")))
    state = {"model": model, "optimizer": optimizer, "steps": 1}
    fabric.save(checkpoint_path, state)

    checkpoint = torch.load(checkpoint_path, weights_only=True)
    assert checkpoint["steps"] == 1
    loaded_state_dict = checkpoint["model"]

    # assert the correct state model was saved
    state_dict = model.state_dict()
    assert set(loaded_state_dict.keys()) == set(state_dict.keys())
    for param_name in state_dict:
        assert torch.equal(loaded_state_dict[param_name], state_dict[param_name].full_tensor().cpu())
    params_before = [p.full_tensor().cpu() for p in model.parameters()]

    # assert the correct optimizer state was saved
    optimizer_state_before = get_optimizer_state_dict(model, optimizer)
    assert set(checkpoint["optimizer"].keys()) == set(optimizer_state_before.keys()) == {"state", "param_groups"}

    # 1. verify the FSDP state can be loaded back into a FSDP model/strategy directly
    strategy = ModelParallelStrategy(_parallelize_feed_forward_fsdp2, data_parallel_size=2, tensor_parallel_size=1)
    fabric = Fabric(accelerator="cuda", strategy=strategy, devices=2)
    fabric.launch()
    model, optimizer = _train(fabric)

    metadata = fabric.load(checkpoint_path, {"model": model, "optimizer": optimizer})
    assert metadata == {"steps": 1}

    params_after = [p.full_tensor() for p in model.parameters()]
    assert all(torch.equal(p0.cpu(), p1.cpu()) for p0, p1 in zip(params_before, params_after))

    optimizer_state_after = get_optimizer_state_dict(model, optimizer)
    optimizer_state_after["param_groups"][0]["betas"] = tuple(optimizer_state_after["param_groups"][0]["betas"])
    assert set(optimizer_state_after.keys()) == set(optimizer_state_before.keys()) == {"state", "param_groups"}
    torch.testing.assert_close(optimizer_state_after["state"], optimizer_state_before["state"], atol=0, rtol=0)
    assert optimizer_state_after["param_groups"] == optimizer_state_before["param_groups"]

    # run a step to verify the optimizer state is correct
    _train(fabric, model, optimizer)

    # 2. verify the FSDP state can be loaded back into a single-device model/strategy
    fabric = Fabric(accelerator="cpu", devices=1)
    model, optimizer = _train(fabric)
    metadata = fabric.load(checkpoint_path, {"model": model, "optimizer": optimizer})
    assert metadata == {"steps": 1}
    params_after = list(model.parameters())
    assert all(torch.equal(p0, p1) for p0, p1 in zip(params_before, params_after))

    # get optimizer state after loading
    normal_checkpoint_path = Path(fabric.broadcast(str(tmp_path / "normal-checkpoint.pt")))
    fabric.save(normal_checkpoint_path, {"model": model, "optimizer": optimizer, "steps": 2})

    optimizer_state_after = torch.load(normal_checkpoint_path, weights_only=True)["optimizer"]
    assert set(optimizer_state_after.keys()) == set(optimizer_state_before.keys()) == {"state", "param_groups"}
    assert torch.equal(
        optimizer_state_after["state"][0]["exp_avg"],
        optimizer_state_before["state"]["_forward_module.w1.weight"]["exp_avg"].full_tensor().cpu(),
    )

    # run a step to verify the optimizer state is correct
    _train(fabric, model, optimizer)

    # 3. verify that a single-device model/strategy states can be loaded into a FSDP model/strategy
    strategy = ModelParallelStrategy(_parallelize_feed_forward_fsdp2, data_parallel_size=2, tensor_parallel_size=1)
    fabric = Fabric(accelerator="cuda", strategy=strategy, devices=2)
    fabric.launch()
    model, optimizer = _train(fabric)

    metadata = fabric.load(normal_checkpoint_path, {"model": model, "optimizer": optimizer})
    assert metadata == {"steps": 2}

    params_after = [p.full_tensor() for p in model.parameters()]
    assert all(torch.equal(p0.cpu(), p1.cpu()) for p0, p1 in zip(params_before, params_after))

    optimizer_state_after = get_optimizer_state_dict(model, optimizer)
    optimizer_state_after["param_groups"][0]["betas"] = tuple(optimizer_state_after["param_groups"][0]["betas"])
    assert set(optimizer_state_after.keys()) == set(optimizer_state_before.keys()) == {"state", "param_groups"}
    torch.testing.assert_close(optimizer_state_after["state"], optimizer_state_before["state"], atol=0, rtol=0)
    assert optimizer_state_after["param_groups"] == optimizer_state_before["param_groups"]

    # run a step to verify the optimizer state is correct
    _train(fabric, model, optimizer)
Test that ModelParallelStrategy saves the full state into a single file with
    `save_distributed_checkpoint=False`.
tests/tests_fabric/strategies/test_model_parallel_integration.py