Add API docs for NeMo Megatron

docs/source/_static/css/custom.css

@@ -255,7 +255,7 @@ article ul {
 	}
 }
 
-@media (min-width: 1400px) {
+@media (min-width: none) {
 	body {
 		font-size: 18px;
 	}

docs/source/conf.py

@@ -28,7 +28,6 @@
 
 sys.path.insert(0, os.path.abspath("../.."))
 sys.path.insert(0, os.path.abspath("../../nemo"))
-sys.path.insert(0, os.path.abspath("../../nemo_text_processing"))
 
 from package_info import __version__
 
@@ -47,18 +46,20 @@
     'hydra',  # hydra-core in requirements, hydra during import
     'dateutil',  # part of core python
     'transformers.tokenization_bert',  # has ., troublesome for this regex
-    'megatron',  # megatron-lm in requirements, megatron in import
     'sklearn',  # scikit_learn in requirements, sklearn in import
     'nemo_text_processing.inverse_text_normalization',  # Not installed automatically
     'nemo_text_processing.text_normalization',  # Not installed automatically
     'attr',  # attrdict in requirements, attr in import
     'torchmetrics',  # inherited from PTL
     'lightning_utilities',  # inherited from PTL
     'apex',
+    'megatron.core',
+    'transformer_engine',
     'joblib',  # inherited from optional code
     'IPython',
     'ipadic',
     'psutil',
+    'regex',
 ]
 
 _skipped_autodoc_mock_imports = ['wrapt', 'numpy']

docs/source/nlp/api.rst

@@ -1,99 +1,142 @@
-NeMo NLP collection API
+NeMo Megatron API
 =======================
 
-Model Classes
--------------
+Pretraining Model Classes
+-------------------------
+
+.. autoclass:: nemo.collections.nlp.models.language_modeling.megatron_base_model.MegatronBaseModel
+    :show-inheritance: 
+    :no-members:
+    :members: __init__, configure_optimizers
+
+.. autoclass:: nemo.collections.nlp.models.language_modeling.megatron_gpt_model.MegatronGPTModel
+    :show-inheritance:
+    :no-members:
+    :members: generate, training_step, validation_step, build_train_valid_test_datasets, setup, on_save_checkpoint, on_load_checkpoint
+
+.. autoclass:: nemo.collections.nlp.models.language_modeling.megatron_bert_model.MegatronBertModel
+    :show-inheritance: 
+    :no-members:
+    :members: training_step, validation_step, build_train_valid_test_datasets, build_LDDL_data, setup, on_save_checkpoint, on_load_checkpoint
+
+.. autoclass:: nemo.collections.nlp.models.language_modeling.megatron_bart_model.MegatronBARTModel
+    :show-inheritance: 
+    :no-members:
+    :members: training_step, validation_step, build_train_valid_test_datasets, setup, on_save_checkpoint, on_load_checkpoint
+
+.. autoclass:: nemo.collections.nlp.models.language_modeling.megatron_retrieval_model.MegatronRetrievalModel
+    :show-inheritance: 
+    :no-members:
+    :members: generate, training_step, validation_step, build_train_valid_test_datasets, setup
+
+.. autoclass:: nemo.collections.nlp.models.language_modeling.megatron_t5_model.MegatronT5Model
+    :show-inheritance: 
+    :no-members:
+    :members: complete, encode, decode, add_special_tokens_to_tokenizer, training_step, validation_step, build_train_valid_test_datasets, setup
+
+Customization Model Classes
+---------------------------
+
+.. autoclass:: nemo.collections.nlp.models.language_modeling.megatron_gpt_sft_model.MegatronGPTSFTModel
+    :show-inheritance: 
+    :no-members:
+    :members: generate, training_step, validation_step, build_train_valid_test_datasets, setup
+
+.. autoclass:: nemo.collections.nlp.models.language_modeling.megatron_gpt_adapter_model.MegatronGPTAdapterLearningModel
+    :show-inheritance: 
+    :no-members:
+    :members: __init__, state_dict, generate, training_step, validation_step, build_train_valid_test_datasets, setup
+
+.. autoclass:: nemo.collections.nlp.models.language_modeling.megatron_gpt_adapter_model.MegatronGPTInfusedAdapterModel
+    :show-inheritance: 
+    :no-members:
+    :members: __init__, state_dict, generate, training_step, validation_step, build_train_valid_test_datasets, setup
+
+.. autoclass:: nemo.collections.nlp.models.language_modeling.megatron_gpt_prompt_learning_model.MegatronGPTPromptLearningModel
+    :show-inheritance: 
+    :no-members:
+    :members: built_virtual_prompt_dataset, generate, training_step, validation_step, build_train_valid_test_datasets, setup
+
+.. autoclass:: nemo.collections.nlp.models.language_modeling.megatron_t5_adapter_model.MegatronT5AdapterLearningModel
+    :show-inheritance: 
+    :no-members:
+    :members: __init__, state_dict, training_step, validation_step, build_train_valid_test_datasets, setup
+
+.. autoclass:: nemo.collections.nlp.models.language_modeling.megatron_t5_adapter_model.MegatronT5AdapterLearningModel
+    :show-inheritance: 
+    :no-members:
+    :members: _add_adapters_to_component, __init__, state_dict, training_step, validation_step, build_train_valid_test_datasets, setup
+
+.. autoclass:: nemo.collections.nlp.models.language_modeling.megatron_t5_adapter_model.MegatronT5InfusedAdapterModel
+    :show-inheritance: 
+    :no-members:
+    :members: _add_adapters_to_component, __init__, state_dict, training_step, validation_step, build_train_valid_test_datasets, setup
 
-.. autoclass:: nemo.collections.nlp.models.TextClassificationModel
-    :show-inheritance:
-    :members: setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact, classifytext
+Modules
+-------
 
-.. autoclass:: nemo.collections.nlp.models.GLUEModel
-    :show-inheritance:
-    :members: setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact
+.. autoclass:: nemo.collections.nlp.modules.common.megatron.module.MegatronModule
+    :show-inheritance: 
 
-.. autoclass:: nemo.collections.nlp.models.PunctuationCapitalizationModel
-    :show-inheritance:
-    :members:
+.. autoclass:: nemo.collections.nlp.modules.common.megatron.module.Float16Module
+    :show-inheritance: 
 
-.. autoclass:: nemo.collections.nlp.models.TokenClassificationModel
-    :show-inheritance:
-    :members: setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact
-
-.. autoclass:: nemo.collections.nlp.models.QAModel
-    :show-inheritance:
-    :members: setup_training_data, setup_optimization, setup_validation_data, setup_test_data, inference, validation_epoch_end, test_epoch_end
 
-.. autoclass:: nemo.collections.nlp.models.DuplexTaggerModel
-    :show-inheritance:
-    :members: setup_training_data, setup_optimization, setup_validation_data, setup_test_data, inference, validation_epoch_end, test_epoch_end
+.. autoclass:: nemo.collections.nlp.models.language_modeling.megatron.gpt_model.GPTModel
+    :show-inheritance: 
+    :no-members:
+    :members: forward
 
-.. autoclass:: nemo.collections.nlp.models.DuplexDecoderModel
-    :show-inheritance:
-    :members: setup_training_data, setup_optimization, setup_validation_data, setup_test_data, inference, validation_epoch_end, test_epoch_end
+.. autoclass:: nemo.collections.nlp.models.language_modeling.megatron.bert_model.BertModel
+    :show-inheritance: 
+    :no-members:
+    :members: forward
 
-.. autoclass:: nemo.collections.nlp.models.BERTLMModel
-    :show-inheritance:
-    :members: setup_training_data, setup_optimization
+.. autoclass:: nemo.collections.nlp.modules.common.megatron.token_level_encoder_decoder.MegatronTokenLevelEncoderDecoderModule
+    :show-inheritance: 
+    :no-members:
+    :members: forward
 
-Modules
--------
+.. autoclass:: nemo.collections.nlp.modules.common.megatron.retrieval_token_level_encoder_decoder.MegatronRetrievalTokenLevelEncoderDecoderModule
+    :show-inheritance: 
+    :no-members:
+    :members: forward
 
-.. autoclass:: nemo.collections.nlp.modules.BertModule
-    :show-inheritance:
-    :members:
-
-.. autoclass:: nemo.collections.nlp.modules.AlbertEncoder
-    :show-inheritance:
-    :members:
 
-.. autoclass:: nemo.collections.nlp.modules.BertEncoder
-    :show-inheritance:
-    :members:
-
-.. autoclass:: nemo.collections.nlp.modules.DistilBertEncoder
-    :show-inheritance:
-    :members:
+Datasets
+--------
 
-.. autoclass:: nemo.collections.nlp.modules.RobertaEncoder
-    :show-inheritance:
-    :members:
+.. autoclass:: nemo.collections.nlp.data.language_modeling.megatron.blendable_dataset.BlendableDataset
+    :show-inheritance: 
 
-.. autoclass:: nemo.collections.nlp.modules.SequenceClassifier
-    :show-inheritance:
-    :members:
+.. autoclass:: nemo.collections.nlp.data.language_modeling.megatron.gpt_dataset.GPTDataset
+    :show-inheritance: 
 
-.. autoclass:: nemo.collections.nlp.modules.SequenceRegression
-    :show-inheritance:
-    :members:
+.. autoclass:: nemo.collections.nlp.data.language_modeling.megatron.gpt_dataset.MockGPTDataset
+    :show-inheritance: 
 
-.. autoclass:: nemo.collections.nlp.modules.SequenceTokenClassifier
-    :show-inheritance:
-    :members:
+.. autoclass:: nemo.collections.nlp.data.language_modeling.megatron.bert_dataset.BertDataset
+    :show-inheritance: 
 
-.. autofunction::  nemo.collections.nlp.modules.get_lm_model
+.. autoclass:: nemo.collections.nlp.data.language_modeling.megatron.base_prompt_learning_dataset.BasePromptLearningDataset
+    :show-inheritance: 
 
-.. autofunction::  nemo.collections.nlp.modules.get_pretrained_lm_models_list
+.. autoclass:: nemo.collections.nlp.data.language_modeling.megatron.gpt_sft_dataset.GPTSFTDataset
+    :show-inheritance: 
 
-.. autofunction::  nemo.collections.nlp.modules.common.megatron.get_megatron_lm_models_list
+.. autoclass:: nemo.collections.nlp.data.language_modeling.megatron.gpt_sft_dataset.GPTSFTChatDataset
+    :show-inheritance: 
 
-Datasets
---------
+.. autoclass:: nemo.collections.nlp.data.language_modeling.megatron.retro_dataset.RETRODataset
+    :show-inheritance: 
 
-.. autoclass:: nemo.collections.nlp.data.token_classification.punctuation_capitalization_dataset.BertPunctuationCapitalizationDataset
-    :show-inheritance:
-    :members:
-    :special-members: __getitem__
+.. autoclass:: nemo.collections.nlp.data.language_modeling.megatron.t5_dataset.T5Dataset
+    :show-inheritance: 
+    :exclude-members: MAX_SEQ_LENGTH_DELTA
 
-.. autofunction:: nemo.collections.nlp.data.token_classification.punctuation_capitalization_tarred_dataset.create_tarred_dataset
+.. autoclass:: nemo.collections.nlp.data.language_modeling.megatron.t5_prompt_learning_dataset.T5PromptLearningDataset
+    :show-inheritance: 
 
-.. autoclass:: nemo.collections.nlp.data.token_classification.punctuation_capitalization_tarred_dataset.BertPunctuationCapitalizationTarredDataset
-    :show-inheritance:
-    :members:
-    :special-members: __iter__
-    :exclude-members: reinforce_type
+.. autoclass:: nemo.collections.nlp.data.language_modeling.megatron.ul2_dataset.UL2Dataset
+    :show-inheritance: 
 
-.. autoclass:: nemo.collections.nlp.data.token_classification.punctuation_capitalization_infer_dataset.BertPunctuationCapitalizationInferDataset
-    :show-inheritance:
-    :members:
-    :special-members: __getitem__

nemo/collections/nlp/models/language_modeling/megatron_base_model.py

@@ -60,18 +60,19 @@
 
 class MegatronBaseModel(NLPModel):
     """
-    Megatron base class
-    It does the following things:
-    1. Initialize the model parallel for nemo given the model parallel parameters.
-    2. Turn on all the nvidia optimizations.
-    3. If `cfg.tokenizer` is available, it loads the tokenizer and pad the vocab to the correct size for tensor model parallelism.
-    4. If using distributed optimizer, configure to be compatible with
-       O2-level optimizations and/or model parallelism.
-    5. Perform gradient clipping: `grad_clip_pl_default` triggers the
-       PyTorch Lightning default implementation, `with_distributed_adam`
-       triggers the distributed optimizer's implementation,
-       `megatron_amp_o2` triggers gradient clipping on the main grads,
-       and otherwise gradient clipping is performed on the model grads.
+    Megatron base class. All NeMo Megatron models inherit from this class.
+
+    - Initialize the model parallel world for nemo.
+    - Turn on all of the nvidia optimizations.
+    - If `cfg.tokenizer` is available, it loads the tokenizer and pad the vocab to the 
+      correct size for tensor model parallelism.
+    - If using distributed optimizer, configure to be compatible
+      with O2 level optimizations and/or model parallelism.
+    - Perform gradient clipping: `grad_clip_pl_default` triggers
+      the PyTorch Lightning default implementation, `with_distributed_adam` triggers
+      the distributed optimizer's implementation, `megatron_amp_o2` triggers gradient clipping on the main grads,
+      and otherwise gradient clipping is performed on the model grads.
+
     """
 
     def __init__(self, cfg: DictConfig, trainer: Trainer, no_lm_init=True):