docstr: If under pipeline parallel setting, replacing the original forward method of huggingface
  to customized forward method, and add this changing to policy.
docstr_tok: ["If","under","pipeline","parallel","setting",",","replacing","the","original","forward","method","of","huggingface","to","customized","forward","method",",","and","add","this","changing","to","policy","."]
code: def set_pipeline_forward(
      self, model_cls: nn.Module, new_forward: Callable, policy: Dict
  ) -> None:
      """
      If under pipeline parallel setting, replacing the original forward method of huggingface
      to customized forward method, and add this changing to policy.
      """
      if self.pipeline_stage_manager is None:
          return
  
      stage_manager = self.pipeline_stage_manager
      if self.model.__class__.__name__ == "BertModel":
          module = self.model
      else:
          module = self.model.bert
  
      if stage_manager.is_interleave:
          layers_per_stage = self.distribute_layers(
              len(module.encoder.layer),
              stage_manager.num_stages * stage_manager.num_model_chunks,
          )
          stage_manager.stage_indices = Policy.get_stage_index(
              layers_per_stage,
              stage_manager.stage,
              num_model_chunks=stage_manager.num_model_chunks,
              num_stages=stage_manager.num_stages,
          )
          method_replacement = {
              "forward": partial(
                  new_forward,
                  stage_manager=stage_manager,
                  shard_config=self.shard_config,
              )
          }
  
      else:
          layers_per_stage = Policy.distribute_layers(
              len(module.encoder.layer), stage_manager.num_stages
          )
          stage_index = Policy.get_stage_index(
              layers_per_stage, stage_manager.stage
          )
          method_replacement = {
              "forward": partial(
                  new_forward,
                  stage_manager=stage_manager,
                  stage_index=stage_index,
                  shard_config=self.shard_config,
              )
          }
  
      self.append_or_create_method_replacement(
          description=method_replacement,
          policy=policy,
          target_key=model_cls,
      )
  
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 7 for calling the BertPolicy obj's underlying member method code with example usage: obj.set_pipeline_forward(model_cls, new_forward, policy) without return types