code: def get_fp32_state_dict_from_zero1_checkpoint(
      checkpoint_dir: str,
  ) -> collections.OrderedDict:
      """
      Simplified version of DeepSpeed's FP32 state dict reconstruction for ZeRO
      stage 1.
  
      :param checkpoint_dir: DeepSpeed checkpoint directory.
      :return: FP32 model state dictionary.
      """
      # Find latest checkpoint
      with utils.smart_open(
          os.path.join(checkpoint_dir, C.DEEPSPEED_LATEST_NAME)
      ) as tag_in:
          tag = tag_in.readline().strip()
      checkpoint_dir = os.path.join(checkpoint_dir, tag)
      # Load state dictionaries from all ranks (sockeye-train workers)
      optim_files = get_optim_files(checkpoint_dir)
      state_dicts = [
          torch.load(fname, map_location=torch.device("cpu"))
          for fname in optim_files
      ]
      # Collect data from state dicts
      _world_size = state_dicts[0][OPTIMIZER_STATE_DICT][
          PARTITION_COUNT
      ]
      world_size = (
          max(_world_size)
          if isinstance(_world_size, list)
          else _world_size
      )
      fp32_flat_groups = [
          state_dict[OPTIMIZER_STATE_DICT][
              SINGLE_PARTITION_OF_FP32_GROUPS
          ]
          for state_dict in state_dicts
      ]
      # Load model state. Stage 1 uses the same state file as stage 2.
      model_state_file = get_model_state_file(
          checkpoint_dir, zero_stage=2
      )
      buffers, param_shapes, _ = parse_model_state(
          file=model_state_file
      )
      # Reconstruct FP32 state dict from all of the above
      return _get_fp32_state_dict_from_zero2_checkpoint(
          world_size, param_shapes, fp32_flat_groups, buffers
      )
  
file_path: sockeye/sockeye/convert_deepspeed.py
imports: import argparse
  import collections
  import gc
  import logging
  import os
  import shutil
  import torch
  from .None import constants
  from .None import log
  from .None import model
  from .None import utils
label_desc: Use image node_id 1 for calling a global function with example usage: get_fp32_state_dict_from_zero1_checkpoint(checkpoint_dir) and returns: _get_fp32_state_dict_from_zero2_checkpoint