megatron.core.parallel_state

Description

Model and data parallel groups.

Functions

destroy_model_parallel()	Set the groups to none.
get_data_parallel_group()	Get the data parallel group the caller rank belongs to.
get_data_parallel_rank()	Return my rank for the data parallel group.
get_data_parallel_src_rank()	Calculate the global rank corresponding to the first local rank in the data parallel group.
get_data_parallel_world_size()	Return world size for the data parallel group.
get_embedding_group()	Get the embedding group the caller rank belongs to.
get_global_memory_buffer()	Return the global GlobalMemoryBuffer object
get_model_parallel_group()	Get the model parallel group the caller rank belongs to.
get_pipeline_model_parallel_first_rank()	Return the global rank of the first process in the pipeline for the current tensor parallel group
get_pipeline_model_parallel_group()	Get the pipeline model parallel group the caller rank belongs to.
get_pipeline_model_parallel_last_rank()	Return the global rank of the last process in the pipeline for the current tensor parallel group
get_pipeline_model_parallel_next_rank()	Return the global rank that follows the caller in the pipeline
get_pipeline_model_parallel_prev_rank()	Return the global rank that preceeds the caller in the pipeline
get_pipeline_model_parallel_rank()	Return my rank for the pipeline model parallel group.
get_pipeline_model_parallel_world_size()	Return world size for the pipeline model parallel group.
get_position_embedding_group()	Get the position embedding group the caller rank belongs to.
get_tensor_model_parallel_group()	Get the tensor model parallel group the caller rank belongs to.
get_tensor_model_parallel_rank()	Return my rank for the tensor model parallel group.
get_tensor_model_parallel_src_rank()	Calculate the global rank corresponding to the first local rank in the tensor model parallel group.
get_tensor_model_parallel_world_size()	Return world size for the tensor model parallel group.
get_virtual_pipeline_model_parallel_rank()	Return the virtual pipeline-parallel rank.
get_virtual_pipeline_model_parallel_world_size()	Return the virtual pipeline-parallel world size.
initialize_model_parallel([...])	Initialize model data parallel groups.
is_pipeline_first_stage([ignore_virtual])	Return True if in the first pipeline model-parallel stage, False otherwise.
is_pipeline_last_stage([ignore_virtual])	Return True if in the last pipeline model-parallel stage, False otherwise.
is_pipeline_stage_after_split([rank])	Return True if pipeline stage executes decoder block for a model with both encoder and decoder.
is_pipeline_stage_at_split()	Return true if pipeline stage executes decoder block and next stage executes encoder block for a model with both encoder and decoder.
is_pipeline_stage_before_split([rank])	Return True if pipeline stage executes encoder block for a model with both encoder and decoder.
is_rank_in_embedding_group([ignore_virtual])	Return true if current rank is in embedding group, False otherwise.
is_rank_in_position_embedding_group()	Return true if current rank is in position embedding group, False otherwise.
model_parallel_is_initialized()	Check if model and data parallel groups are initialized.
set_pipeline_model_parallel_rank(rank)	Set pipeline model parallel rank.
set_pipeline_model_parallel_split_rank(rank)	Set pipeline model parallel split rank.
set_pipeline_model_parallel_world_size(...)	Set the pipeline model parallel size
set_tensor_model_parallel_rank(rank)	Set tensor model parallel rank.
set_tensor_model_parallel_world_size(world_size)	Set the tensor model parallel size
set_virtual_pipeline_model_parallel_rank(rank)	Set the virtual pipeline-parallel rank.