code: def module_policy(
      self,
  ) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
      from colossalai.shardformer.modeling.chatglm2_6b.modeling_chatglm import (
          ChatGLMModel,
          CoreAttention,
          GLMBlock,
      )
  
      policy = {}
  
      if self.shard_config.enable_fused_normalization:
          if self.model.config.rmsnorm:
              norm_cls = col_nn.FusedRMSNorm
          else:
              norm_cls = col_nn.FusedLayerNorm
      else:
          if self.model.config.rmsnorm:
              norm_cls = col_nn.RMSNorm
          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[ChatGLMModel] = ModulePolicyDescription(
              attribute_replacement={},
              sub_module_replacement=[
                  SubModuleReplacementDescription(
                      suffix="embedding.word_embeddings",
                      target_module=col_nn.VocabParallelEmbedding1D,
                  )
              ],
          )
  
          policy[GLMBlock] = ModulePolicyDescription(
              attribute_replacement={
                  "self_attention.num_attention_heads_per_partition": self.model.config.num_attention_heads
                  // self.shard_config.tensor_parallel_size,
                  "self_attention.projection_size": (
                      self.model.config.kv_channels
                      * self.model.config.num_attention_heads
                  )
                  // self.shard_config.tensor_parallel_size,
                  "self_attention.qkv_hidden_size": (
                      self.model.config.kv_channels
                      * self.model.config.num_attention_heads
                      * 3
                  )
                  // self.shard_config.tensor_parallel_size,
                  "self_attention.core_attention.num_attention_heads_per_partition": self.model.config.num_attention_heads
                  // self.shard_config.tensor_parallel_size,
                  "self_attention.core_attention.hidden_size_per_partition": self.model.config.kv_channels
                  * self.model.config.num_attention_heads
                  // self.shard_config.tensor_parallel_size,
              },
              param_replacement=[],
              sub_module_replacement=[
                  SubModuleReplacementDescription(
                      suffix="self_attention.query_key_value",
                      target_module=col_nn.Linear1D_Col,
                      kwargs={
                          "seq_parallel": use_sequence_parallel,
                          "seq_parallel_dim": 0,
                          "overlap": overlap,
                      },
                  ),
                  SubModuleReplacementDescription(
                      suffix="self_attention.dense",
                      target_module=col_nn.Linear1D_Row,
                      kwargs={
                          "seq_parallel": use_sequence_parallel,
                          "seq_parallel_dim": 0,
                      },
                  ),
                  SubModuleReplacementDescription(
                      suffix="self_attention.core_attention.attention_dropout",
                      target_module=col_nn.DropoutForParallelInput,
                  ),
              ],
          )
      # optimization configuration
      self.append_or_create_submodule_replacement(
          description=[
              SubModuleReplacementDescription(
                  suffix="input_layernorm",
                  target_module=norm_cls,
                  kwargs={"sp_partial_derived": use_sequence_parallel},
              ),
              SubModuleReplacementDescription(
                  suffix="post_attention_layernorm",
                  target_module=norm_cls,
                  kwargs={"sp_partial_derived": use_sequence_parallel},
              ),
          ],
          policy=policy,
          target_key=GLMBlock,
      )
  
      if self.model.config.post_layer_norm:
          self.append_or_create_submodule_replacement(
              description=[
                  SubModuleReplacementDescription(
                      suffix="encoder.final_layernorm",
                      target_module=norm_cls,
                  )
              ],
              policy=policy,
              target_key=ChatGLMModel,
          )
  
      # use flash attention
      if self.shard_config.enable_flash_attention:
          self.append_or_create_method_replacement(
              description={
                  "forward": get_flash_core_attention_forward(),
              },
              policy=policy,
              target_key=CoreAttention,
          )
  
      # use sequence parallel
      if use_sequence_parallel:
          self.append_or_create_method_replacement(
              description={
                  "forward": get_chatglm_sequence_parallel_forward_fn(
                      self.shard_config
                  )
              },
              policy=policy,
              target_key=ChatGLMModel,
          )
  
      # use jit fused operator
      if self.shard_config.enable_jit_fused:
          self.append_or_create_method_replacement(
              description={
                  "forward": get_jit_fused_glm_block_forward(),
                  "dropout_add": get_jit_fused_dropout_add_func(),
              },
              policy=policy,
              target_key=GLMBlock,
          )
  
      return policy
  
imports: from functools import partial
  from typing import Callable, Dict, List, Union
  import torch.nn
  from torch import Tensor
  import colossalai.shardformer.layer
  from colossalai.shardformer.modeling.chatglm2 import ChatGLMPipelineForwards
  from colossalai.shardformer.modeling.chatglm2_6b.modeling_chatglm import ChatGLMForConditionalGeneration, ChatGLMModel
  from ..modeling.chatglm2 import get_chatglm_sequence_parallel_forward_fn, get_flash_core_attention_forward, get_jit_fused_glm_block_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 ChatGLMPolicy obj's underlying member method code with example usage: obj.module_policy() and returns: policy