Remove more nemo2 and unused code. (#584)

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

Author: oyilmaz-nvidia

.github/labeler.yml

@@ -1,6 +1,5 @@
 export:
 - nemo_export/**/*
-- scripts/export/**/*
 - tests/functional_tests/tests_trtllm/test_export.py
 - tests/functional_tests/tests_onnx_trt/test_export.py
 - tests/unit_tests/export/**/*
@@ -32,7 +31,6 @@ TensorRT-LLM:
 - nemo_export/trt_llm/**/*
 - nemo_export/tensorrt_llm*.py
 - tests/functional_tests/tests_trtllm/**/*
-- scripts/export/export_hf_to_nemo2.py
 - docs/llm/**/optimized/**/*
 - docs/mm/**/optimized/**/*
 

nemo_export/onnx_llm_exporter.py

@@ -26,7 +26,6 @@
 from nemo_export.utils import (
     get_example_inputs,
     get_model_device_type,
-    is_nemo2_checkpoint,
     validate_fp8_network,
 )
 from nemo_export_deploy_common.import_utils import (
@@ -157,10 +156,9 @@ def __init__(
                 raise ValueError("A model was also passed but it will be overridden.")
 
             if Path(self.model_name_or_path).is_dir():
-                if is_nemo2_checkpoint(self.model_name_or_path):
-                    raise NotImplementedError("NeMo 2.0 checkpoint will be supported later.")
-                else:
-                    self._load_hf_model()
+                self._load_hf_model()
+            else:
+                raise ValueError("The model_name_or_path is not a valid directory.")
 
         self.model.to(self.device)
 

nemo_export/tensorrt_llm.py

@@ -49,10 +49,7 @@
     unload_engine,
 )
 from nemo_export.trt_llm.utils import determine_quantization_settings, is_rank
-from nemo_export.utils import (
-    is_nemo2_checkpoint,
-    prepare_directory_for_export,
-)
+from nemo_export.utils import prepare_directory_for_export
 from nemo_export.utils.constants import TRTLLM_ENGINE_DIR
 from nemo_export_deploy_common.import_utils import (
     MISSING_TENSORRT_LLM_MSG,
@@ -704,7 +701,7 @@ def export(
                 )
 
             # Check if it's a NeMo2 checkpoint
-            if not (Path(nemo_checkpoint_path).exists() and is_nemo2_checkpoint(nemo_checkpoint_path)):
+            if not (Path(nemo_checkpoint_path).exists() and (Path(nemo_checkpoint_path) / "context").is_dir()):
                 raise Exception(
                     f"Standard NeMo export failed and checkpoint is not a NeMo2 checkpoint. "
                     f"HF fallback only works with NeMo2 checkpoints. "

nemo_export/utils/__init__.py

@@ -22,8 +22,6 @@
 from nemo_export.utils.utils import (
     get_example_inputs,
     get_model_device_type,
-    is_nemo2_checkpoint,
-    is_nemo_tarfile,
     prepare_directory_for_export,
     torch_dtype_from_precision,
     validate_fp8_network,
@@ -35,8 +33,6 @@
     "load_sharded_metadata_torch_dist",
     "load_sharded_metadata_zarr",
     "nemo_to_path",
-    "is_nemo2_checkpoint",
-    "is_nemo_tarfile",
     "prepare_directory_for_export",
     "torch_dtype_from_precision",
     "get_model_device_type",

nemo_export/utils/utils.py

@@ -21,19 +21,6 @@
 from transformers import PreTrainedTokenizerBase
 
 
-def is_nemo2_checkpoint(checkpoint_path: str) -> bool:
-    """Checks if the checkpoint is in NeMo 2.0 format.
-
-    Args:
-        checkpoint_path (str): Path to a checkpoint.
-
-    Returns:
-        bool: True if the path points to a NeMo 2.0 checkpoint; otherwise false.
-    """
-    ckpt_path = Path(checkpoint_path)
-    return (ckpt_path / "context").is_dir()
-
-
 def prepare_directory_for_export(
     model_dir: Union[str, Path],
     delete_existing_files: bool,
@@ -64,19 +51,6 @@ def prepare_directory_for_export(
     model_path.mkdir(parents=True, exist_ok=True)
 
 
-def is_nemo_tarfile(path: str) -> bool:
-    """Checks if the path exists and points to packed NeMo 1 checkpoint.
-
-    Args:
-        path (str): Path to possible checkpoint.
-
-    Returns:
-        bool: NeMo 1 checkpoint exists and is in '.nemo' format.
-    """
-    checkpoint_path = Path(path)
-    return checkpoint_path.exists() and checkpoint_path.suffix == ".nemo"
-
-
 # Copied from nemo.collections.nlp.parts.utils_funcs to avoid introducing extra NeMo dependencies:
 def torch_dtype_from_precision(precision: Union[int, str], megatron_amp_O2: bool = True) -> torch.dtype:
     """Mapping from PyTorch Lighthing (PTL) precision types to corresponding PyTorch parameter data type.

scripts/export/export_hf_to_nemo2.py

@@ -31,21 +31,6 @@ def temp_dir(self):
         # Cleanup after test
         shutil.rmtree(temp_dir)
 
-    @pytest.mark.run_only_on("GPU")
-    def test_is_nemo2_checkpoint(self, temp_dir):
-        from nemo_export.utils.utils import is_nemo2_checkpoint
-
-        # Test with non-existent path
-        assert not is_nemo2_checkpoint("/non/existent/path")
-
-        # Test with directory without context folder
-        os.makedirs(os.path.join(temp_dir, "no_context"))
-        assert not is_nemo2_checkpoint(os.path.join(temp_dir, "no_context"))
-
-        # Test with valid NeMo 2.0 checkpoint
-        os.makedirs(os.path.join(temp_dir, "valid_ckpt", "context"))
-        assert is_nemo2_checkpoint(os.path.join(temp_dir, "valid_ckpt"))
-
     @pytest.mark.run_only_on("GPU")
     def test_prepare_directory_for_export(self, temp_dir):
         from nemo_export.utils.utils import prepare_directory_for_export
@@ -76,25 +61,6 @@ def test_prepare_directory_for_export(self, temp_dir):
         prepare_directory_for_export(model_dir, delete_existing_files=False, subdir="subdir")
         assert os.path.exists(os.path.join(model_dir, "subdir"))
 
-    @pytest.mark.run_only_on("GPU")
-    def test_is_nemo_tarfile(self, temp_dir):
-        from nemo_export.utils.utils import is_nemo_tarfile
-
-        # Test with non-existent file
-        assert not is_nemo_tarfile("/non/existent/file.nemo")
-
-        # Test with non-nemo file
-        test_file = os.path.join(temp_dir, "test.txt")
-        with open(test_file, "w") as f:
-            f.write("test")
-        assert not is_nemo_tarfile(test_file)
-
-        # Test with .nemo file
-        nemo_file = os.path.join(temp_dir, "test.nemo")
-        with open(nemo_file, "w") as f:
-            f.write("test")
-        assert is_nemo_tarfile(nemo_file)
-
     @pytest.mark.run_only_on("GPU")
     def test_torch_dtype_from_precision(self):
         from nemo_export.utils.utils import torch_dtype_from_precision