code: def _assert_saved_model_is_equal(fabric, model, checkpoint_path):
      """Convert the saved checkpoint to a single file with the model weights consolidated to easily verify the full
      weights in float32 precision."""
      from deepspeed.utils.zero_to_fp32 import (
          convert_zero_checkpoint_to_fp32_state_dict,
      )
  
      assert isinstance(fabric.strategy, DeepSpeedStrategy)
  
      # carry out the check only on rank 0
      if fabric.is_global_zero:
          if fabric.strategy.config["zero_optimization"]["stage"] in (
              2,
              3,
          ):
              single_ckpt_path = checkpoint_path / "single_model.pt"
              # the tag is hardcoded in DeepSpeedStrategy
              convert_zero_checkpoint_to_fp32_state_dict(
                  checkpoint_path, single_ckpt_path, tag="checkpoint"
              )
              state_dict = torch.load(single_ckpt_path)
          else:
              # 'checkpoint' is the tag, hardcoded in DeepSpeedStrategy
              single_ckpt_path = (
                  checkpoint_path
                  / "checkpoint"
                  / "mp_rank_00_model_states.pt"
              )
              state_dict = torch.load(single_ckpt_path)["module"]
  
          model = model.cpu()
  
          # assert model parameters are identical after loading
          for orig_param, saved_model_param in zip(
              model.parameters(), state_dict.values()
          ):
              # perform the equality check in the same precision
              saved_model_param = saved_model_param.cpu().to(
                  orig_param.dtype
              )
              assert torch.equal(orig_param, saved_model_param)
  
      fabric.barrier()
  
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 9 for calling a global function with example usage: _assert_saved_model_is_equal(fabric, model, checkpoint_path) without return types