Merge branch 'parallel-nemo' into 'main'

Changes for NeMo/lightning compatibility

See merge request ADLR/megatron-lm!104

megatron/__init__.py

@@ -30,7 +30,7 @@ from .global_vars import get_tokenizer
 from .global_vars import get_tensorboard_writer
 from .global_vars import get_adlr_autoresume
 from .global_vars import get_timers
-
+from .initialize  import initialize_megatron
 
 def print_rank_0(message):
     """If distributed is initialized print only on rank 0."""
@@ -38,4 +38,4 @@ def print_rank_0(message):
         if torch.distributed.get_rank() == 0:
             print(message, flush=True)
     else:
-        print(message, flush=True)
+        print(message, flush=True)
\ No newline at end of file

megatron/arguments.py

@@ -322,7 +322,10 @@ def _add_distributed_args(parser):
                        'to use.')
     group.add_argument('--local_rank', type=int, default=None,
                        help='local rank passed from distributed launcher.')
-
+    group.add_argument('--lazy-mpu-init', type=bool, required=False,
+                       help='If set to True, initialize_megatron() skips DDP initialization'
+                       ' and returns function to complete it instead'
+                       'This is for external DDP manager.' )
     return parser
 
 

megatron/initialize.py

@@ -26,7 +26,7 @@ from megatron import get_args
 from megatron import get_tensorboard_writer
 from megatron import mpu
 from megatron.global_vars import set_global_variables
-
+from megatron.mpu import set_model_parallel_rank, set_model_parallel_world_size
 
 def initialize_megatron(extra_args_provider=None, args_defaults={},
                         ignore_unknown_args=False, allow_no_cuda=False):
@@ -34,38 +34,52 @@ def initialize_megatron(extra_args_provider=None, args_defaults={},
     set autoresume and random seeds.
     `allow_no_cuda` should not be set unless using megatron for cpu only 
     data processing. In general this arg should not be set unless you know 
-    what you are doing."""
+    what you are doing.
+    Returns a function to finalize distributed env initialization 
+    (optionally, only when args.lazy_mpu_init == True)
+
+"""
     if not allow_no_cuda:
         # Make sure cuda is available.
         assert torch.cuda.is_available(), 'Megatron requires CUDA.'
 
-    # This is temporary WAR to make simple case like pytest calling with same args twice
-    # Need to implement clean factory init.
-    if mpu.model_parallel_is_initialized():
-        return
-    
-    
     # Parse args, build tokenizer, and set adlr-autoresume,
     # tensorboard-writer, and timers.
     set_global_variables(extra_args_provider=extra_args_provider,
                          args_defaults=args_defaults,
                          ignore_unknown_args=ignore_unknown_args)
 
-    # Pytorch distributed.
-    _initialize_distributed()
-
-    # Autoresume.
-    _init_autoresume()
+    # torch.distributed initialization
+    def finish_mpu_init():
+        args = get_args()
+        # Pytorch distributed.
+        _initialize_distributed()
+        
+        # Random seeds for reproducibility.
+        if args.rank == 0:
+            print('> setting random seeds to {} ...'.format(args.seed))
+        _set_random_seed(args.seed)
 
-    # Random seeds for reproducibility.
     args = get_args()
-    if args.rank == 0:
-        print('> setting random seeds to {} ...'.format(args.seed))
-    _set_random_seed(args.seed)
-
-    # Write arguments to tensorboard.
-    _write_args_to_tensorboard()
-
+    if  args.lazy_mpu_init:
+        # delayed initialization of DDP-related stuff
+        # We only set basic DDP globals    
+        set_model_parallel_world_size(args.model_parallel_size)
+        # and return function for external DDP manager to call when it has DDP initialized
+        set_model_parallel_rank(args.rank)    
+        return finish_mpu_init
+    else:
+        # Megatron's MPU is the master. Complete initialization right away.
+        finish_mpu_init()
+        
+        # Autoresume.
+        _init_autoresume()
+        
+        # Write arguments to tensorboard.
+        _write_args_to_tensorboard()
+        # No continuation function
+        return None
+        
 
 def _initialize_distributed():
     """Initialize torch.distributed and mpu."""
@@ -79,11 +93,6 @@ def _initialize_distributed():
                   'skipping initialization ...', flush=True)
         args.rank = torch.distributed.get_rank()
         args.world_size = torch.distributed.get_world_size()
-        if device_count > 0:
-            device = torch.cuda.current_device()
-            local_rank = args.rank % device_count
-            assert local_rank == device, \
-                'expected local-rank to be the same as rank % device-count.'
 
     else:
 

megatron/model/__init__.py

@@ -15,6 +15,7 @@
 
 from .distributed import *
 from .bert_model import BertModel
-from megatron.model.realm_model import ICTBertModel
+from .realm_model import ICTBertModel
 from .gpt2_model import GPT2Model
 from .utils import get_params_for_weight_decay_optimization
+from .language_model import get_language_model

megatron/model/language_model.py

@@ -24,7 +24,7 @@ from megatron.module import MegatronModule
 from megatron.model.transformer import ParallelTransformer
 from megatron.model.utils import openai_gelu, erf_gelu
 from megatron.model.utils import get_linear_layer
-
+from megatron.model.utils import init_method_normal, scaled_init_method_normal
 
 def parallel_lm_logits(input_, word_embeddings_weight, parallel_output,
                        bias=None):
@@ -44,7 +44,7 @@ def parallel_lm_logits(input_, word_embeddings_weight, parallel_output,
 
 
 def get_language_model(attention_mask_func, num_tokentypes, add_pooler,
-                       init_method, scaled_init_method):
+                       init_method=None, scaled_init_method=None):
     """Build language model and return along with the key to save."""
     args = get_args()
 
@@ -55,6 +55,12 @@ def get_language_model(attention_mask_func, num_tokentypes, add_pooler,
     elif args.onnx_safe:
         gelu = erf_gelu
     
+    if init_method is None:
+        init_method = init_method_normal(args.init_method_std)
+
+    if scaled_init_method is None:
+        scaled_init_method = scaled_init_method_normal(args.init_method_std, args.num_layers)
+
     # Language model.
     language_model = TransformerLanguageModel(
         attention_mask_func=attention_mask_func,

megatron/mpu/__init__.py

@@ -27,9 +27,9 @@ from .initialize import get_data_parallel_group
 from .initialize import get_data_parallel_rank
 from .initialize import get_data_parallel_world_size
 from .initialize import get_model_parallel_group
-from .initialize import get_model_parallel_rank
+from .initialize import get_model_parallel_rank, set_model_parallel_rank
 from .initialize import get_model_parallel_src_rank
-from .initialize import get_model_parallel_world_size
+from .initialize import get_model_parallel_world_size, set_model_parallel_world_size
 from .initialize import initialize_model_parallel
 from .initialize import model_parallel_is_initialized
 

megatron/mpu/mappings.py

@@ -15,20 +15,19 @@
 
 import torch
 
-from .initialize import get_model_parallel_group
+from .initialize import get_model_parallel_group, get_model_parallel_world_size, get_model_parallel_rank
 from .utils import split_tensor_along_last_dim
 
 
 def _reduce(input_):
     """All-reduce the the input tensor across model parallel group."""
-    group = get_model_parallel_group()
 
     # Bypass the function if we are using only 1 GPU.
-    if torch.distributed.get_world_size(group=group) == 1:
+    if get_model_parallel_world_size()==1:
         return input_
 
     # All-reduce.
-    torch.distributed.all_reduce(input_, group=group)
+    torch.distributed.all_reduce(input_, group=get_model_parallel_group())
 
     return input_
 
@@ -36,18 +35,17 @@ def _reduce(input_):
 def _split(input_):
     """Split the tensor along its last dimension and keep the
     corresponding slice."""
-    group = get_model_parallel_group()
 
+    world_size = get_model_parallel_world_size()
     # Bypass the function if we are using only 1 GPU.
-    if torch.distributed.get_world_size(group=group) == 1:
+    if world_size==1:
         return input_
 
     # Split along last dimension.
-    world_size = torch.distributed.get_world_size(group=group)
     input_list = split_tensor_along_last_dim(input_, world_size)
 
     # Note: torch.split does not create contiguous tensors by default.
-    rank = torch.distributed.get_rank(group=group)
+    rank = get_model_parallel_rank()
     output = input_list[rank].contiguous()
 
     return output
@@ -55,16 +53,15 @@ def _split(input_):
 
 def _gather(input_):
     """Gather tensors and concatinate along the last dimension."""
-    group = get_model_parallel_group()
 
+    world_size = get_model_parallel_world_size()
     # Bypass the function if we are using only 1 GPU.
-    if torch.distributed.get_world_size(group=group) == 1:
+    if world_size==1:
         return input_
 
     # Size and dimension.
     last_dim = input_.dim() - 1
-    rank = torch.distributed.get_rank(group=group)
-    world_size = torch.distributed.get_world_size(group=group)
+    rank = get_model_parallel_rank()
 
     tensor_list = [torch.empty_like(input_) for _ in range(world_size)]
     tensor_list[rank] = input_