class_docstr: A wrapper that is used to replace the DDP module provided by Megatron, which
  is adapted to enable the sophiscated parallel and reduction strategies in
  Fast MoE.
docstr: Keep consitency with Megatron
code: def state_dict_for_save_checkpoint(self, *args, **kwargs):
      r"""
      Keep consitency with Megatron
      """
      return self.module.state_dict_for_save_checkpoint(*args, **kwargs)
  
file_path: fastmoe/fmoe/megatron/distributed.py
imports: import torch
  from fmoe.distributed import DistributedGroupedDataParallel
label_desc: Use image node_id 4 for calling the DistributedDataParallel obj's underlying member method code with example usage: obj.state_dict_for_save_checkpoint() and returns: self