docstr: Reshape the Embedding layer to make the embedding dimension divisible by world_size
docstr_tok: ["Reshape","the","Embedding","layer","to","make","the","embedding","dimension","divisible","by","world_size"]
code: def preprocess(self):
      # reshape the embedding layer
      r"""
      Reshape the Embedding layer to make the embedding dimension divisible by world_size
      """
      if self.shard_config.enable_tensor_parallelism:
          vocab_size = self.model.config.vocab_size
          world_size = self.shard_config.tensor_parallel_size
          if vocab_size % world_size != 0:
              new_vocab_size = (
                  vocab_size + world_size - vocab_size % world_size
              )
              self.model.resize_token_embeddings(new_vocab_size)
      return self.model
  
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 2 for calling the T5BasePolicy obj's underlying member method code with example usage: obj.preprocess() and returns: self