nvidia-modelopt 0.27.0__py3-none-win_amd64.whl

modelopt/__init__.py

@@ -0,0 +1,42 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Nvidia TensorRT Model Optimizer (modelopt)."""
+
+import sys as _sys
+import warnings as _warnings
+from importlib.metadata import version as _version
+
+__version__ = _version("nvidia-modelopt")
+
+
+try:
+    # Import from local source if available
+    from . import core
+
+    __core_version__ = __version__
+except ImportError:
+    # Import from nvidia-modelopt-core wheel installation
+    import modelopt_core as core  # type: ignore[no-redef]
+
+    _sys.modules["modelopt.core"] = core
+    __core_version__ = _version("nvidia-modelopt-core")
+
+# Versions need to be the same for compatibility
+if __version__.split(".")[:2] != __core_version__.split(".")[:2]:
+    _warnings.warn(
+        f"Version mismatch between nvidia-modelopt ({__version__}) and nvidia-modelopt-core"
+        f" ({__core_version__}). Please ensure both versions are the same for compatibility."
+    )

modelopt/deploy/__init__.py

@@ -0,0 +1,16 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Model Optimizer's deployment package."""

modelopt/deploy/llm/__init__.py

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""LLM deployment utils with tensorrt_llm."""
+
+from mpi4py import MPI
+
+# Pre load MPI libs to avoid tensorrt_llm importing failures.
+print(f"Loaded mpi lib {MPI.__file__} successfully")
+
+# Pre import tensorrt_llm
+try:
+    import tensorrt_llm
+except Exception as e:
+    print(
+        "tensorrt_llm package is not installed. Please build or install tensorrt_llm package"
+        " properly before calling the llm deployment API."
+    )
+    raise (e)
+
+from .generate import *  # noqa

modelopt/deploy/llm/generate.py

@@ -0,0 +1,297 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""A wrapper over the TensorRT-LLM high level API runner."""
+
+import json
+from pathlib import Path
+from typing import Any, Iterable, Optional, Union
+
+import tensorrt_llm
+import torch
+from packaging.version import parse
+from tensorrt_llm import SamplingParams
+from tensorrt_llm._torch.pyexecutor.config import PyTorchConfig
+from tensorrt_llm.bindings.executor import DecodingConfig
+from tensorrt_llm.llmapi import KvCacheConfig as TRT_KvCacheConfig
+from tensorrt_llm.llmapi.llm import LLM as TRT_LLM
+from tensorrt_llm.llmapi.tokenizer import TokenizerBase, TransformersTokenizer
+
+
+def _sanitize_temperature_and_top_p(temperature, top_p):
+    assert temperature >= 0.0, "Temperature must be greater than 0.0."
+
+    # TRT LLM acccepts temperature values only greater than 0.0
+    temperature = max(temperature, 0.001)
+
+    kwargs = {"temperature": temperature}
+    if top_p is not None:
+        # cpp executor only supports topP.value() > 0.f
+        top_p = 1e-4 if top_p == 0 else top_p
+        kwargs["top_p"] = top_p
+
+    return kwargs
+
+
+class LLM(TRT_LLM):
+    """A wrapper over the ``tensorrt_llm.llmapi.llm.LLM`` for LLM profiling and validation."""
+
+    def _build_trt_llm_from_config(
+        self, config, engine_dir, tokenizer, kv_cache_config, medusa_choices
+    ):
+        build_config = config["build_config"]
+        world_size = config.get("pretrained_config", {}).get("mapping", {}).get("world_size", 1)
+        max_tokens_kv_cache = build_config["max_seq_len"] * build_config["max_batch_size"]
+
+        trt_kv_cache_config = TRT_KvCacheConfig(enable_block_reuse=False)
+
+        # If not specified, free_gpu_memory_fraction is set to the default TRT LLM value 0.9
+        trt_kv_cache_config.free_gpu_memory_fraction = kv_cache_config.get(
+            "free_gpu_memory_fraction", 0.9
+        )
+
+        # If not specified, max_tokens is set to the max value calculated above.
+        trt_kv_cache_config.max_tokens = kv_cache_config.get("max_tokens", max_tokens_kv_cache)
+
+        kwargs = {}
+        if medusa_choices is not None:
+            decoding_config = DecodingConfig()
+            decoding_config.medusa_choices = medusa_choices
+            kwargs["decoding_config"] = decoding_config
+            assert world_size == 1, "decoding_config does not support multi TP in HLAPI."
+
+        if tokenizer is None:
+            # Assume the tokenizer is stored in the engine_dir if not specified.
+            tokenizer = engine_dir
+
+        # CustomSentencePieceTokenizer will not be recognized by llmapi, wrapping it around TransformersTokenizer
+        if type(tokenizer).__name__ in ["CustomSentencePieceTokenizer"]:
+            tokenizer = TransformersTokenizer(tokenizer)
+
+        super().__init__(
+            model=engine_dir,
+            tokenizer=tokenizer,
+            kv_cache_config=trt_kv_cache_config,
+            **kwargs,
+        )
+
+    def _build_torch_llm_from_config(self, checkpoint_dir, tokenizer, tp, trust_remote_code):
+        pytorch_config = PyTorchConfig(use_cuda_graph=True)
+
+        kwargs = {}
+        if tokenizer is not None:
+            kwargs["tokenizer"] = tokenizer
+
+        if tp < 1:
+            tp = torch.cuda.device_count()
+
+        # Sometimes 90% of the GPU memory is not enough for the TRT LLM torch engine.
+        trt_kv_cache_config = TRT_KvCacheConfig(
+            enable_block_reuse=False, free_gpu_memory_fraction=0.85
+        )
+
+        super().__init__(
+            backend="pytorch",
+            model=checkpoint_dir,
+            tensor_parallel_size=tp,
+            trust_remote_code=trust_remote_code,
+            enable_chunked_prefill=True,
+            pytorch_backend_config=pytorch_config,
+            kv_cache_config=trt_kv_cache_config,
+            **kwargs,
+        )
+
+    def __init__(
+        self,
+        checkpoint_dir: Union[str, Path],
+        tokenizer: Optional[Union[str, Path, TokenizerBase]] = None,
+        kv_cache_config: dict[str, Union[int, float]] = {},
+        medusa_choices: Any = None,
+        tp: int = 0,
+        trust_remote_code: bool = False,
+    ):
+        """Initializes the LLM runner class.
+
+        Args:
+            engine_dir: the directory path of the TensorRT-LLM engine.
+            tokenizer: the tokenizer. For example, a tokenizer from the Huggingface model.
+            kv_cache_config: the kv cache config as a dict. Please refer to
+                https://nvidia.github.io/TensorRT-LLM/performance/performance-tuning-guide/
+            medusa_choices: The medusa choices for the decoding config.
+            tp: the tensor parallel size (for the torch backend). If 0, it will be set to the number of GPUs.
+            trust_remote_code: whether to trust the remote code (for the torch backend).
+        """
+        assert parse(tensorrt_llm.__version__) >= parse("0.17.0")
+
+        with open(Path(checkpoint_dir) / "config.json", "r") as config_file:
+            config = json.load(config_file)
+
+            if "build_config" in config:
+                self._build_trt_llm_from_config(
+                    config, checkpoint_dir, tokenizer, kv_cache_config, medusa_choices
+                )
+
+                self._is_torch = False
+                self._max_seq_len = self.args.build_config.max_seq_len
+                self._max_beam_width = self.args.build_config.max_beam_width
+                self._gather_context_logits = self.args.build_config.gather_context_logits
+            else:
+                assert medusa_choices is None, (
+                    "medusa_choices is not supported with the torch llmapi"
+                )
+
+                self._build_torch_llm_from_config(checkpoint_dir, tokenizer, tp, trust_remote_code)
+                self._is_torch = True
+                self._max_seq_len = config["max_position_embeddings"]
+                self._max_beam_width = 1
+                self._gather_context_logits = False
+
+    @property
+    def max_seq_len(self):
+        """Get the max sequence length from the LLM instance."""
+        return self._max_seq_len
+
+    @property
+    def max_beam_width(self):
+        """Get the max beam width from the LLM instance."""
+        return self._max_beam_width
+
+    @property
+    def gather_context_logits(self):
+        """Returns whether the context_logits can be returned from the LLM instance."""
+        return self._gather_context_logits
+
+    def _generate(
+        self,
+        prompts: Union[Iterable[str], Iterable[list[int]]],
+        max_new_tokens: int,
+        temperature: float = 1.0,
+        top_p: float = None,
+        stop_words: list[str] = None,
+    ):
+        assert temperature >= 0.0, "Temperature must be greater than 0.0."
+
+        # TODO: Remove this once torch backend supports stop words
+        if self._is_torch:
+            stop_words = None
+
+        beam_width = self.max_beam_width
+        kwargs = _sanitize_temperature_and_top_p(temperature, top_p)
+        sampling_config = SamplingParams(
+            max_tokens=max_new_tokens, beam_width=beam_width, stop=stop_words, **kwargs
+        )
+
+        return self.generate(prompts, sampling_params=sampling_config, use_tqdm=False)
+
+    def generate_tokens(
+        self,
+        prompts: Union[Iterable[str], Iterable[list[int]]],
+        max_new_tokens: int,
+        temperature: float = 1.0,
+        top_p: float = None,
+        stop_words: list[str] = None,
+    ) -> Union[list[list[int]], list[list[list[int]]]]:
+        """Generates the tokens based on the input prompts.
+
+        Args:
+            prompts: The input prompts. Could be a list of strings or token lists.
+            max_new_tokens: The max output token length.
+            temperature: The sampling temperature.
+            top_p: The nucleus sampling parameter.
+            stop_words: A list of words that the generate stops on.
+
+        Returns:
+            a list of output token lists if max_beam_width is 1 or a 3D list with shape [batch, beam, sequence_len].
+        """
+        outputs = self._generate(prompts, max_new_tokens, temperature, top_p, stop_words)
+        output_tokens = []
+
+        beam_width = self.max_beam_width
+        for request_output in outputs:
+            token_ids = [
+                completion_output.token_ids for completion_output in request_output.outputs
+            ]
+            if beam_width == 1:
+                # each output_text is a single list of tokens.
+                output_tokens += token_ids
+            else:
+                # each output_text is a list of beam searched token lists.
+                output_tokens.append(token_ids)
+
+        return output_tokens
+
+    def generate_text(
+        self,
+        prompts: Union[Iterable[str], Iterable[list[int]]],
+        max_new_tokens: int,
+        temperature: float = 1.0,
+        top_p: float = None,
+        stop_words: list[str] = None,
+    ) -> Union[list[str], list[list[str]]]:
+        """Generates the text based on the input prompts.
+
+        Args:
+            prompts: The input prompts. Could be a list of strings or token lists.
+            max_new_tokens: The max output token length.
+            temperature: The sampling temperature
+            stop_words: A list of words the generate will stop on.
+
+        Returns:
+            a list of output text strings if max_beam_width is 1 or a 2D list with shape [batch, beam].
+        """
+        outputs = self._generate(prompts, max_new_tokens, temperature, top_p, stop_words)
+        output_texts = []
+
+        beam_width = self.max_beam_width
+        for request_output in outputs:
+            texts = [completion_output.text for completion_output in request_output.outputs]
+            if beam_width == 1:
+                # each output_text is a single text string.
+                output_texts += texts
+            else:
+                # each output_text is a list of beam searched texts
+                output_texts.append(texts)
+
+        return output_texts
+
+    def generate_context_logits(
+        self,
+        prompts: Union[Iterable[str], Iterable[list[int]]],
+        temperature: float = 1.0,
+        top_p: float = None,
+    ) -> list[torch.tensor]:
+        """Generates the context logits based on the input prompts.
+
+        Args:
+            prompts: The input prompts. Could be a list of strings or token lists.
+            temperature: The sampling temperature.
+            top_p: The nucleus sampling parameter.
+
+        Returns:
+            a tensor list of the context_logits.
+        """
+        assert self.gather_context_logits, (
+            "Please enable gather_context_logits flag when building the engine."
+        )
+        assert temperature >= 0.0, "Temperature must be greater than 0.0."
+
+        kwargs = _sanitize_temperature_and_top_p(temperature, top_p)
+        kwargs["return_context_logits"] = True
+
+        sampling_config = SamplingParams(max_tokens=1, beam_width=1, **kwargs)
+
+        outputs = self.generate(prompts, sampling_params=sampling_config, use_tqdm=False)
+
+        return [output.context_logits for output in outputs]

modelopt/deploy/llm/nemo_utils.py

@@ -0,0 +1,184 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""The utils to support Nemo models."""
+
+import warnings
+from pathlib import Path
+
+import numpy as np
+import torch
+from transformers import GPT2Tokenizer, PreTrainedTokenizer, T5Tokenizer
+
+try:
+    from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer
+
+    sentence_piece_tokenizer_available = True
+except Exception:
+    sentence_piece_tokenizer_available = False
+    SentencePieceTokenizer = PreTrainedTokenizer
+
+
+class CustomSentencePieceTokenizer(SentencePieceTokenizer):
+    """Custom tokenizer based on Nemo SentencePieceTokenizer.
+
+    This extension of SentencePieceTokenizer is to make API consistent with HuggingFace tokenizers
+    in order to run evaluation tools in examples/tensorrt_llm/scripts/nemo_example.sh script.
+    """
+
+    def __init__(self, *args, **kwargs):
+        """Constructor method with extra check for non-legacy SentencePieceTokenizer variant."""
+        super().__init__(*args, **kwargs)
+        assert not self.legacy, "Only non-legacy tokenizer is supported"
+
+    @property
+    def pad_token(self):
+        """pad_token."""
+        return self.pad_id
+
+    @property
+    def eos_token(self):
+        """eos_token."""
+        return self.eos_id
+
+    @property
+    def pad_token_id(self):
+        """pad_token_id."""
+        return self.pad_id
+
+    @property
+    def eos_token_id(self):
+        """eos_token_id."""
+        return self.eos_id
+
+    def encode(self, text, return_tensors=None, max_length=None, **kwargs):
+        """Method introduced for HF tokenizers API consistency for evaluation scripts.
+
+        Note: kwargs other than return_tensors and max_length are ignored.
+        """
+        output = self.tokenizer.encode_as_ids(text)
+        if max_length is not None:
+            if isinstance(text, str):
+                output = output[:max_length]
+            if isinstance(text, list):
+                output = [x[:max_length] for x in output]
+        if return_tensors == "pt":
+            # Only plain text input is supported since for list of strings some padding needs to be introduced
+            assert isinstance(text, str), (
+                "Returning 'pt' tensors is only supported for simple text input"
+            )
+            output = torch.LongTensor(output).reshape((1, -1))
+        return output
+
+    def batch_encode_plus(self, texts, **kwargs):
+        """Method introduced for HF tokenizers API consistency for evaluation scripts.
+
+        Note: kwargs are ignored.
+        """
+        assert isinstance(texts, list), f"Expected list of texts got {texts}"
+        return {"input_ids": self.tokenizer.encode_as_ids(texts)}
+
+    def decode(self, ids, **kwargs):
+        """MMethod introduced for HF tokenizers API consistency for evaluation scripts.
+
+        Note: kwargs are ignored.
+        """
+        if isinstance(ids, np.ndarray) or torch.is_tensor(ids):
+            ids = ids.tolist()
+        return self.tokenizer.decode(ids)
+
+    def batch_decode(self, ids, **kwargs):
+        """Method introduced for HF tokenizers API consistency for evaluation scripts."""
+        return self.decode(ids, **kwargs)
+
+
+def _build_tokenizer(tokenizer_config: dict):
+    if tokenizer_config["library"] == "sentencepiece":
+        # Model Optimizer modification.
+        # Turn off legacy model by default: See https://github.com/huggingface/transformers/pull/24622
+        tokenizer = T5Tokenizer(tokenizer_config["model"], extra_ids=0, legacy=False)
+    elif "GPT2" in tokenizer_config["type"]:
+        tokenizer = GPT2Tokenizer(tokenizer_config["vocab_file"], tokenizer_config["merge_file"])
+    else:
+        raise ValueError(f"Tokenizer type {tokenizer_config['library']} not handled")
+
+    if tokenizer.bos_token_id is None:
+        tokenizer.add_special_tokens({"bos_token": "<s>"})
+    if tokenizer.eos_token_id is None:
+        tokenizer.add_special_tokens({"eos_token": "</s>"})
+
+    return tokenizer
+
+
+def get_tokenzier(tokenizer_dir_or_path: Path) -> PreTrainedTokenizer:
+    """Loads the tokenizer from the decoded NEMO weights dir."""
+    # TODO: Remove this function sometime in the future
+    warnings.warn(
+        (
+            "Function get_tokenzier is deprecated and may be removed soon. "
+            "Please consider using get_nemo_tokenizer instead."
+        ),
+        DeprecationWarning,
+    )
+    model_path = (
+        tokenizer_dir_or_path / "tokenizer.model"
+        if tokenizer_dir_or_path.is_dir()
+        else tokenizer_dir_or_path
+    )
+    tokenizer_config = {"library": "sentencepiece", "model": str(model_path)}
+    return _build_tokenizer(tokenizer_config)
+
+
+def get_nemo_tokenizer(tokenizer_cfg_path: str):
+    """Build tokenizer from Nemo tokenizer config.
+
+    Refer to the logic of get_nmt_tokenizer function on how to instantiate tokenizers in Nemo, see
+    https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/modules/common/tokenizer_utils.py.
+    """
+    from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer
+    from omegaconf import OmegaConf
+
+    print(f"Initializing tokenizer from tokenizer config {tokenizer_cfg_path}")
+    tokenizer_cfg = OmegaConf.load(tokenizer_cfg_path)
+
+    library = tokenizer_cfg.library
+    legacy = tokenizer_cfg.get("sentencepiece_legacy", library == "sentencepiece")
+    special_tokens_dict = tokenizer_cfg.get("special_tokens", {})
+
+    if library == "huggingface":
+        print(f"Getting HuggingFace AutoTokenizer with pretrained_model_name: {tokenizer_cfg.type}")
+        tokenizer = AutoTokenizer(
+            pretrained_model_name=tokenizer_cfg.type,
+            vocab_file=tokenizer_cfg.get("vocab_file", None),
+            merges_file=tokenizer_cfg.get("merge_file", None),
+            use_fast=tokenizer_cfg.get("use_fast", False),
+            **special_tokens_dict,
+        )
+        print("Unwrapping HuggingFace tokenizer from Nemo AutoTokenizer")
+        # TODO: Either Nemo tokenizer API alignment is needed https://github.com/NVIDIA/NeMo/pull/8818
+        #  or alternatively we can use only HuggingFace tokenizers in future. Current workaround
+        #  proposed here is to unwrap the HuggingFace tokenizer from Nemo AutoTokenizer class:
+        tokenizer = tokenizer.tokenizer
+    elif library == "sentencepiece":
+        print(f"Getting SentencePiece with model: {tokenizer_cfg.model}")
+        if not sentence_piece_tokenizer_available:
+            warnings.warn("Cannot import nemo package, falling back to HF PreTrainedTokenizer!")
+        tokenizer = CustomSentencePieceTokenizer(model_path=tokenizer_cfg.model, legacy=legacy)
+    else:
+        raise NotImplementedError(
+            "Currently we only support 'huggingface' and 'sentencepiece' tokenizer libraries."
+        )
+
+    return tokenizer

modelopt/onnx/__init__.py

@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Model optimization subpackage for onnx."""
+
+import sys
+import warnings
+
+MIN_PYTHON_VERSION = (3, 10)
+
+try:
+    from . import quantization  # noqa: F401
+except ImportError as e:
+    raise ImportError(f"{e}\nPlease install optional ``[onnx]`` dependencies.")
+
+
+# Check the current Python version
+if sys.version_info < MIN_PYTHON_VERSION:
+    warnings.warn(
+        f"This package requires Python {MIN_PYTHON_VERSION[0]}.{MIN_PYTHON_VERSION[1]} or higher. "
+        f"You are using Python {sys.version_info[0]}.{sys.version_info[1]}.",
+        UserWarning,
+    )