code: def parse_model_states(files):
      zero_model_states = []
      for file in files:
          state_dict = torch.load(file, map_location=device)
  
          if BUFFER_NAMES not in state_dict:
              raise ValueError(
                  f"{file} is not a model state checkpoint"
              )
          buffer_names = state_dict[BUFFER_NAMES]
          if debug:
              print("Found buffers:", buffer_names)
  
          # recover just the buffers while restoring them to fp32 if they were saved in fp16
          buffers = {
              k: v.float()
              for k, v in state_dict["module"].items()
              if k in buffer_names
          }
          param_shapes = state_dict[PARAM_SHAPES]
  
          # collect parameters that are included in param_shapes
          param_names = []
          for s in param_shapes:
              for name in s.keys():
                  param_names.append(name)
  
          # update with frozen parameters
          frozen_param_shapes = state_dict.get(
              FROZEN_PARAM_SHAPES, None
          )
          if frozen_param_shapes is not None:
              if debug:
                  print(
                      f"Found frozen_param_shapes: {frozen_param_shapes}"
                  )
              param_names += list(frozen_param_shapes.keys())
  
          # handle shared params
          shared_params = [
              [k, v] for k, v in state_dict["shared_params"].items()
          ]
  
          ds_version = state_dict.get(DS_VERSION, None)
  
          frozen_param_fragments = state_dict.get(
              FROZEN_PARAM_FRAGMENTS, None
          )
  
          z_model_state = zero_model_state(
              buffers=buffers,
              param_shapes=param_shapes,
              shared_params=shared_params,
              ds_version=ds_version,
              frozen_param_shapes=frozen_param_shapes,
              frozen_param_fragments=frozen_param_fragments,
          )
          zero_model_states.append(z_model_state)
  
      return zero_model_states
  
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 7 for calling a global function with example usage: parse_model_states(files) and returns: zero_model_states