code: def add_lm_head_policy(self, base_policy):
      from transformers.models.bert.modeling_bert import (
          BertLMPredictionHead,
      )
  
      # optimize for tensor parallelism
      if self.shard_config.enable_tensor_parallelism:
          self.append_or_create_submodule_replacement(
              description=SubModuleReplacementDescription(
                  suffix="decoder",
                  target_module=col_nn.Linear1D_Col,
                  kwargs={"gather_output": True},
              ),
              policy=base_policy,
              target_key=BertLMPredictionHead,
          )
  
      # optimize with fused normalization
      if self.shard_config.enable_fused_normalization:
          # Handle bert lm prediction head
          self.append_or_create_submodule_replacement(
              description=SubModuleReplacementDescription(
                  suffix="transform.LayerNorm",
                  target_module=col_nn.FusedLayerNorm,
              ),
              policy=base_policy,
              target_key=BertLMPredictionHead,
          )
      return base_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 4 for calling the BertPolicy obj's underlying member method code with example usage: obj.add_lm_head_policy(base_policy) and returns: base_policy