code: def save_checkpoint(iteration, model, optimizer, lr_scheduler):
      """Save a model checkpoint with expert parallel"""
      # TODO: update patch
      from megatron import get_args
      from megatron import mpu
      from megatron import print_rank_last
      from megatron import utils
  
      expert_dp_comm = "none"
  
      if mpu.get_data_parallel_rank() == 0:
          # at dp rank 0, we still follows the native load_checkpoint by megatron
          from megatron.checkpointing import (
              save_checkpoint as save_checkpoint_native,
          )
  
          save_checkpoint_native(
              iteration, model, optimizer, lr_scheduler
          )
          return
  
      args = get_args()
  
      # Only rank zero of the data parallel writes to the disk.
      try:
          model = utils.unwrap_model(model)
      except AttributeError:
          # fallback to the old way of unwrapping a model
          if hasattr(model, "module"):
              model = model.module
          model = [
              model,
          ]
  
      print_rank_last(
          "saving checkpoint at iteration {:7d} to {}".format(
              iteration, args.save
          )
      )
  
      # Arguments, iteration, and model.
      state_dict = {}
      assert (
          len(model) == 1
      ), "FMoE does not support interleaved pipelining, i.e., only supports len(model) == 1 for now."
      state_dict["model"] = model[0].state_dict_for_save_checkpoint(
          keep_vars=(mpu.get_data_parallel_rank() > 0)
      )
  
      def extract_expert_param(state_dict, expert_dp_comm="none"):
          state_dict_new = state_dict.__class__()
          for k, v in state_dict.items():
              # megatron uses both dict and OrderedDict in its state_dict
              if isinstance(v, (OrderedDict, dict)):
                  v_new = extract_expert_param(v, expert_dp_comm)
                  if len(v_new) > 0:
                      state_dict_new[k] = v_new
              elif (
                  hasattr(v, "dp_comm") and v.dp_comm == expert_dp_comm
              ):
                  state_dict_new[k] = v.detach()
          return state_dict_new
  
      state_dict["model"] = extract_expert_param(
          state_dict["model"], expert_dp_comm
      )
  
      # Optimizer stuff.
      if not args.no_save_optim:
          if optimizer is not None:
              state_dict["optimizer"] = optimizer.state_dict()
              param_global_idx = 0
              for param_group in optimizer.optimizer.param_groups:
                  for param in param_group["params"]:
                      if not (
                          hasattr(param, "dp_comm")
                          and param.dp_comm == expert_dp_comm
                      ):
                          # this parameter is not an expert parameter
                          # thus there is no need to save its state in current rank
                          # since it has been saved by data parallel rank 0
                          if args.fp16:
                              # fp16 optimizer may have empty state due to overflow
                              state_dict["optimizer"]["optimizer"][
                                  "state"
                              ].pop(param_global_idx, None)
                          else:
                              state_dict["optimizer"]["state"].pop(
                                  param_global_idx
                              )
                      param_global_idx += 1
              if args.fp16:
                  state_dict["optimizer"]["optimizer"].pop(
                      "param_groups"
                  )
                  # fp32_from_fp16_params in state_dict is not a copy
                  # but a reference to optimizer.fp32_from_fp16_params,
                  # changing it in state_dict will change
                  # optimizer.fp32_from_fp16_params as well
                  # thus we create an empty fp32_from_fp16_params in state_dict
                  # and only insert expert parameters.
                  fp32_from_fp16_params = state_dict["optimizer"][
                      "fp32_from_fp16_params"
                  ]
                  state_dict["optimizer"]["fp32_from_fp16_params"] = []
                  for param_group in fp32_from_fp16_params:
                      param_group_copy = []
                      for param in param_group:
                          param_copy = (
                              param
                              if hasattr(param, "dp_comm")
                              and param.dp_comm == expert_dp_comm
                              else None
                          )
                          param_group_copy.append(param_copy)
                      state_dict["optimizer"][
                          "fp32_from_fp16_params"
                      ].append(param_group_copy)
              else:
                  state_dict["optimizer"].pop("param_groups")
  
      # Save.
      checkpoint_name = get_fmoe_checkpoint_name(args.save, iteration)
      from megatron.checkpointing import ensure_directory_exists
      from megatron.checkpointing import get_checkpoint_tracker_filename
  
      ensure_directory_exists(checkpoint_name)
      torch.save(state_dict, checkpoint_name)
  
      # Wait so everyone is done (necessary)
      torch.distributed.barrier()
      if torch.distributed.get_rank() == 0:
          print(
              "  successfully saved checkpoint at iteration {:7d} to {}".format(
                  iteration, args.save
              ),
              flush=True,
          )
      # And update the latest iteration
      if torch.distributed.get_rank() == 0:
          tracker_filename = get_checkpoint_tracker_filename(args.save)
          with open(tracker_filename, "w") as f:
              f.write(str(iteration))
      # Wait so everyone is done (not necessary)
      torch.distributed.barrier()
  
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 2 for calling a global function with example usage: save_checkpoint(iteration, model, optimizer, lr_scheduler) and returns: state_dict_new