Remove nemo from in-framework deployment

Signed-off-by: Onur Yilmaz <oyilmaz@nvidia.com>

Author: oyilmaz-nvidia

.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/**/*
 

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/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,28 +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,
@@ -194,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,
         )

nemo_deploy/llm/megatronllm_deployable_ray.py

@@ -28,7 +28,7 @@
 from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
 
 from ..ray_utils import find_available_port
-from .megatronllm_deployable import MegatronLLMDeployableNemo2
+from .megatronllm_deployable import MegatronLLMDeployable
 
 LOGGER = logging.getLogger("NeMo")
 
@@ -44,7 +44,7 @@ class ModelWorker:
 
     def __init__(
         self,
-        nemo_checkpoint_filepath: str,
+        megatron_checkpoint_filepath: str,
         rank: int,
         world_size: int,
         tensor_model_parallel_size: int,
@@ -59,9 +59,7 @@ def __init__(
         legacy_ckpt: bool = False,
         max_batch_size: int = 32,
         random_seed: Optional[int] = None,
-        megatron_checkpoint_filepath: str = None,
         model_type: str = "gpt",
-        model_format: str = "nemo",
         micro_batch_size: Optional[int] = None,
         **model_config_kwargs,
     ):
@@ -83,8 +81,8 @@ def __init__(
             LOGGER.info(f"Replica {replica_id} - MASTER_ADDR: {os.environ['MASTER_ADDR']}")
 
         try:
-            self.model = MegatronLLMDeployableNemo2(
-                nemo_checkpoint_filepath=nemo_checkpoint_filepath,
+            self.model = MegatronLLMDeployable(
+                megatron_checkpoint_filepath=megatron_checkpoint_filepath,
                 num_devices=world_size,
                 num_nodes=world_size // torch.cuda.device_count(),
                 tensor_model_parallel_size=tensor_model_parallel_size,
@@ -96,9 +94,7 @@ def __init__(
                 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,
             )
@@ -128,28 +124,26 @@ class MegatronRayDeployable:
 
     def __init__(
         self,
-        nemo_checkpoint_filepath: str,
+        megatron_checkpoint_filepath: str,
         num_gpus: int = 1,
         tensor_model_parallel_size: int = 1,
         pipeline_model_parallel_size: int = 1,
         context_parallel_size: int = 1,
         expert_model_parallel_size: int = 1,
-        model_id: str = "nemo-model",
+        model_id: str = "megatron-model",
         enable_cuda_graphs: bool = False,
         enable_flash_decode: bool = False,
         legacy_ckpt: bool = False,
         max_batch_size: int = 32,
         random_seed: Optional[int] = None,
-        megatron_checkpoint_filepath: str = None,
         model_type: str = "gpt",
-        model_format: str = "nemo",
         micro_batch_size: Optional[int] = None,
         **model_config_kwargs,
     ):
         """Initialize the distributed Megatron LLM model deployment.
 
         Args:
-            nemo_checkpoint_filepath (str): Path to the .nemo checkpoint file.
+            megatron_checkpoint_filepath (str): Path to the Megatron checkpoint directory.
             num_gpus (int): Number of GPUs to use for the deployment
             tensor_model_parallel_size (int): Size of tensor model parallelism.
             pipeline_model_parallel_size (int): Size of pipeline model parallelism.
@@ -194,7 +188,7 @@ def __init__(
 
             # Common arguments for rank 0 worker
             rank_0_kwargs = dict(
-                nemo_checkpoint_filepath=nemo_checkpoint_filepath,
+                megatron_checkpoint_filepath=megatron_checkpoint_filepath,
                 rank=0,
                 world_size=num_gpus,
                 tensor_model_parallel_size=tensor_model_parallel_size,
@@ -209,9 +203,7 @@ def __init__(
                 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,
             )
@@ -233,7 +225,7 @@ def __init__(
             # Create remaining workers in parallel
             for rank in range(1, num_gpus):
                 worker = ModelWorker.remote(
-                    nemo_checkpoint_filepath=nemo_checkpoint_filepath,
+                    megatron_checkpoint_filepath=megatron_checkpoint_filepath,
                     rank=rank,
                     world_size=num_gpus,
                     tensor_model_parallel_size=tensor_model_parallel_size,
@@ -247,9 +239,7 @@ def __init__(
                     enable_flash_decode=enable_flash_decode,
                     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/utils.py

@@ -12,14 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 import typing
-from pathlib import Path
 
 import numpy as np
 import torch
 
-from nemo_export.tarutils import TarPath
 from nemo_export_deploy_common.import_utils import MISSING_PIL_MSG, MISSING_TRITON_MSG, UnavailableError
 
 try:
@@ -42,10 +39,6 @@
     HAVE_TRITON = False
 
 
-NEMO2 = "NEMO 2.0"
-NEMO1 = "NEMO 1.0"
-
-
 def typedict2tensor(
     typedict_class,
     overwrite_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None,
@@ -98,29 +91,6 @@ def _get_tensor_params(type_):
     )
 
 
-def nemo_checkpoint_version(path: str) -> str:
-    """Determines the version of a NeMo checkpoint from its file structure.
-
-    Examines the provided checkpoint path to determine if it follows the NeMo 2.0
-    or NeMo 1.0 format based on the presence of 'context' and 'weights' directories.
-
-    Args:
-        path (str): Path to the NeMo checkpoint file or directory
-
-    Returns:
-        str: Version string - either NEMO2 or NEMO1 constant indicating the checkpoint version
-    """
-    if os.path.isdir(path):
-        path = Path(path)
-    else:
-        path = TarPath(path)
-
-    if (path / "context").exists() and (path / "weights").exists():
-        return NEMO2
-    else:
-        return NEMO1
-
-
 def str_list2numpy(str_list: typing.List[str]) -> np.ndarray:
     """Converts a list of strings to a numpy array of UTF-8 encoded bytes.