code: def module_policy(self):
      from transformers.models.whisper.modeling_whisper import (
          WhisperAttention,
          WhisperDecoder,
          WhisperDecoderLayer,
          WhisperEncoder,
          WhisperEncoderLayer,
      )
  
      policy = {}
  
      if self.shard_config.enable_fused_normalization:
          norm_cls = col_nn.FusedLayerNorm
      else:
          norm_cls = col_nn.LayerNorm
  
      if self.shard_config.enable_sequence_parallelism:
          self.shard_config.enable_sequence_parallelism = False
          warnings.warn(
              "Whisper dosen't support sequence parallelism now, will ignore the sequence parallelism flag."
          )
  
      # TODO using the jit fused add_and_dropout affect the accuracy
      if self.shard_config.enable_jit_fused:
          self.shard_config.enable_jit_fused = False
          warnings.warn(
              "Whisper dosen't support jit fused operator now, will ignore the jit fused operator flag."
          )
  
      if self.shard_config.enable_tensor_parallelism:
          policy[WhisperEncoderLayer] = ModulePolicyDescription(
              attribute_replacement={
                  "self_attn.embed_dim": self.model.config.d_model
                  // self.shard_config.tensor_parallel_size,
                  "self_attn.num_heads": self.model.config.encoder_attention_heads
                  // self.shard_config.tensor_parallel_size,
              },
              sub_module_replacement=[
                  SubModuleReplacementDescription(
                      suffix="self_attn.q_proj",
                      target_module=col_nn.Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="self_attn.k_proj",
                      target_module=col_nn.Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="self_attn.v_proj",
                      target_module=col_nn.Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="self_attn.out_proj",
                      target_module=col_nn.Linear1D_Row,
                  ),
                  SubModuleReplacementDescription(
                      suffix="fc1",
                      target_module=col_nn.Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="fc2",
                      target_module=col_nn.Linear1D_Row,
                  ),
              ],
          )
  
          policy[WhisperDecoderLayer] = ModulePolicyDescription(
              attribute_replacement={
                  "self_attn.embed_dim": self.model.config.d_model
                  // self.shard_config.tensor_parallel_size,
                  "self_attn.num_heads": self.model.config.decoder_attention_heads
                  // self.shard_config.tensor_parallel_size,
                  "encoder_attn.embed_dim": self.model.config.d_model
                  // self.shard_config.tensor_parallel_size,
                  "encoder_attn.num_heads": self.model.config.encoder_attention_heads
                  // self.shard_config.tensor_parallel_size,
              },
              sub_module_replacement=[
                  SubModuleReplacementDescription(
                      suffix="self_attn.q_proj",
                      target_module=col_nn.Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="self_attn.k_proj",
                      target_module=col_nn.Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="self_attn.v_proj",
                      target_module=col_nn.Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="self_attn.out_proj",
                      target_module=col_nn.Linear1D_Row,
                  ),
                  SubModuleReplacementDescription(
                      suffix="encoder_attn.q_proj",
                      target_module=col_nn.Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="encoder_attn.k_proj",
                      target_module=col_nn.Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="encoder_attn.v_proj",
                      target_module=col_nn.Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="encoder_attn.out_proj",
                      target_module=col_nn.Linear1D_Row,
                  ),
                  SubModuleReplacementDescription(
                      suffix="fc1",
                      target_module=col_nn.Linear1D_Col,
                  ),
                  SubModuleReplacementDescription(
                      suffix="fc2",
                      target_module=col_nn.Linear1D_Row,
                  ),
              ],
          )
  
          policy[WhisperDecoder] = ModulePolicyDescription(
              sub_module_replacement=[
                  SubModuleReplacementDescription(
                      suffix="embed_tokens",
                      target_module=col_nn.VocabParallelEmbedding1D,
                  ),
              ]
          )
  
      # optimization configuration
      # Handle encoder layer
      self.append_or_create_submodule_replacement(
          description=[
              SubModuleReplacementDescription(
                  suffix="self_attn_layer_norm",
                  target_module=norm_cls,
              ),
              SubModuleReplacementDescription(
                  suffix="final_layer_norm",
                  target_module=norm_cls,
              ),
          ],
          policy=policy,
          target_key=WhisperEncoderLayer,
      )
  
      # Handle decoder layer
      self.append_or_create_submodule_replacement(
          description=[
              SubModuleReplacementDescription(
                  suffix="self_attn_layer_norm",
                  target_module=norm_cls,
              ),
              SubModuleReplacementDescription(
                  suffix="final_layer_norm",
                  target_module=norm_cls,
              ),
          ],
          policy=policy,
          target_key=WhisperDecoderLayer,
      )
  
      # handle encoder layer
      self.append_or_create_submodule_replacement(
          description=[
              SubModuleReplacementDescription(
                  suffix="layer_norm",
                  target_module=norm_cls,
              )
          ],
          policy=policy,
          target_key=WhisperEncoder,
      )
  
      # handle decoder layer
      self.append_or_create_submodule_replacement(
          description=[
              SubModuleReplacementDescription(
                  suffix="layer_norm",
                  target_module=norm_cls,
              )
          ],
          policy=policy,
          target_key=WhisperDecoder,
      )
  
      # enable flash attention
      if self.shard_config.enable_flash_attention:
          self.append_or_create_method_replacement(
              description={
                  "forward": get_whisper_flash_attention_forward(),
              },
              policy=policy,
              target_key=WhisperAttention,
          )
  
      # use jit fused operator
      if self.shard_config.enable_jit_fused:
          self.append_or_create_method_replacement(
              description={
                  "forward": get_jit_fused_whisper_decoder_layer_forward(),
                  "dropout_add": get_jit_fused_dropout_add_func(),
              },
              policy=policy,
              target_key=WhisperDecoderLayer,
          )
          self.append_or_create_method_replacement(
              description={
                  "forward": get_jit_fused_whisper_encoder_layer_forward(),
                  "dropout_add": get_jit_fused_dropout_add_func(),
              },
              policy=policy,
              target_key=WhisperEncoderLayer,
          )
  
      return policy
  
imports: import warnings
  from functools import partial
  from typing import Callable, Dict, List, Tuple
  import numpy
  import torch.nn
  from torch import Tensor
  import colossalai.shardformer.layer
  from ..modeling.jit import get_jit_fused_dropout_add_func
  from ..modeling.whisper import WhisperPipelineForwards, get_jit_fused_whisper_decoder_layer_forward, get_jit_fused_whisper_encoder_layer_forward, get_whisper_flash_attention_forward
  from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
label_desc: Use image node_id 4 for calling the WhisperPolicy obj's underlying member method code with example usage: obj.module_policy() and returns: policy