code: def extract_zero_shards(dir, ds_checkpoint, indices_3D):
      pp_index, tp_index, dp_index = indices_3D
      sd = ds_checkpoint.get_zero_checkpoint_state(
          pp_index=pp_index, tp_index=tp_index, dp_index=dp_index
      )
  
      # pprint(f"Processing {dp_index=} {pp_index=}, {tp_index=}")
  
      optim_sd = sd[OPTIMIZER_STATE_DICT]
      param_slice_mappings = optim_sd[PARAM_SLICE_MAPPINGS]
      universal_checkpoint_info = ds_checkpoint.get_checkpoint_info(
          UNIVERSAL_CHECKPOINT_INFO
      )
      pipeline_replicated_params = universal_checkpoint_info.get(
          PIPELINE_REPLICATED_PARAMETER_PATTERNS, []
      )
      # print(f'{pipeline_replicated_params=}')
  
      # dict
      state_groups = optim_sd[BASE_OPTIMIZER_STATE]["state"]
      # list
      fp32_groups = optim_sd[SINGLE_PARTITION_OF_FP32_GROUPS]
      param_groups_cnt = len(state_groups)
  
      for param_group_id in range(param_groups_cnt):
          flat_state = dict(
              exp_avg=state_groups[param_group_id]["exp_avg"],
              exp_avg_sq=state_groups[param_group_id]["exp_avg_sq"],
              fp32=fp32_groups[param_group_id],
          )
  
          for name, fragment_mapping in param_slice_mappings[
              param_group_id
          ].items():
              if pp_index > 0 and any(
                  re.match(pattern, name)
                  for pattern in pipeline_replicated_params
              ):
                  # Skip tied weights that are replicated in first and last pp stages
                  continue
  
              # pprint(f"dpt{dp_index}{pp_index}{tp_index} {param_group_id} {name} => {fragment_mapping.start}:{fragment_mapping.numel}")
              for state_key in flat_state.keys():
                  dump_param_fragment(
                      dir,
                      tp_index,
                      dp_index,
                      state_key,
                      flat_state[state_key],
                      name,
                      fragment_mapping.start,
                      fragment_mapping.numel,
                  )
  
imports: from functools import partial
  import argparse
  import glob
  import itertools
  import multiprocessing
  import os
  import re
  import shutil
  import torch
  import tqdm
  from deepspeed.checkpoint import DeepSpeedCheckpoint
  from deepspeed.checkpoint import OPTIMIZER_STATE_DICT, BASE_OPTIMIZER_STATE, SINGLE_PARTITION_OF_FP32_GROUPS, PARAM_SLICE_MAPPINGS, PARAM_SHAPES, PARAM, CAT_DIM, PARAM_N_SUB_PARAMS, VOCAB_TENSOR, UNIVERSAL_CHECKPOINT_INFO, VOCABULARY_PARAMETER_PATTERNS, PIPELINE_REPLICATED_PARAMETER_PATTERNS, TP_REPLICATED_PARAMETER_PATTERNS, PARAMETER_TO_AVERAGE_PATTERNS, PARAMETER_WITH_ROW_PARALLELISM_PATTERNS, PARAMETER_WITH_2_SUB_PARAMS_CAT_DIM_0
label_desc: Use image node_id 4 for calling a global function with example usage: extract_zero_shards(dir, ds_checkpoint, indices_3D) without return types