code: def _extract_zero_shard_files(args, ds_checkpoint, temp_dir):
      _3d_range_list = list(
          itertools.product(
              range(ds_checkpoint.pp_degree),
              range(ds_checkpoint.tp_degree),
              range(ds_checkpoint.dp_degree),
          )
      )
      # pprint(f'{_3d_range_list=}')
      work_chunks = list(
          _get_chunks(_3d_range_list, args.num_extract_workers)
      )
      # pprint(f'{work_chunks=}')
  
      # extract_zero_shards(temp_dir, ds_checkpoint, _3d_range_list[0])
      do_work = partial(extract_zero_shards, temp_dir, ds_checkpoint)
      _do_parallel_work(do_work, work_chunks, args.num_extract_workers)
  
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 10 for calling a global function with example usage: _extract_zero_shard_files(args, ds_checkpoint, temp_dir) without return types