code: def get_model_state_file(checkpoint_dir, zero_stage):
      if not os.path.isdir(checkpoint_dir):
          raise FileNotFoundError(
              f"Directory '{checkpoint_dir}' doesn't exist"
          )
  
      # there should be only one file
      if zero_stage <= 2:
          file = os.path.join(
              checkpoint_dir, "mp_rank_00_model_states.pt"
          )
      elif zero_stage == 3:
          file = os.path.join(
              checkpoint_dir,
              "zero_pp_rank_0_mp_rank_00_model_states.pt",
          )
  
      if not os.path.exists(file):
          raise FileNotFoundError(
              f"can't find model states file at '{file}'"
          )
  
      return file
  
imports: import argparse
  import torch
  import glob
  import math
  import os
  import re
  from collections import OrderedDict
  from dataclasses import dataclass
  from deepspeed.utils import logger
  from deepspeed.checkpoint.constants import DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS, FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES, FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS
label_desc: Use image node_id 3 for calling a global function with example usage: get_model_state_file(checkpoint_dir, zero_stage) and returns: file