code: def module_policy(
      self,
  ) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
      from transformers.models.mistral.modeling_mistral import (
          MistralAttention,
          MistralDecoderLayer,
          MistralModel,
      )
  
      policy = {}
  
      if self.shard_config.enable_sequence_parallelism:
          self.shard_config.enable_sequence_parallelism = False
          warnings.warn(
              "Mistral dosen't support sequence parallelism now, will ignore the sequence parallelism flag."
          )
  
      if self.shard_config.enable_tensor_parallelism:
          decoder_attribute_replacement = {
              "self_attn.hidden_size": self.model.config.hidden_size
              // self.shard_config.tensor_parallel_size,
              "self_attn.num_heads": self.model.config.num_attention_heads
              // self.shard_config.tensor_parallel_size,
              "self_attn.num_key_value_heads": self.model.config.num_key_value_heads
              // self.shard_config.tensor_parallel_size,
          }
  
          policy[MistralDecoderLayer] = ModulePolicyDescription(
              attribute_replacement=decoder_attribute_replacement,
              sub_module_replacement=[
                  SubModuleReplacementDescription(
                      suffix="self_attn.q_proj",
                      target_module=Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="self_attn.k_proj",
                      target_module=Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="self_attn.v_proj",
                      target_module=Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="self_attn.o_proj",
                      target_module=Linear1D_Row,
                  ),
                  SubModuleReplacementDescription(
                      suffix="mlp.gate_proj",
                      target_module=Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="mlp.up_proj",
                      target_module=Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="mlp.down_proj",
                      target_module=Linear1D_Row,
                  ),
              ],
          )
  
          self.append_or_create_submodule_replacement(
              description=SubModuleReplacementDescription(
                  suffix="embed_tokens",
                  target_module=VocabParallelEmbedding1D,
              ),
              policy=policy,
              target_key=MistralModel,
          )
  
      # 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,
                  ),
              ],
              policy=policy,
              target_key=MistralDecoderLayer,
          )
  
          self.append_or_create_submodule_replacement(
              description=SubModuleReplacementDescription(
                  suffix="norm",
                  target_module=FusedRMSNorm,
              ),
              policy=policy,
              target_key=MistralModel,
          )
  
      if self.shard_config.enable_flash_attention:
          self.append_or_create_method_replacement(
              description={
                  "forward": get_mistral_flash_attention_forward(),
              },
              policy=policy,
              target_key=MistralAttention,
          )
  
      return policy
  
imports: import warnings
  from typing import Dict, Union
  import torch.nn
  from colossalai.shardformer.layer import FusedRMSNorm, Linear1D_Col, Linear1D_Row, VocabParallelEmbedding1D
  from ..modeling.mistral import get_mistral_flash_attention_forward
  from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
label_desc: Use image node_id 3 for calling the MistralPolicy obj's underlying member method code with example usage: obj.module_policy() and returns: policy