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 not self.pipeline_stage_manager:
          raise ValueError(
              "set_pipeline_forward method can only be called when pipeline parallel is enabled."
          )
      stage_manager = self.pipeline_stage_manager
  
      encoder = self.model.encoder
      decoder = getattr(self.model, "decoder", None)
  
      num_encoder_layers = len(encoder.block)
      num_decoder_layers = len(decoder.block) if decoder else 0
  
      (
          layers_per_stage,
          decoder_starting_stage,
      ) = T5BasePolicy.distribute_t5_layers(
          num_encoder_layers,
          num_decoder_layers,
          stage_manager.num_stages,
      )
      stage_index = T5BasePolicy.get_t5_stage_index(
          layers_per_stage, stage_manager.stage, decoder_starting_stage
      )
  
      method_replacement = {
          "forward": partial(
              new_forward,
              stage_manager=stage_manager,
              stage_index=stage_index,
              decoder_starting_stage=decoder_starting_stage,
          )
      }
      self.append_or_create_method_replacement(
          description=method_replacement,
          policy=policy,
          target_key=model_cls,
      )
  
imports: import warnings
  from functools import partial
  from typing import Callable, Dict, List, Tuple
  import numpy
  from torch import Tensor, nn
  from colossalai.shardformer.layer import DropoutForParallelInput, Embedding1D, FusedRMSNorm, Linear1D_Col, Linear1D_Row, RMSNorm, VocabParallelEmbedding1D
  from colossalai.shardformer.policies.base_policy import ModulePolicyDescription
  from ..modeling.jit import get_jit_fused_dropout_add_func
  from ..modeling.t5 import T5PipelineForwards, get_jit_fused_T5_layer_ff_forward, get_t5_flash_attention_forward, get_T5_layer_cross_attention_forward, get_T5_layer_self_attention_forward
  from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
label_desc: Use image node_id 8 for calling the T5BasePolicy obj's underlying member method code with example usage: obj.set_pipeline_forward(model_cls, new_forward, policy) without return types