code: def module_policy(self):
      from transformers.models.bert.modeling_bert import (
          BertEmbeddings,
          BertLayer,
          BertModel,
          BertOutput,
          BertSelfAttention,
          BertSelfOutput,
      )
  
      policy = {}
  
      if self.shard_config.enable_fused_normalization:
          norm_cls = col_nn.FusedLayerNorm
      else:
          norm_cls = col_nn.LayerNorm
  
      use_sequence_parallel = (
          self.shard_config.enable_sequence_parallelism
      )
      overlap = self.shard_config.enable_sequence_overlap
      if self.shard_config.enable_tensor_parallelism:
          policy[BertLayer] = ModulePolicyDescription(
              attribute_replacement={
                  "attention.self.all_head_size": self.model.config.hidden_size
                  // self.shard_config.tensor_parallel_size,
                  "crossattention.self.all_head_size": self.model.config.hidden_size
                  // self.shard_config.tensor_parallel_size,
                  "attention.self.num_attention_heads": self.model.config.num_attention_heads
                  // self.shard_config.tensor_parallel_size,
                  "crossattention.self.num_attention_heads": self.model.config.num_attention_heads
                  // self.shard_config.tensor_parallel_size,
              },
              sub_module_replacement=[
                  SubModuleReplacementDescription(
                      suffix="attention.self.query",
                      target_module=col_nn.Linear1D_Col,
                      kwargs={
                          "seq_parallel": use_sequence_parallel,
                          "overlap": overlap,
                      },
                  ),
                  SubModuleReplacementDescription(
                      suffix="attention.self.key",
                      target_module=col_nn.Linear1D_Col,
                      kwargs={
                          "seq_parallel": use_sequence_parallel,
                          "overlap": overlap,
                      },
                  ),
                  SubModuleReplacementDescription(
                      suffix="attention.self.value",
                      target_module=col_nn.Linear1D_Col,
                      kwargs={
                          "seq_parallel": use_sequence_parallel,
                          "overlap": overlap,
                      },
                  ),
                  SubModuleReplacementDescription(
                      suffix="attention.self.dropout",
                      target_module=col_nn.DropoutForParallelInput,
                  ),
                  SubModuleReplacementDescription(
                      suffix="attention.output.dense",
                      target_module=col_nn.Linear1D_Row,
                      kwargs={"seq_parallel": use_sequence_parallel},
                  ),
                  SubModuleReplacementDescription(
                      suffix="attention.output.dropout",
                      target_module=col_nn.DropoutForParallelInput,
                  ),
                  SubModuleReplacementDescription(
                      suffix="intermediate.dense",
                      target_module=col_nn.Linear1D_Col,
                      kwargs={
                          "seq_parallel": use_sequence_parallel,
                          "overlap": overlap,
                      },
                  ),
                  SubModuleReplacementDescription(
                      suffix="output.dense",
                      target_module=col_nn.Linear1D_Row,
                      kwargs={"seq_parallel": use_sequence_parallel},
                  ),
                  SubModuleReplacementDescription(
                      suffix="output.dropout",
                      target_module=col_nn.DropoutForParallelInput,
                  ),
              ],
          )
  
          policy[BertEmbeddings] = ModulePolicyDescription(
              sub_module_replacement=[
                  SubModuleReplacementDescription(
                      suffix="word_embeddings",
                      target_module=col_nn.VocabParallelEmbedding1D,
                  ),
                  SubModuleReplacementDescription(
                      suffix="dropout",
                      target_module=col_nn.DropoutForReplicatedInput,
                  ),
              ]
          )
  
      if use_sequence_parallel:
          self.append_or_create_method_replacement(
              description={
                  "forward": bert_sequence_parallel_forward_fn(
                      self.shard_config
                  )
              },
              policy=policy,
              target_key=BertModel,
          )
  
      # optimization configuration
      # Handle bert layer
      self.append_or_create_submodule_replacement(
          description=[
              SubModuleReplacementDescription(
                  suffix="attention.output.LayerNorm",
                  target_module=norm_cls,
                  kwargs={"sp_partial_derived": use_sequence_parallel},
              ),
              SubModuleReplacementDescription(
                  suffix="output.LayerNorm",
                  target_module=norm_cls,
                  kwargs={"sp_partial_derived": use_sequence_parallel},
              ),
          ],
          policy=policy,
          target_key=BertLayer,
      )
      # handle embedding layer
      self.append_or_create_submodule_replacement(
          description=[
              SubModuleReplacementDescription(
                  suffix="LayerNorm",
                  target_module=norm_cls,
              )
          ],
          policy=policy,
          target_key=BertEmbeddings,
      )
  
      # use flash attention
      if self.shard_config.enable_flash_attention:
          self.append_or_create_method_replacement(
              description={
                  "forward": get_bert_flash_attention_forward(),
              },
              policy=policy,
              target_key=BertSelfAttention,
          )
  
      # use jit operator
      if self.shard_config.enable_jit_fused:
          self.append_or_create_method_replacement(
              description={
                  "forward": get_jit_fused_bert_self_output_forward(),
                  "dropout_add": get_jit_fused_dropout_add_func(),
              },
              policy=policy,
              target_key=BertSelfOutput,
          )
          self.append_or_create_method_replacement(
              description={
                  "forward": get_jit_fused_bert_output_forward(),
                  "dropout_add": get_jit_fused_dropout_add_func(),
              },
              policy=policy,
              target_key=BertOutput,
          )
  
      return policy
  
imports: from functools import partial
  from typing import Callable, Dict, List
  import torch.nn
  from torch import Tensor
  from torch.nn import Module
  import colossalai.shardformer.layer
  from ..modeling.bert import BertPipelineForwards, bert_sequence_parallel_forward_fn, get_bert_flash_attention_forward, get_jit_fused_bert_output_forward, get_jit_fused_bert_self_output_forward
  from ..modeling.jit import get_jit_fused_dropout_add_func
  from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
label_desc: Use image node_id 3 for calling the BertPolicy obj's underlying member method code with example usage: obj.module_policy() and returns: policy