def test_deepspeed_save_load_checkpoint_zero_3(stage, tmp_path):
    """Test that DeepSpeed stage 1, 2, and 3 model checkpoints can be saved and loaded successfully."""
    from deepspeed import DeepSpeedEngine

    fabric = Fabric(accelerator="cuda", devices=2, strategy=DeepSpeedStrategy(stage=stage), precision="bf16-mixed")
    fabric.launch()

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

    with fabric.init_module():
        model = BoringModel()

    optimizer = torch.optim.SGD(model.parameters(), lr=0.0001)
    model, optimizer = fabric.setup(model, optimizer)
    assert isinstance(model._forward_module, DeepSpeedEngine)

    # TODO(fabric): The dtype on the model is not correct, should be torch.bfloat16
    assert model.dtype == torch.float32
    assert next(model.parameters()).dtype == torch.bfloat16

    # dummy training step
    output = model(torch.randn(1, 32).to(fabric.device))
    loss = output.sum()
    fabric.backward(loss)
    optimizer.step()
    optimizer.zero_grad()

    state = {"model": model, "optimizer": optimizer, "steps": 1}
    fabric.save(checkpoint_path, state)

    # re-init all objects and resume
    fabric = Fabric(accelerator="cuda", devices=2, strategy=DeepSpeedStrategy(stage=stage), precision="bf16")
    fabric.launch()
    with fabric.init_module():
        model = BoringModel()

    optimizer = torch.optim.SGD(model.parameters(), lr=0.0001)
    model, optimizer = fabric.setup(model, optimizer)
    state = {"model": model, "optimizer": optimizer, "steps": 0}

    metadata = fabric.load(checkpoint_path, state)

    # check user data in state reloaded
    assert state["steps"] == 1
    # the remainder of the deepspeed checkpoint contains metadata
    assert "ds_version" in metadata

    _assert_saved_model_is_equal(fabric, model, checkpoint_path)
Test that DeepSpeed stage 1, 2, and 3 model checkpoints can be saved and loaded successfully.
tests/tests_fabric/strategies/test_deepspeed_integration.py