def METHOD_NAME(self, prefix=''):
    METHOD_NAME = {}
    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,
    )
    METHOD_NAME[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,
        )
        METHOD_NAME[sharded_position_embeddings_key] = sharded_position_embeddings_tensor
    return METHOD_NAME