code: def merge_state_dict(state_dict_rank0, state_dict_local, fp16):
      """merge two state dicts, one from data parallel rank 0,
      another only contains expert states"""
      # from megatron import print_rank_last
  
      def merge_model(state_dict_rank0, state_dict_local):
          for k, v in state_dict_local.items():
              # megatron uses both dict and OrderedDict in its state_dict
              if isinstance(v, (OrderedDict, dict)):
                  merge_model(state_dict_rank0[k], v)
              else:
                  state_dict_rank0[k] = v
  
      merge_model(state_dict_rank0["model"], state_dict_local["model"])
  
      optimizer_rank0 = (
          state_dict_rank0["optimizer"]["optimizer"]
          if fp16
          else state_dict_rank0["optimizer"]
      )
      optimizer_local = (
          state_dict_local["optimizer"]["optimizer"]
          if fp16
          else state_dict_local["optimizer"]
      )
  
      for k, v in optimizer_local["state"].items():
          optimizer_rank0["state"][k] = v
  
      if fp16:
          for group_idx, param_group in enumerate(
              state_dict_local["optimizer"]["fp32_from_fp16_params"]
          ):
              for param_in_group_idx, param in enumerate(param_group):
                  if param is not None:
                      state_dict_rank0["optimizer"][
                          "fp32_from_fp16_params"
                      ][group_idx][param_in_group_idx] = param
  
      return state_dict_rank0
  
file_path: thu-pacman-faster-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 3 for calling a global function with example usage: merge_state_dict(state_dict_rank0, state_dict_local, fp16) and returns: state_dict_rank0