code: def get_moe_epsize_param_dict(
      model: nn.Module,
  ) -> Dict[int, List[nn.Parameter]]:
      """Returns a parameter dictionary, the key of which is the expert parallel
      size of every parameter. Since the parameters in data parallelism is replicated
      in each GPU, we set their ep_size to 1.
  
      Args:
          model (:class:`torch.nn.Module`): A pyTorch `nn.Module` from which we get dict.
      """
      epsize_param_dict = dict()
      for param in model.parameters():
          if not is_moe_tensor(param):
              ep_size = 1  # set ep_size to 1 for dp parameters
          else:
              ep_size = get_ep_size(param)
          if ep_size not in epsize_param_dict:
              epsize_param_dict[ep_size] = []
          epsize_param_dict[ep_size].append(param)
  
      return epsize_param_dict
  
imports: import contextlib
  import os
  from typing import Any, Callable, Dict, List, Optional, Tuple
  import torch
  import torch.distributed
  import torch.nn
  import torch.nn.functional
  from colossalai.moe.manager import MOE_MANAGER
  from colossalai.tensor.moe_tensor.api import get_dp_group, get_dp_group_ranks, get_ep_size, is_moe_tensor
  from colossalai.utils import get_current_device
label_desc: Use image node_id 6 for calling a global function with example usage: get_moe_epsize_param_dict(model) and returns: epsize_param_dict