from .parallel_state import ( initialize_model_parallel, get_tensor_model_parallel_world_size, get_pipeline_model_parallel_world_size, get_data_parallel_world_size, ) from megatron.core import tensor_parallel
