code: def dump_param_fragment(
      dir,
      tp_index,
      dp_index,
      state_name,
      state_flat_tensor,
      param_name,
      offset,
      numel,
  ):
      global cnt  # temp hack
  
      param_base_path = os.path.join(dir, param_name, str(tp_index))
      os.makedirs(param_base_path, exist_ok=True)
  
      cnt += 1
      counter = f"{dp_index:0>2d}"
  
      path = os.path.join(param_base_path, f"{state_name}.{counter}")
  
      # print(f"{param_name}: {offset}: {numel} => {path}")
  
      t = state_flat_tensor.narrow(0, offset, numel).clone()
      _save_checkpoint(path, t)
  
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 5 for calling a global function with example usage: dump_param_fragment(dir, tp_index, dp_index, state_name, state_flat_tensor, param_name, offset, numel) without return types