Merge branch 'main' into athitten/chat_endpoint

Author: athitten

.github/labeler.yml

@@ -33,7 +33,7 @@ TensorRT-LLM:
 - nemo_export/trt_llm/**/*
 - nemo_export/tensorrt_llm*.py
 - tests/functional_tests/tests_trtllm/**/*
-- scripts/export/export_to_trt_llm.py
+- scripts/export/export_hf_to_nemo2.py
 - docs/llm/**/optimized/**/*
 - docs/mm/**/optimized/**/*
 

.github/workflows/release-docs.yml

@@ -0,0 +1,74 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+name: Release docs
+on:
+  workflow_dispatch:
+    inputs:
+      dry-run:
+        description: Whether to run the workflow in dry-run mode
+        required: true
+        type: boolean
+        default: true
+      version-number:
+        description: Version number to release this as (use `latest` for main branch)
+        required: true
+        type: string
+      notify-emails:
+        description: Email addresses to send the notification to. Format as "me@me.com,you@you.com".
+        required: true
+        type: string
+      aws-region:
+        description: AWS region
+        required: false
+        type: string
+        default: us-east-1
+
+jobs:
+  build-docs:
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_build_docs.yml@v0.67.0
+
+  publish-docs:
+    runs-on: ubuntu-latest
+    needs: [build-docs]
+    steps:
+      - uses: actions/checkout@v6
+        with:
+          repository: NVIDIA-NeMo/FW-CI-templates
+          ref: v0.67.2
+          path: FW-CI-templates
+
+      - uses: ./FW-CI-templates/.github/actions/publish-docs
+        # This workflow runs either on main, or on a version tag. Any other git ref will lead
+        # to an error.
+        # If its on main, it will publish to "latest" directory in Akamai.
+        # If its on a versioned tag, it will extract the version number from the tag (strip `v` prefix)
+        # and publish to the versioned directory in Akamai.
+        with:
+          dry-run: ${{ inputs.dry-run }}
+          artifacts-name: docs-html
+          artifacts-path: _build/html
+          emails-csv: ${{ inputs.notify-emails && format('{0},{1}', vars.docs_release_emails, inputs.notify-emails) || vars.docs_release_emails }}
+          overwrite-latest-on-tag: false
+          run-on-version-tag-only: ${{ github.ref_name != 'main' }}
+          request-name: export-deploy-publish-docs-${{ github.run_id }}
+          aws-region: ${{ inputs.aws-region }}
+          aws-role-to-assume: ${{ secrets.AWS_ASSUME_ROLE_ARN }}
+          aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+          akamai-host: ${{ secrets.AKAMAI_HOST }}
+          akamai-client-token: ${{ secrets.AKAMAI_CLIENT_TOKEN }}
+          akamai-client-secret: ${{ secrets.AKAMAI_CLIENT_SECRET }}
+          akamai-access-token: ${{ secrets.AKAMAI_ACCESS_TOKEN }}
+          s3-target-root: ${{ secrets.S3_BUCKET_NAME }}
+          s3-target-path: nemo/export-deploy

README.md

@@ -242,13 +242,13 @@ In order to run examples with NeMo models, a NeMo checkpoint is required. Please
    huggingface-cli login
    ```
 
-4. Run the following Python code to generate the NeMo 2.0 checkpoint:
+4. Run the following Python code to generate the NeMo checkpoint:
 
-   ```shell
-   python scripts/export/export_hf_to_nemo2.py \
-    --hf_model meta-llama/Llama-3.2-1B \
-    --output_path /opt/checkpoints/hf_llama32_1B_nemo2 \
-    --config Llama32Config1B
+   ```python
+   from nemo.collections import llm
+   
+   # Example: Converting Hugging Face model to NeMo format
+   # See NeMo documentation for detailed instructions
    ```
 
 ## 🚀 Export and Deploy Examples
@@ -297,9 +297,9 @@ from nemo_export.tensorrt_llm import TensorRTLLM
 from nemo_deploy import DeployPyTriton
 
 # Export model to TensorRT-LLM
-exporter = TensorRTLLM(model_dir="/tmp/hf_llama32_1B_nemo2")
+exporter = TensorRTLLM(model_dir="/tmp/llama32_1B_nemo")
 exporter.export(
-    nemo_checkpoint_path="/opt/checkpoints/hf_llama32_1B_nemo2",
+    nemo_checkpoint_path="/opt/checkpoints/llama32_1B_nemo",
     tensor_parallelism_size=1,
 )
 
@@ -328,8 +328,8 @@ from nemo_deploy import DeployPyTriton
 # Export model to vLLM
 exporter = vLLMExporter()
 exporter.export(
-    nemo_checkpoint="/opt/checkpoints/hf_llama32_1B_nemo2",
-    model_dir="/tmp/hf_llama32_1B_nemo2",
+    nemo_checkpoint="/opt/checkpoints/llama32_1B_nemo",
+    model_dir="/tmp/llama32_1B_nemo",
     tensor_parallel_size=1,
 )
 
@@ -355,10 +355,10 @@ You can also deploy NeMo and Hugging Face models directly using Triton Inference
 
 ```python
 from nemo_deploy import DeployPyTriton
-from nemo_deploy.nlp.megatronllm_deployable import MegatronLLMDeployableNemo2
+from nemo_deploy.llm.megatronllm_deployable import MegatronLLMDeployable
 
-model = MegatronLLMDeployableNemo2(
-    nemo_checkpoint_filepath="/opt/checkpoints/hf_llama32_1B_nemo2",
+model = MegatronLLMDeployable(
+    nemo_checkpoint_filepath="/opt/checkpoints/llama32_1B_nemo",
     num_devices=1,
     num_nodes=1,
 )

nemo_deploy/deploy_ray.py

@@ -171,13 +171,13 @@ def _stop(self):
 
     def deploy_inframework_model(
         self,
-        nemo_checkpoint: str,
+        megatron_checkpoint: str,
         num_gpus: int = 1,
         tensor_model_parallel_size: int = 1,
         pipeline_model_parallel_size: int = 1,
         expert_model_parallel_size: int = 1,
         context_parallel_size: int = 1,
-        model_id: str = "nemo-model",
+        model_id: str = "megatron-model",
         num_cpus_per_replica: float = 8,
         num_replicas: int = 1,
         enable_cuda_graphs: bool = False,
@@ -186,13 +186,11 @@ def deploy_inframework_model(
         max_batch_size: int = 32,
         random_seed: Optional[int] = None,
         test_mode: bool = False,
-        megatron_checkpoint_filepath: str = None,
         model_type: str = "gpt",
-        model_format: str = "nemo",
         micro_batch_size: Optional[int] = None,
         **model_config_kwargs,
     ):
-        """Deploy an inframework NeMo/Megatron model using Ray Serve.
+        """Deploy an inframework Megatron model using Ray Serve.
 
         This method handles the complete deployment lifecycle including:
         - Starting Ray Serve
@@ -201,22 +199,20 @@ def deploy_inframework_model(
         - Keeping the deployment running until interrupted
 
         Args:
-            nemo_checkpoint (str): Path to the .nemo checkpoint file.
+            megatron_checkpoint (str): Path to the Megatron checkpoint directory.
             num_gpus (int, optional): Number of GPUs per node. Defaults to 1.
             tensor_model_parallel_size (int, optional): Tensor model parallel size. Defaults to 1.
             pipeline_model_parallel_size (int, optional): Pipeline model parallel size. Defaults to 1.
             expert_model_parallel_size (int, optional): Expert model parallel size. Defaults to 1.
             context_parallel_size (int, optional): Context parallel size. Defaults to 1.
-            model_id (str, optional): Model identifier for API responses. Defaults to "nemo-model".
+            model_id (str, optional): Model identifier for API responses. Defaults to "megatron-model".
             num_cpus_per_replica (float, optional): CPUs per model replica. Defaults to 8.
             num_replicas (int, optional): Number of replicas for deployment. Defaults to 1.
             enable_cuda_graphs (bool, optional): Enable CUDA graphs. Defaults to False.
             enable_flash_decode (bool, optional): Enable Flash Attention decode. Defaults to False.
             legacy_ckpt (bool, optional): Use legacy checkpoint format. Defaults to False.
             test_mode (bool, optional): Enable test mode. Defaults to False.
-            megatron_checkpoint_filepath (str, optional): Path to the Megatron checkpoint file. Defaults to None.
             model_type (str, optional): Type of model to load. Defaults to "gpt".
-            model_format (str, optional): Format of model to load. Defaults to "nemo".
             micro_batch_size (Optional[int], optional): Micro batch size for model execution. Defaults to None.
 
         Raises:
@@ -244,7 +240,7 @@ def deploy_inframework_model(
                 num_replicas=num_replicas,
                 ray_actor_options={"num_cpus": num_cpus_per_replica},
             ).bind(
-                nemo_checkpoint_filepath=nemo_checkpoint,
+                megatron_checkpoint_filepath=megatron_checkpoint,
                 num_gpus=gpus_per_replica,
                 tensor_model_parallel_size=tensor_model_parallel_size,
                 pipeline_model_parallel_size=pipeline_model_parallel_size,
@@ -256,9 +252,7 @@ def deploy_inframework_model(
                 legacy_ckpt=legacy_ckpt,
                 max_batch_size=max_batch_size,
                 random_seed=random_seed,
-                megatron_checkpoint_filepath=megatron_checkpoint_filepath,
                 model_type=model_type,
-                model_format=model_format,
                 micro_batch_size=micro_batch_size,
                 **model_config_kwargs,
             )

nemo_deploy/llm/megatronllm_deployable.py

@@ -27,10 +27,8 @@
 from nemo_deploy import ITritonDeployable
 from nemo_deploy.llm.inference.inference_base import create_mcore_engine
 from nemo_deploy.utils import (
-    NEMO2,
     broadcast_list,
     cast_output,
-    nemo_checkpoint_version,
     str_ndarray2list,
 )
 from nemo_export_deploy_common.import_utils import MISSING_TRITON_MSG, UnavailableError, null_decorator
@@ -54,73 +52,16 @@
 LOGGER = logging.getLogger("NeMo")
 
 
-class MegatronLLMDeploy:
-    """A factory class for creating deployable instances of Megatron LLM models.
-
-    This class provides a method to get the appropriate deployable instance
-    based on the version of the NeMo checkpoint model used.
-    """
-
-    @staticmethod
-    def get_deployable(
-        nemo_checkpoint_filepath: str,
-        num_devices: int = None,
-        num_nodes: int = None,
-        tensor_model_parallel_size: int = 1,
-        pipeline_model_parallel_size: int = 1,
-        expert_model_parallel_size: int = 1,
-        context_parallel_size: int = 1,
-        max_batch_size: int = 32,
-        random_seed: Optional[int] = None,
-        enable_flash_decode: bool = False,
-        enable_cuda_graphs: bool = False,
-        legacy_ckpt: bool = False,
-    ):
-        """Returns the appropriate deployable instance for the given NeMo checkpoint.
-
-        Args:
-            nemo_checkpoint_filepath (str): Path to the .nemo checkpoint file.
-            num_devices (int): Number of devices to use for deployment.
-            num_nodes (int): Number of nodes to use for deployment.
-            tensor_model_parallel_size (int): Size of the tensor model parallelism.
-            pipeline_model_parallel_size (int): Size of the pipeline model parallelism.
-            context_parallel_size (int): Size of the context parallelism.
-            enable_flash_decode (bool): Whether to enable flash decode for inference.
-            enable_cuda_graphs (bool): Whether to enable CUDA graphs for inference.
-            legacy_ckpt (bool): Whether to use legacy checkpoint format. Defaults to False.
-
-        Returns:
-            ITritonDeployable: An instance of a deployable class compatible with Triton inference server.
-        """
-        if nemo_checkpoint_version(nemo_checkpoint_filepath) == NEMO2:
-            return MegatronLLMDeployableNemo2(
-                num_devices=num_devices,
-                num_nodes=num_nodes,
-                nemo_checkpoint_filepath=nemo_checkpoint_filepath,
-                tensor_model_parallel_size=tensor_model_parallel_size,
-                pipeline_model_parallel_size=pipeline_model_parallel_size,
-                context_parallel_size=context_parallel_size,
-                expert_model_parallel_size=expert_model_parallel_size,
-                max_batch_size=max_batch_size,
-                random_seed=random_seed,
-                enable_flash_decode=enable_flash_decode,
-                enable_cuda_graphs=enable_cuda_graphs,
-                legacy_ckpt=legacy_ckpt,
-            )
-        else:
-            raise Exception("Only NeMo 2.0 checkpoint is supported.")
-
-
 def dict_to_str(messages):
     """Serializes dict to str."""
     return json.dumps(messages)
 
 
-class MegatronLLMDeployableNemo2(ITritonDeployable):
-    """Triton inference server compatible deploy class for a .nemo model file.
+class MegatronLLMDeployable(ITritonDeployable):
+    """Triton inference server compatible deploy class for a Megatron model checkpoint.
 
     Args:
-        nemo_checkpoint_filepath (str): path for the nemo checkpoint.
+        megatron_checkpoint_filepath (str): path for the megatron checkpoint.
         num_devices (int): number of GPUs.
         num_nodes (int): number of nodes.
         tensor_model_parallel_size (int): tensor parallelism.
@@ -136,17 +77,15 @@ class MegatronLLMDeployableNemo2(ITritonDeployable):
         enable_flash_decode (bool): enable flash decode for inference. Defaults to False.
         enable_cuda_graphs (bool): enable CUDA graphs for inference. Defaults to False.`
         legacy_ckpt (bool): use legacy checkpoint format. Defaults to False.
-        megatron_checkpoint_filepath (str): path for the megatron checkpoint.
-        model_type (str): type of model to load. Defaults to "gpt".(Only for Megatron models)
-        model_format (str): format of model to load. Defaults to "nemo".
-        micro_batch_size (Optional[int]): micro batch size for model execution. Defaults to None.(Only for Megatron models)
+        model_type (str): type of model to load. Defaults to "gpt".
+        micro_batch_size (Optional[int]): micro batch size for model execution. Defaults to None.
     """
 
     def __init__(
         self,
+        megatron_checkpoint_filepath: str,
         num_devices: int = None,
         num_nodes: int = None,
-        nemo_checkpoint_filepath: str = None,
         tensor_model_parallel_size: int = 1,
         pipeline_model_parallel_size: int = 1,
         context_parallel_size: int = 1,
@@ -159,27 +98,20 @@ def __init__(
         max_batch_size: int = 8,
         random_seed: Optional[int] = None,
         legacy_ckpt: bool = False,
-        megatron_checkpoint_filepath: str = None,
         model_type: str = "gpt",
-        model_format: str = "nemo",
         micro_batch_size: Optional[int] = None,
         **model_config_kwargs,
     ):
         if not HAVE_TRITON:
             raise UnavailableError(MISSING_TRITON_MSG)
-        if model_format == "nemo":
-            checkpoint_filepath = nemo_checkpoint_filepath
-        elif model_format == "megatron":
-            if model_type not in ["gpt", "mamba"]:
-                raise ValueError(f"Model type {model_type} not supported for Megatron models.")
-            checkpoint_filepath = megatron_checkpoint_filepath
-        else:
-            raise ValueError(f"Model format {model_format} not supported.")
+
+        if model_type not in ["gpt", "mamba"]:
+            raise ValueError(f"Model type {model_type} not supported for Megatron models.")
 
         self.mcore_engine, self.inference_wrapped_model, self.mcore_tokenizer = create_mcore_engine(
             num_devices=num_devices,
             num_nodes=num_nodes,
-            path=Path(checkpoint_filepath),
+            path=Path(megatron_checkpoint_filepath),
             params_dtype=params_dtype,
             inference_batch_times_seqlen_threshold=inference_batch_times_seqlen_threshold,
             inference_max_seq_length=inference_max_seq_length,
@@ -193,7 +125,7 @@ def __init__(
             enable_cuda_graphs=enable_cuda_graphs,
             legacy_ckpt=legacy_ckpt,
             model_type=model_type,
-            model_format=model_format,
+            model_format="megatron",
             micro_batch_size=micro_batch_size,
             **model_config_kwargs,
         )