Documentation for ONNX export of Megatron Models (#6914)

* add Megatron ONNX export guide Signed-off-by: Asfiya Baig <[email protected]> * fix formatting Signed-off-by: Asfiya Baig <[email protected]> * include megatron_onnx_export in api.rst Signed-off-by: Asfiya Baig <[email protected]> * include megatron_onnx_export in index.rst Signed-off-by: Asfiya Baig <[email protected]> * update installation section Signed-off-by: Asfiya Baig <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * revert changes to megatron_ckpt_to_nemo.py Signed-off-by: Asfiya Baig <[email protected]> * address comments Signed-off-by: Asfiya Baig <[email protected]> --------- Signed-off-by: Asfiya Baig <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Eric Harper <[email protected]>
NVIDIA · Jun 25, 2023 · c4e677a · c4e677a
1 parent 74cbbb2
commit c4e677a
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 0 deletions.
diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -45,6 +45,7 @@ NVIDIA NeMo User Guide
    nlp/machine_translation/machine_translation
    nlp/text_normalization/intro
    nlp/api
+   nlp/megatron_onnx_export
    nlp/models
 
 

diff --git a/docs/source/nlp/api.rst b/docs/source/nlp/api.rst
@@ -140,3 +140,13 @@ Datasets
 .. autoclass:: nemo.collections.nlp.data.language_modeling.megatron.ul2_dataset.UL2Dataset
     :show-inheritance: 
 
+Exportable Model Classes
+-------------------------
+
+.. autoclass:: nemo.collections.nlp.models.language_modeling.megatron_gpt_model.MegatronGPTExportableModel
+    :show-inheritance: 
+
+.. toctree::
+   :maxdepth: 1
+
+   megatron_onnx_export
diff --git a/docs/source/nlp/megatron_onnx_export.rst b/docs/source/nlp/megatron_onnx_export.rst
@@ -0,0 +1,47 @@
+.. _megatron_onnx_export:
+
+ONNX Export of Megatron Models
+====================================
+
+This guide demonstrates the usage of the ONNX export functionality for Megatron models.
+
+Requirements
+-----------------
+Set up the development environment by launching the latest `NeMo container <https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo/tags>`_
+
+The minimum version requirements for NeMo and TransformerEngine are below
+
+.. code-block:: bash 
+
+    nemo > 1.19
+    transformer_engine > 0.10
+
+Export to ONNX
+-----------------
+The export script supports the ONNX export of models with .nemo and .ckpt file extensions. The script also supports the export of the following types of models: GPT, T5, BERT, BART, NMT, RETRO.
+Commands for both file formats are discussed in the following sections. The model type used for the examples is GPT.
+
+
+Export using .nemo file
+^^^^^^^^^^^^^^^^^^^^^^^^
+A model with .nemo file extension can be exported using the command below
+
+.. code-block:: bash 
+
+    python3 examples/nlp/language_modeling/megatron_export.py \
+        model_type=gpt \
+        onnx_model_file=gpt_126m.onnx \
+        gpt_model_file=gpt_126m.nemo
+
+Export using .ckpt file
+^^^^^^^^^^^^^^^^^^^^^^^^
+A model with .ckpt file extension can be exported using the command below
+
+.. code-block:: bash 
+
+    python3 examples/nlp/language_modeling/megatron_export.py \
+        model_type=gpt \
+        onnx_model_file=gpt_126m.onnx \
+        checkpoint_dir=./gpt_126m/ \
+        checkpoint_name=model_weights.ckpt \
+        hparams_file=./gpt_126m/hparams.yaml