Move layers from mpu to core.tensor_parallel.

megatron/training.py

@@ -20,6 +20,7 @@ from megatron import get_num_microbatches
 from megatron import is_last_rank
 from megatron import update_num_microbatches
 from megatron import mpu
+from megatron import core
 from megatron import print_rank_0
 from megatron import print_rank_last
 from megatron.checkpointing import load_checkpoint
@@ -257,7 +258,7 @@ def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap
     # are set for all params so the optimizer can use them.
     for model_module in model:
         for param in model_module.parameters():
-            mpu.set_defaults_if_not_set_tensor_model_parallel_attributes(param)
+            core.tensor_parallel.set_defaults_if_not_set_tensor_model_parallel_attributes(param)
 
     # Print number of parameters.
     if mpu.get_data_parallel_rank() == 0:

megatron/utils.py

@@ -14,7 +14,7 @@ from megatron import get_args
 from megatron import get_adlr_autoresume
 from megatron import mpu
 from megatron.model.module import param_is_not_shared
-from megatron.mpu.layers import param_is_not_tensor_parallel_duplicate
+from megatron.core.tensor_parallel import param_is_not_tensor_parallel_duplicate
 
 
 def unwrap_model(model, module_instances=(torchDDP)):