Move Parallelism usage from Apex -> Megatron Core

Dockerfile

@@ -42,8 +42,13 @@ RUN apt-get update && \
   libavdevice-dev && \
   rm -rf /var/lib/apt/lists/*
 
-WORKDIR /tmp/
+WORKDIR /workspace/
+# Install Megatron-core
+RUN git clone https://github.com/aklife97/Megatron-LM.git && \
+  cd Megatron-LM && \
+  pip install -e .
 
+WORKDIR /tmp/
 # TODO: Remove once this Apex commit (2/24/23) is included in PyTorch
 # container
 RUN git clone https://github.com/NVIDIA/apex.git && \

Jenkinsfile

@@ -57,6 +57,14 @@ pipeline {
       }
     }
 
+    // TODO: remove when pip package is available
+    stage('Megatron Core installation') {
+      steps {
+        sh 'git clone https://github.com/aklife97/Megatron-LM.git && \
+            cd Megatron-LM && \
+            pip install -e .'
+      }
+    }
 
     stage('PyTorch Lightning version') {
       steps {

examples/nlp/language_modeling/megatron_ckpt_to_nemo.py

@@ -28,7 +28,7 @@
 from argparse import ArgumentParser
 
 import torch
-from apex.transformer import parallel_state
+from megatron.core import parallel_state
 from pytorch_lightning.plugins.environments import TorchElasticEnvironment
 from pytorch_lightning.trainer.trainer import Trainer
 
@@ -121,9 +121,9 @@ def convert(local_rank, rank, world_size, args):
     app_state.model_parallel_size = app_state.tensor_model_parallel_size * app_state.pipeline_model_parallel_size
 
     parallel_state.initialize_model_parallel(
-        tensor_model_parallel_size_=app_state.tensor_model_parallel_size,
-        pipeline_model_parallel_size_=app_state.pipeline_model_parallel_size,
-        pipeline_model_parallel_split_rank_=app_state.pipeline_model_parallel_split_rank,
+        tensor_model_parallel_size=app_state.tensor_model_parallel_size,
+        pipeline_model_parallel_size=app_state.pipeline_model_parallel_size,
+        pipeline_model_parallel_split_rank=app_state.pipeline_model_parallel_split_rank,
     )
 
     app_state.pipeline_model_parallel_rank = parallel_state.get_pipeline_model_parallel_rank()

examples/nlp/language_modeling/megatron_gpt_eval.py

@@ -33,11 +33,13 @@
 from nemo.utils.model_utils import inject_model_parallel_rank
 
 try:
-    from apex.transformer import parallel_state
+    from megatron.core import parallel_state
+
+    HAVE_MEGATRON_CORE = True
 
-    HAVE_APEX = True
 except (ImportError, ModuleNotFoundError):
-    HAVE_APEX = False
+
+    HAVE_MEGATRON_CORE = False
 
 """
 This is the script to run GPT text generation.

examples/nlp/language_modeling/megatron_gpt_prompt_learning_eval.py

@@ -14,7 +14,7 @@
 
 import torch
 import torch.multiprocessing as mp
-from apex.transformer import parallel_state
+from megatron.core import parallel_state
 from omegaconf import OmegaConf
 from omegaconf.omegaconf import open_dict
 from pytorch_lightning.trainer.trainer import Trainer

examples/nlp/language_modeling/megatron_lm_ckpt_to_nemo.py

@@ -42,6 +42,7 @@
 from typing import Any, Optional
 
 import torch
+from megatron.core import parallel_state
 from pytorch_lightning.core.saving import _load_state as ptl_load_state
 from pytorch_lightning.core.saving import load_hparams_from_tags_csv, load_hparams_from_yaml
 from pytorch_lightning.trainer.trainer import Trainer

examples/nlp/language_modeling/megatron_retro_eval.py

@@ -25,11 +25,13 @@
 from nemo.core.config import hydra_runner
 
 try:
-    from apex.transformer import parallel_state
+    from megatron.core import parallel_state
+
+    HAVE_MEGATRON_CORE = True
 
-    HAVE_APEX = True
 except (ImportError, ModuleNotFoundError):
-    HAVE_APEX = False
+
+    HAVE_MEGATRON_CORE = False
 
 """
 This is the script to run RETRO Model text generation.

examples/nlp/language_modeling/megatron_t5_prompt_learning_eval.py

@@ -26,11 +26,11 @@
 from nemo.utils.app_state import AppState
 
 try:
-    from apex.transformer import parallel_state
+    from megatron.core import parallel_state
 
-    HAVE_APEX = True
+    HAVE_MEGATRON_CORE = True
 except (ImportError, ModuleNotFoundError):
-    HAVE_APEX = False
+    HAVE_MEGATRON_CORE = False
 
 
 if not torch.cuda.is_available():

examples/nlp/language_modeling/tuning/megatron_gpt_adapter_eval.py

@@ -15,7 +15,7 @@
 
 import torch
 import torch.multiprocessing as mp
-from apex.transformer import parallel_state
+from megatron.core import parallel_state
 from omegaconf import OmegaConf
 from omegaconf.omegaconf import open_dict
 from pytorch_lightning.trainer.trainer import Trainer

examples/nlp/language_modeling/tuning/megatron_gpt_ia3_eval.py

@@ -15,7 +15,7 @@
 
 import torch
 import torch.multiprocessing as mp
-from apex.transformer import parallel_state
+from megatron.core import parallel_state
 from omegaconf import OmegaConf
 from omegaconf.omegaconf import open_dict
 from pytorch_lightning.trainer.trainer import Trainer

examples/nlp/language_modeling/tuning/megatron_t5_adapter_eval.py

@@ -15,7 +15,7 @@
 
 import torch
 import torch.multiprocessing as mp
-from apex.transformer import parallel_state
+from megatron.core import parallel_state
 from omegaconf import OmegaConf
 from omegaconf.omegaconf import open_dict
 from pytorch_lightning.trainer.trainer import Trainer

examples/nlp/language_modeling/tuning/megatron_t5_ia3_eval.py

@@ -15,7 +15,7 @@
 
 import torch
 import torch.multiprocessing as mp
-from apex.transformer import parallel_state
+from megatron.core import parallel_state
 from omegaconf import OmegaConf
 from omegaconf.omegaconf import open_dict
 from pytorch_lightning.trainer.trainer import Trainer

examples/nlp/machine_translation/megatron_nmt_training.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 
+import torch.multiprocessing as mp
 from omegaconf.omegaconf import OmegaConf, open_dict
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import ModelSummary
@@ -33,6 +34,8 @@
 from nemo.utils import logging
 from nemo.utils.exp_manager import exp_manager
 
+mp.set_start_method("spawn", force=True)
+
 
 @hydra_runner(config_path="conf", config_name="aayn_base_megatron")
 def main(cfg) -> None:

nemo/collections/nlp/data/glue_benchmark/glue_benchmark_dataset.py

@@ -382,6 +382,7 @@ def __init__(
         max_seq_length_decoder: int = 128,
         use_cache: bool = True,
         prefix_override: str = None,
+        pad_to_max_length: bool = True,
     ):
         """
         Processes GLUE datasets
@@ -392,10 +393,12 @@ def __init__(
             max_seq_length: max sequence length minus 2 for [CLS] and [SEP]
             use_cache: whether to use data cache
             prefix_override: if you want to override default prompt for this task specify this via a string.
+            pad_to_max_length: If true, pad to the maximum length.
         """
         super().__init__(file_name, task_name, tokenizer, max_seq_length, use_cache, compute_features=False)
         self.max_seq_length = max_seq_length
         self.max_seq_length_decoder = max_seq_length_decoder
+        self.pad_to_max_length = pad_to_max_length
         self.processor = processors[self.task_name]()
         self.prefix_override = prefix_override
         self.features = self.convert_examples_to_features()
@@ -412,9 +415,16 @@ def collate_fn(self, batch):
         dec_input = [item['text_dec'] for item in batch]
         labels = [item['labels'] for item in batch]
 
-        max_dec_input_length = max([len(item) for item in dec_input]) if dec_input else 0
         max_enc_query_length = max([len(item) for item in enc_query]) if enc_query else 0
+        max_dec_input_length = max([len(item) for item in dec_input]) if dec_input else 0
         max_label_length = max([len(item) for item in labels]) if labels else 0
+        if self.pad_to_max_length:
+            assert max_enc_query_length <= self.max_seq_length
+            assert max_dec_input_length <= self.max_seq_length_decoder
+            assert max_label_length <= self.max_seq_length_decoder
+            max_enc_query_length = self.max_seq_length
+            max_dec_input_length = self.max_seq_length_decoder
+            max_label_length = self.max_seq_length_decoder
 
         loss_mask = [([1] * (len(item))) + ([0] * (max_label_length - len(item))) for item in labels]
         enc_query = [item + [self.tokenizer.pad_id] * (max_enc_query_length - len(item)) for item in enc_query]
@@ -488,10 +498,18 @@ def __init__(
         use_cache: bool = True,
         prefix_override: str = None,
         lang_list: List[str] = None,
+        pad_to_max_length: bool = True,
     ):
         self.lang_list = set(lang_list)
         super().__init__(
-            file_name, task_name, tokenizer, max_seq_length, max_seq_length_decoder, use_cache, prefix_override
+            file_name,
+            task_name,
+            tokenizer,
+            max_seq_length,
+            max_seq_length_decoder,
+            use_cache,
+            prefix_override,
+            pad_to_max_length,
         )
         if len(lang_list) <= 0 or lang_list is None:
             raise ValueError(f"Found an empty or None lang_list for {self.task_name}")

nemo/collections/nlp/data/language_modeling/megatron/dataset_utils.py

@@ -54,13 +54,13 @@
 from nemo.utils.get_rank import is_global_rank_zero
 
 try:
-    from apex.transformer import parallel_state
+    from megatron.core import parallel_state
 
-    HAVE_APEX = True
+    HAVE_MEGATRON_CORE = True
 
 except (ImportError, ModuleNotFoundError):
 
-    HAVE_APEX = False
+    HAVE_MEGATRON_CORE = False
 
 
 DSET_TYPE_BERT = 'standard_bert'

nemo/collections/nlp/data/language_modeling/megatron/gpt_dataset.py

@@ -32,13 +32,13 @@
 from nemo.utils import logging
 
 try:
-    from apex.transformer import parallel_state
+    from megatron.core import parallel_state
 
-    HAVE_APEX = True
+    HAVE_MEGATRON_CORE = True
 
 except (ImportError, ModuleNotFoundError):
 
-    HAVE_APEX = False
+    HAVE_MEGATRON_CORE = False
 
 
 def build_dataset(cfg, trainer, data_prefix, data_impl, num_samples, seq_length, seed, skip_warmup, tokenizer, name):
@@ -303,9 +303,9 @@ def __init__(
         seed,
         drop_last=True,
     ):
-        if not HAVE_APEX:
+        if not HAVE_MEGATRON_CORE:
             raise ImportError(
-                "Apex was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
+                "megatron-core was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
             )
 
         super().__init__()
@@ -432,9 +432,9 @@ class MockGPTDataset(Dataset):
     def __init__(
         self, cfg, tokenizer, name, num_samples, seq_length, seed,
     ):
-        if not HAVE_APEX:
+        if not HAVE_MEGATRON_CORE:
             raise ImportError(
-                "Apex was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
+                "Megatron core was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
             )
 
         super().__init__()

nemo/collections/nlp/data/language_modeling/megatron/megatron_batch_samplers.py

@@ -69,9 +69,9 @@ def __init__(
             micro_batch_size: The size of each micro batch.
             global_batch_size: The size of global batch.
             data_parallel_rank: The value you can obtain via
-                `parallel_state.get_data_parallel_rank()` of apex.transformer.
+                `parallel_state.get_data_parallel_rank()` of megatron.core.
             data_parallel_size: The value you can obtain via
-                `parallel_state.get_data_parallel_world_size()` of apex.transformer.
+                `parallel_state.get_data_parallel_world_size()` of megatron.core.
         """
         # Sanity checks.
         if total_samples <= 0:

nemo/collections/nlp/data/language_modeling/megatron/retro_dataset.py

@@ -37,11 +37,13 @@
 from nemo.utils import logging
 
 try:
-    from apex.transformer import parallel_state
+    from megatron.core import parallel_state
+
+    HAVE_MEGATRON_CORE = True
 
-    HAVE_APEX = True
 except (ImportError, ModuleNotFoundError):
-    HAVE_APEX = False
+
+    HAVE_MEGATRON_CORE = False
 
 __all__ = [
     "RETRODataset",
@@ -76,9 +78,9 @@ def __init__(
         knn_index: KNNIndex,
         retrieval_index: MMapRetrievalIndexedDataset,
     ):
-        if not HAVE_APEX:
+        if not HAVE_MEGATRON_CORE:
             raise ImportError(
-                "Apex was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
+                "megatron-core was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
             )
 
         super().__init__()

nemo/collections/nlp/models/language_modeling/megatron/__init__.py

@@ -17,8 +17,8 @@
 try:
     from nemo.collections.nlp.models.language_modeling.megatron.gpt_model import GPTModel
 
-    HAVE_APEX = True
+    HAVE_MEGATRON_CORE = True
 except (ImportError, ModuleNotFoundError):
-    HAVE_APEX = False
+    HAVE_MEGATRON_CORE = False
 
 # from nemo.collections.nlp.models.language_modeling.megatron.t5_model import T5Model

nemo/collections/nlp/models/language_modeling/megatron/bert_model.py

@@ -31,16 +31,26 @@
 )
 
 try:
-    from apex.transformer import parallel_state, tensor_parallel
     from apex.transformer.enums import AttnMaskType
     from apex.transformer.tensor_parallel.layers import set_tensor_model_parallel_attributes
 
     HAVE_APEX = True
 except (ImportError, ModuleNotFoundError):
+
     HAVE_APEX = False
+
     # fake missing classes with None attributes
     AttnMaskType = ApexGuardDefaults()
 
+try:
+    from megatron.core import parallel_state, tensor_parallel
+
+    HAVE_MEGATRON_CORE = True
+
+except (ImportError, ModuleNotFoundError):
+
+    HAVE_MEGATRON_CORE = False
+
 
 def bert_extended_attention_mask(attention_mask):
     # We create a 3D attention mask from a 2D tensor mask.

nemo/collections/nlp/models/language_modeling/megatron/gpt_model.py

@@ -26,15 +26,26 @@
 )
 
 try:
-    from apex.transformer import parallel_state, tensor_parallel
     from apex.transformer.enums import AttnMaskType
 
     HAVE_APEX = True
+
 except (ImportError, ModuleNotFoundError):
-    HAVE_APEX = False
+
     # fake missing classes with None attributes
     AttnMaskType = ApexGuardDefaults()
 
+    HAVE_APEX = False
+
+try:
+    from megatron.core import parallel_state, tensor_parallel
+
+    HAVE_MEGATRON_CORE = True
+
+except (ImportError, ModuleNotFoundError):
+
+    HAVE_MEGATRON_CORE = False
+
 
 def post_language_model_processing(
     lm_output,