docstr: Distribute t5 layers into stages when pipeline parallel is used.
  Return the layer distribution as a list and the starting stage of decoder.
  If decoder doesn't exist, returned decoder starting stage is set to num_encoder_layers.
docstr_tok: ["Distribute","t5","layers","into","stages","when","pipeline","parallel","is","used",".","Return","the","layer","distribution","as","a","list","and","the","starting","stage","of","decoder",".","If","decoder","does","n't","exist",",","returned","decoder","starting","stage","is","set","to","num_encoder_layers","."]
code: def distribute_t5_layers(
      num_encoder_layers: int, num_decoder_layers: int, num_stages: int
  ) -> Tuple[List[int], int]:
      """
      Distribute t5 layers into stages when pipeline parallel is used.
      Return the layer distribution as a list and the starting stage of decoder.
      If decoder doesn't exist, returned decoder starting stage is set to num_encoder_layers.
      """
  
      # number of encoder layers must be a positive integer
      if num_encoder_layers <= 0:
          raise ValueError(
              "The number of encoder layers for T5 must be a positive integer."
          )
  
      # number of layers should be large enough to fill in every stage
      if num_encoder_layers + num_decoder_layers < num_stages:
          raise ValueError(
              "The total number of layers can't be smaller than number of stages."
          )
  
      # in the case of T5EncoderModel, set decoder starting stage to num_stages since it doesn't exist
      if num_decoder_layers == 0:
          return (
              Policy.distribute_layers(num_encoder_layers, num_stages),
              num_stages,
          )
  
      # the number of stages distributed between encoder and decoder is optmized in this way:
      # num_encoder_stages = argmin(abs(num_encoder_layers / encoder_stages - num_decoder_layers / decoder_stages))
      #                   s.t. num_encoder_stages + num_decoder_stages = num_stages, num_encoder_stages >= 1, num_decoder_stages >= 1
      def objective(num_encoder_stages):
          return abs(
              num_encoder_layers / num_encoder_stages
              - num_decoder_layers / (num_stages - num_encoder_stages)
          )
  
      num_encoder_stages = (
          np.argmin([objective(i) for i in range(1, num_stages)]) + 1
      )
      num_decoder_stages = num_stages - num_encoder_stages
  
      encoder_distribution = Policy.distribute_layers(
          num_encoder_layers, num_encoder_stages
      )
      decoder_distribution = Policy.distribute_layers(
          num_decoder_layers, num_decoder_stages
      )
      return (
          encoder_distribution + decoder_distribution,
          num_encoder_stages,
      )
  
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 5 for calling the T5BasePolicy obj's underlying member method code with example usage: obj.distribute_t5_layers(num_encoder_layers, num_decoder_layers, num_stages) and returns: unknown,  num_encoder_stages, Policy,  num_stages, abs