code: def module_policy(self):
      from transformers.models.t5.modeling_t5 import (
          T5Attention,
          T5DenseActDense,
          T5DenseGatedActDense,
          T5LayerCrossAttention,
          T5LayerFF,
          T5LayerSelfAttention,
          T5Stack,
      )
  
      policy = {}
  
      if self.shard_config.enable_fused_normalization:
          norm_cls = FusedRMSNorm
      else:
          norm_cls = RMSNorm
  
      if self.shard_config.enable_sequence_parallelism:
          self.shard_config.enable_sequence_parallelism = False
          warnings.warn(
              "T5 dosen't support sequence parallelism now, will ignore the sequence parallelism flag."
          )
  
      if self.shard_config.enable_tensor_parallelism:
          policy[T5Stack] = ModulePolicyDescription(
              sub_module_replacement=[
                  SubModuleReplacementDescription(
                      suffix="dropout",
                      target_module=DropoutForParallelInput,
                  ),
                  SubModuleReplacementDescription(
                      suffix="embed_tokens",
                      target_module=VocabParallelEmbedding1D,
                  ),
              ]
          )
          policy[T5LayerSelfAttention] = ModulePolicyDescription(
              sub_module_replacement=[
                  SubModuleReplacementDescription(
                      suffix="dropout",
                      target_module=DropoutForParallelInput,
                  ),
              ]
          )
          policy[T5LayerCrossAttention] = ModulePolicyDescription(
              sub_module_replacement=[
                  SubModuleReplacementDescription(
                      suffix="dropout",
                      target_module=DropoutForParallelInput,
                  )
              ]
          )
          policy[T5Attention] = ModulePolicyDescription(
              attribute_replacement={
                  "d_model": self.model.config.d_model
                  // self.shard_config.tensor_parallel_size,
                  "n_heads": self.model.config.num_heads
                  // self.shard_config.tensor_parallel_size,
                  "inner_dim": self.model.config.num_heads
                  * self.model.config.d_kv
                  // self.shard_config.tensor_parallel_size,
              },
              sub_module_replacement=[
                  SubModuleReplacementDescription(
                      suffix="q",
                      target_module=Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="k",
                      target_module=Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="v",
                      target_module=Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="o",
                      target_module=Linear1D_Row,
                  ),
                  SubModuleReplacementDescription(
                      suffix="relative_attention_bias",
                      target_module=Embedding1D,
                      kwargs=dict(gather_output=False),
                      ignore_if_not_exist=True,
                  ),
              ],
          )
          policy[T5LayerFF] = ModulePolicyDescription(
              sub_module_replacement=[
                  SubModuleReplacementDescription(
                      suffix="dropout",
                      target_module=DropoutForParallelInput,
                  ),
              ]
          )
          policy[T5DenseGatedActDense] = ModulePolicyDescription(
              sub_module_replacement=[
                  SubModuleReplacementDescription(
                      suffix="wi_0 ",
                      target_module=Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="wi_1",
                      target_module=Linear1D_Row,
                  ),
                  SubModuleReplacementDescription(
                      suffix="wo",
                      target_module=Linear1D_Col,
                      kwargs=dict(gather_output=True),
                  ),
                  SubModuleReplacementDescription(
                      suffix="dropout",
                      target_module=DropoutForParallelInput,
                  ),
              ]
          )
          policy[T5DenseActDense] = ModulePolicyDescription(
              sub_module_replacement=[
                  SubModuleReplacementDescription(
                      suffix="wi",
                      target_module=Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="wo",
                      target_module=Linear1D_Row,
                  ),
                  SubModuleReplacementDescription(
                      suffix="dropout",
                      target_module=DropoutForParallelInput,
                  ),
              ]
          )
  
      # optimization configuration
      self.append_or_create_submodule_replacement(
          description=SubModuleReplacementDescription(
              suffix="layer_norm",
              target_module=norm_cls,
          ),
          policy=policy,
          target_key=T5LayerFF,
      )
      self.append_or_create_submodule_replacement(
          description=SubModuleReplacementDescription(
              suffix="layer_norm", target_module=norm_cls
          ),
          policy=policy,
          target_key=T5LayerSelfAttention,
      )
      self.append_or_create_submodule_replacement(
          description=SubModuleReplacementDescription(
              suffix="layer_norm", target_module=norm_cls
          ),
          policy=policy,
          target_key=T5LayerCrossAttention,
      )
      self.append_or_create_submodule_replacement(
          description=SubModuleReplacementDescription(
              suffix="final_layer_norm", target_module=norm_cls
          ),
          policy=policy,
          target_key=T5Stack,
      )
  
      # use flash attention
      if self.shard_config.enable_flash_attention:
          self.append_or_create_method_replacement(
              description={
                  "forward": get_t5_flash_attention_forward(),
              },
              policy=policy,
              target_key=T5Attention,
          )
  
      # use jit operator
      if self.shard_config.enable_jit_fused:
          self.append_or_create_method_replacement(
              description={
                  "forward": get_jit_fused_T5_layer_ff_forward(),
                  "dropout_add": get_jit_fused_dropout_add_func(),
              },
              policy=policy,
              target_key=T5LayerFF,
          )
          self.append_or_create_method_replacement(
              description={
                  "forward": get_T5_layer_self_attention_forward(),
                  "dropout_add": get_jit_fused_dropout_add_func(),
              },
              policy=policy,
              target_key=T5LayerSelfAttention,
          )
          self.append_or_create_method_replacement(
              description={
                  "forward": get_T5_layer_cross_attention_forward(),
                  "dropout_add": get_jit_fused_dropout_add_func(),
              },
              policy=policy,
              target_key=T5LayerCrossAttention,
          )
  
      return policy
  
imports: import warnings
  from functools import partial
  from typing import Callable, Dict, List, Tuple
  import numpy
  from torch import Tensor, nn
  from colossalai.shardformer.layer import DropoutForParallelInput, Embedding1D, FusedRMSNorm, Linear1D_Col, Linear1D_Row, RMSNorm, VocabParallelEmbedding1D
  from colossalai.shardformer.policies.base_policy import ModulePolicyDescription
  from ..modeling.jit import get_jit_fused_dropout_add_func
  from ..modeling.t5 import T5PipelineForwards, get_jit_fused_T5_layer_ff_forward, get_t5_flash_attention_forward, get_T5_layer_cross_attention_forward, get_T5_layer_self_attention_forward
  from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
label_desc: Use image node_id 3 for calling the T5BasePolicy obj's underlying member method code with example usage: obj.module_policy() and returns: policy