megatron.core.parallel_state.get_embedding_group

megatron.core.parallel_state.get_embedding_group()

Get the embedding group the caller rank belongs to.