code: 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)
  
imports: import os
  from copy import deepcopy
  from unittest import mock
  from unittest.mock import ANY
  import pytest
  import torch
  import torch.nn
  import torch.nn.functional
  from lightning.fabric import Fabric
  from lightning.fabric.plugins import DeepSpeedPrecision
  from lightning.fabric.strategies import DeepSpeedStrategy
  from torch.utils.data import DataLoader
  from tests_fabric.helpers.models import RandomDataset, RandomIterableDataset
  from tests_fabric.helpers.runif import RunIf
  from tests_fabric.test_fabric import BoringModel
label_desc: Use image node_id 10 for calling a global function with example usage: test_deepspeed_save_load_checkpoint_zero_3(stage, tmp_path) without return types