class_docstr: Language model embeddings.
  
  Arguments:
      config (TransformerConfig): config object with all necessary configs for TransformerBlock
      vocab_size (int): vocabulary size
      max_sequence_length (int): maximum size of sequence. This
                           is used for positional embedding
      add_position_embedding (bool): Add a position embedding.
      embedding_dropout_prob (float): dropout probability for embeddings
      num_tokentypes (int): Set to 0 without binary head, and 2 with a binary head . Defaults to 0.
class_docstr_tok: ["Language","model","embeddings",".","Arguments",":","config","(","TransformerConfig",")",":","config","object","with","all","necessary","configs","for","TransformerBlock","vocab_size","(","int",")",":","vocabulary","size","max_sequence_length","(","int",")",":","maximum","size","of","sequence",".","This","is","used","for","positional","embedding","add_position_embedding","(","bool",")",":","Add","a","position","embedding",".","embedding_dropout_prob","(","float",")",":","dropout","probability","for","embeddings","num_tokentypes","(","int",")",":","Set","to","0","without","binary","head",",","and","2","with","a","binary","head",".","Defaults","to","0","."]
code: def sharded_state_dict(self, prefix=""):
      sharded_state_dict = {}
  
      word_embeddings_prefix = f"{prefix}word_embeddings."
      word_embeddings_state_dict = self.word_embeddings.state_dict(
          prefix=word_embeddings_prefix, keep_vars=True
      )
  
      sharded_word_embeddings_key = f"{word_embeddings_prefix}weight"
      sharded_word_embeddings_tensor = (
          make_tp_sharded_tensor_for_checkpoint(
              tensor=word_embeddings_state_dict[
                  sharded_word_embeddings_key
              ],
              key=sharded_word_embeddings_key,
              allow_shape_mismatch=True,
          )
      )
      sharded_state_dict[
          sharded_word_embeddings_key
      ] = sharded_word_embeddings_tensor
  
      if self.add_position_embedding:
          position_embeddings_prefix = f"{prefix}position_embeddings."
          position_embeddings_state_dict = (
              self.position_embeddings.state_dict(
                  prefix=position_embeddings_prefix, keep_vars=True
              )
          )
          sharded_position_embeddings_key = (
              f"{position_embeddings_prefix}weight"
          )
          sharded_position_embeddings_tensor = (
              make_sharded_tensor_for_checkpoint(
                  tensor=position_embeddings_state_dict[
                      sharded_position_embeddings_key
                  ],
                  key=sharded_position_embeddings_key,
              )
          )
          sharded_state_dict[
              sharded_position_embeddings_key
          ] = sharded_position_embeddings_tensor
  
      return sharded_state_dict
  
file_path: megatron-lm/megatron/core/models/common/embeddings/language_model_embedding.py
imports: from typing import Literal, Optional
  import torch
  from torch import Tensor
  from megatron.core import tensor_parallel
  from megatron.core.transformer.module import MegatronModule
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.utils import make_sharded_tensor_for_checkpoint, make_tp_sharded_tensor_for_checkpoint
label_desc: Use image node_id 4 for calling the LanguageModelEmbedding obj's underlying member method code with example usage: obj.sharded_state_dict(prefix) and returns: sharded_state_dict