code: def get_fmoe_checkpoint_name(
      checkpoints_path, iteration, release=False, data_parallel_rank=-1
  ):
      """A unified checkpoint name, allowing specifying a data parallel rank"""
      from megatron import mpu
      from megatron.checkpointing import get_checkpoint_name
  
      if data_parallel_rank == -1:
          data_parallel_rank = mpu.get_data_parallel_rank()
      if data_parallel_rank == 0:
          return get_checkpoint_name(
              checkpoints_path, iteration, release
          )
  
      if release:
          directory = "release"
      else:
          directory = "iter_{:07d}".format(iteration)
      # Use both the tensor and pipeline MP rank.
      if mpu.get_pipeline_model_parallel_world_size() == 1:
          return os.path.join(
              checkpoints_path,
              directory,
              "mp_rank_{:02d}_dp_rank_{:04d}".format(
                  mpu.get_tensor_model_parallel_rank(),
                  data_parallel_rank,
              ),
              "model_optim_rng.pt",
          )
      return os.path.join(
          checkpoints_path,
          directory,
          "mp_rank_{:02d}_{:03d}_dp_rank_{:04d}".format(
              mpu.get_tensor_model_parallel_rank(),
              mpu.get_pipeline_model_parallel_rank(),
              data_parallel_rank,
          ),
          "model_optim_rng.pt",
      )
  
file_path: smart-moe/fmoe/megatron/checkpoint.py
imports: import os
  import sys
  import random
  from collections import OrderedDict
  import numpy
  import torch
label_desc: Use image node_id 1 for calling a global function with example usage: get_fmoe_checkpoint_name(checkpoints_path, iteration, release, data_parallel_rank) and returns: os, get_checkpoint_name, os