code: def module_policy(
      self,
  ) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
      policy = {}
  
      if self.shard_config.enable_sequence_parallelism:
          self.shard_config.enable_sequence_parallelism = False
          raise NotImplementedError(
              "openmoe dosen't support sequence parallelism now, will ignore the sequence parallelism flag."
          )
  
      if self.shard_config.enable_tensor_parallelism:
          raise NotImplementedError(
              "Tensor parallelism is not supported for openmoe model now."
          )
  
      # optimization configuration
      if self.shard_config.enable_fused_normalization:
          self.append_or_create_submodule_replacement(
              description=[
                  SubModuleReplacementDescription(
                      suffix="input_layernorm",
                      target_module=FusedRMSNorm,
                  ),
                  SubModuleReplacementDescription(
                      suffix="post_attention_layernorm",
                      target_module=FusedRMSNorm,
                  ),
                  SubModuleReplacementDescription(
                      suffix="pre_extra_mlp_layernorm",
                      target_module=FusedRMSNorm,
                      ignore_if_not_exist=True,
                  ),
              ],
              policy=policy,
              target_key=OpenMoeDecoderLayer,
          )
  
          self.append_or_create_submodule_replacement(
              description=SubModuleReplacementDescription(
                  suffix="norm",
                  target_module=FusedRMSNorm,
              ),
              policy=policy,
              target_key=OpenMoeModel,
          )
  
      if self.shard_config.enable_flash_attention:
          raise NotImplementedError(
              "Flash attention has already been replaced in openmoe."
          )
  
      return policy
  
imports: import warnings
  from functools import partial
  from typing import Callable, Dict, List, Optional, Union
  import torch
  import torch.nn
  import torch.nn.functional
  from torch import Tensor
  from torch.nn import Module
  from transformers.modeling_outputs import CausalLMOutputWithPast
  from transformers.utils import logging
  from colossalai.moe.manager import MOE_MANAGER
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from colossalai.shardformer.layer import FusedRMSNorm, Linear1D_Col
  from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
  from .modeling_openmoe import OpenMoeDecoderLayer, OpenMoeForCausalLM, OpenMoeModel
label_desc: Use image node_id 3 for calling the OpenMoePolicy obj's underlying member method code with example usage: obj.module_policy() and returns: policy