diff --git a/docs/source/developer-guide/ci-overview.md b/docs/source/developer-guide/ci-overview.md index 8df1a6a00160..e708dce631a6 100644 --- a/docs/source/developer-guide/ci-overview.md +++ b/docs/source/developer-guide/ci-overview.md @@ -98,8 +98,8 @@ Each line contains the fully qualified test name followed by an optional specific hardware family. Example: ```text -examples/test_openai.py::test_llm_openai_triton_1gpu SKIP (https://nvbugspro.nvidia.com/bug/4963654) -full:GH200/examples/test_qwen2audio.py::test_llm_qwen2audio_single_gpu[qwen2_audio_7b_instruct] SKIP (arm is not supported) +accuracy/test_disaggregated_serving.py::TestDeepSeekV32Exp::test_auto_dtype[False] SKIP (https://nvbugs/6120535) +full:A100/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6422318) ``` Changes to `waives.txt` should include a bug link or brief explanation so other diff --git a/docs/source/features/auto_deploy/advanced/testing_strategy.md b/docs/source/features/auto_deploy/advanced/testing_strategy.md index d65841862ee6..b60eada7e722 100644 --- a/docs/source/features/auto_deploy/advanced/testing_strategy.md +++ b/docs/source/features/auto_deploy/advanced/testing_strategy.md @@ -117,8 +117,8 @@ Format: `path/to/test_file.py::test_function_name[param_id]` Example from `l0_a30.yml`: ```yaml -- accuracy/test_cli_flow.py::TestLlama3_1_8BInstruct::test_medusa_fp8_prequantized -- examples/test_multimodal.py::test_llm_multimodal_general[Qwen2-VL-7B-Instruct-pp:1-tp:1-float16-bs:1-cpp_e2e:False-nb:4] +- accuracy/test_llm_api_pytorch.py::TestLlama3_1_8B::test_auto_dtype +- unittest/_torch/modeling -k "modeling_llama" ``` ### Example: Adding an Accuracy Test diff --git a/examples/cpp_library/CMakeLists.txt b/examples/cpp_library/CMakeLists.txt deleted file mode 100644 index c60ff48d4cde..000000000000 --- a/examples/cpp_library/CMakeLists.txt +++ /dev/null @@ -1,55 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. cmake needs this line -cmake_minimum_required(VERSION 3.1) -# cmake_minimum_required(VERSION 2.8) - -# Enable C++11 -set(CMAKE_CXX_STANDARD 14) -set(CMAKE_CXX_STANDARD_REQUIRED TRUE) - -# Define project name -set(TARGET_NAME trt_llm_plugins_cpp_load_example) -project(${TARGET_NAME}) - -set(CMAKE_VERBOSE_MAKEFILE 1) - -# Compile options -set(CMAKE_C_FLAGS "-Wall -pthread ") -set(CMAKE_C_FLAGS_DEBUG "-g -O0") -set(CMAKE_C_FLAGS_RELEASE "-O2") -set(CMAKE_CXX_FLAGS "${CMAKE_C_FLAGS} -lstdc++") -set(CMAKE_CXX_FLAGS_DEBUG ${CMAKE_C_FLAGS_DEBUG}) -set(CMAKE_CXX_FLAGS_RELEASE ${CMAKE_C_FLAGS_RELEASE}) - -set(CMAKE_BUILD_TYPE release) -# set(CMAKE_BUILD_TYPE debug) - -find_package(CUDA REQUIRED) -message(STATUS "CUDA library status:") -message(STATUS " config: ${CUDA_DIR}") -message(STATUS " version: ${CUDA_VERSION}") -message(STATUS " libraries: ${CUDA_LIBRARIES}") -message(STATUS " include path: ${CUDA_INCLUDE_DIRS}") - -# Declare the executable target built from your sources -add_executable(${TARGET_NAME} main.cpp) - -# Link your application with CUDA libraries -target_link_libraries(${TARGET_NAME} LINK_PRIVATE ${CUDA_LIBRARIES}) -target_link_libraries(${TARGET_NAME} LINK_PRIVATE cudnn) -target_link_libraries(${TARGET_NAME} LINK_PRIVATE nvinfer) -target_link_libraries(${TARGET_NAME} LINK_PRIVATE nvinfer_plugin_tensorrt_llm) - -target_include_directories(${TARGET_NAME} PUBLIC /usr/local/cuda/include) diff --git a/examples/cpp_library/build.sh b/examples/cpp_library/build.sh deleted file mode 100755 index a384e5cad4a0..000000000000 --- a/examples/cpp_library/build.sh +++ /dev/null @@ -1,25 +0,0 @@ -#!/usr/bin/env bash - -BUILD_DIR="build" -SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) - -rm -rf ${BUILD_DIR} && mkdir -p ${BUILD_DIR} - -pushd ${BUILD_DIR} - -cmake \ - -DCMAKE_BUILD_TYPE=Release \ - .. - -make -j"$(grep -c ^processor /proc/cpuinfo)" - -export LD_LIBRARY_PATH="${SCRIPT_DIR}:${LD_LIBRARY_PATH}" - -# Test Lib -echo -echo "--------------------------------------------------------------------" -./trt_llm_plugins_cpp_load_example -echo "--------------------------------------------------------------------" -echo - -popd diff --git a/examples/cpp_library/main.cpp b/examples/cpp_library/main.cpp deleted file mode 100644 index 7613a75a140c..000000000000 --- a/examples/cpp_library/main.cpp +++ /dev/null @@ -1,79 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-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. - */ -#include -#include -#include - -#include "tensorrt_llm_libutils.h" - -int main(int argc, char* argv[]) -{ - class TRTLogger : public nvinfer1::ILogger - { - public: - void log(nvinfer1::ILogger::Severity severity, char const* msg) noexcept override - { - if (severity <= nvinfer1::ILogger::Severity::kERROR) - std::cerr << "[TensorRT LLM ERR]: " << msg << std::endl; - else if (severity == nvinfer1::ILogger::Severity::kWARNING) - std::cerr << "[TensorRT LLM WARNING]: " << msg << std::endl; - else - std::cout << "[TensorRT LLM LOG]: " << msg << std::endl; - } - }; - - TRTLogger* trtLogger = new TRTLogger(); - - std::string libname = "libtensorrt_llm_plugin.so"; - - /* =============== initLibNvInferPlugins =============== */ - - typedef bool (*initLibNvInferPlugins_sig)(void*, void const*); - - auto initLibNvInferPlugins = getTrtLLMFunction( - /*libFileSoName=*/libname, - /*symbol=*/"initLibNvInferPlugins"); - - std::cout << std::endl; - - std::string libNamespace = "tensorrt_llm"; - char const* libNamespace_cstr = libNamespace.data(); - - bool status1 = initLibNvInferPlugins(trtLogger, libNamespace_cstr); - std::cout << "Success Status: " << status1 << std::endl << std::endl; - - bool status2 = initLibNvInferPlugins(trtLogger, libNamespace_cstr); - std::cout << "Success Status: " << status2 << std::endl; - - /* =============== getInferLibVersion =============== */ - - std::cout << std::endl; - std::cout << "--------------------------------------------------------------------" << std::endl; - - typedef int32_t (*getInferLibVersion_sig)(); - - auto getInferLibVersion = getTrtLLMFunction( - /*libFileSoName=*/libname, - /*symbol=*/"getInferLibVersion"); - - std::cout << std::endl; - - int32_t version = getInferLibVersion(); - std::cout << "Version: " << version << std::endl; - - return 0; -} diff --git a/examples/cpp_library/tensorrt_llm_libutils.h b/examples/cpp_library/tensorrt_llm_libutils.h deleted file mode 100644 index aa60444eefa3..000000000000 --- a/examples/cpp_library/tensorrt_llm_libutils.h +++ /dev/null @@ -1,62 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-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. - */ -#if !defined(_WIN32) -#include -#endif // !defined(_WIN32) -#include -#include -#include - -#include "NvInfer.h" - -template -tSymbolSignature getTrtLLMFunction(std::string libFileSoName, std::string symbol) -{ -#if !defined(_WIN32) - std::cout << "Trying to load " << libFileSoName << " ..." << std::endl; - - // 1. Defining a handle to the library - void* handle = dlopen(libFileSoName.c_str(), RTLD_LAZY | RTLD_GLOBAL); - - // 2. Check for errors - char const* dl_error1 = dlerror(); - if (!handle) - { - throw std::runtime_error("Cannot open library: " + std::string(dl_error1)); - } - - // 3. Load actual queried `symbol` - std::cout << "Loading symbol `" << symbol << "` ..." << std::endl; - - tSymbolSignature symbolFctn = nullptr; - *(void**) (&symbolFctn) = dlsym(handle, symbol.c_str()); - - // 4. Check for errors - char const* dl_error2 = dlerror(); - if (dl_error2) - { - dlclose(handle); - throw std::runtime_error("Cannot load symbol '" + symbol + "': " + std::string(dl_error2)); - } - - return symbolFctn; -#else // on windows - throw std::runtime_error( - "`tSymbolSignature getTrtLLMFunction(std::string, std::string)` is not implemented on Windows."); - return nullptr; -#endif // !defined(_WIN32) -} diff --git a/examples/dora/README.md b/examples/dora/README.md deleted file mode 100644 index 21f70768a257..000000000000 --- a/examples/dora/README.md +++ /dev/null @@ -1,65 +0,0 @@ -# DoRA - -This document shows how to run a model using DoRA adapters. -DoRA is a PEFT strategy extending LoRA. It is fully supported in the Huggingface `peft` library. For a more detailed description please refer to the DoRA [paper](https://arxiv.org/abs/2402.09353) or official [repo](https://github.com/NVlabs/DoRA). - -## Support Matrix - * FP16/BF16 (over arbitrary precision of the base model). - * Supports adapters from Huggingface `peft` or from the official NVlabs [checkpoints](https://huggingface.co/sliuau/DoRA-weights). - * Multiple adapters (+ mixed LoRA/DoRA setups). - * inflight loading of new adapters to a preloaded base model. - * C++ and python runtime. - * Tensor parallelism and Pipeline parallelism. - -## Usage -Using DoRA is almost exactly the same as using LoRA in TRTLLM, with an additional preprocessing step. -While the official DoRA paper describes the magnitude normalization as part of the execution flow, it can be performed once beforehand to boost inference performance. - -Start by obtaining a local copy of your desired DoRA adapter **and** your base model. We'll use the official NVlabs checkpoint for LLaMA3-8B as an example: - -``` bash -git clone https://huggingface.co/sliuau/DoRA-weights -git clone https://huggingface.co/meta-llama/Meta-Llama-3-8B -``` - -Next, use the [normalize_weights.py](./normalize_weights.py) script to normalize the DoRA magnitude vectors in the adapter checkpoint. -The script requires access to both the local adapter weights and the local base model weights: - -``` bash -export NORMALIZED_DORA_ADAPTER=path/to/normalized/adapter/ckpt - -python ./normalize_weights.py -i DoRA-weights/llama_dora_commonsense_checkpoints/LLama3-8B/dora_r32 -b Meta-Llama-3-8B -o $NORMALIZED_DORA_ADAPTER -``` - -The script will create a new adapter checkpoint, with normalized DoRA vectors, in the provided path. - -Now we may convert our Llama checkpoint and build our TRT engine as described in the Llama [examples](../models/core/llama/README.md). When doing so, ensure you pass `--dora_plugin=enable` to the `trtllm-build` command, as well as enabling the lora plugin: - -``` bash -export CHECKPOINT_DIR=path/to/trtllm/ckpt -export ENGIRE_DIR=path/to/trtllm/engine - -python ../models/core/llama/convert_checkpoint.py --model_dir Meta-Llama-3-8B \ - --output_dir $CHECKPOINT_DIR \ - --dtype float16 - -trtllm-build --checkpoint_dir $CHECKPOINT_DIR \ - --output_dir $ENGINE_DIR \ - --gemm_plugin=auto \ - --lora_plugin=auto \ - --dora_plugin=enable \ - --lora_dir $NORMALIZED_DORA_ADAPTER -``` - -If you wish, you may provide additional LoRA / DoRA adapters to `trtllm-build`. - -**NOTE**: if you omit `--dora_plugin=enable`, you will not receive any warning even if you provide a DoRA adapter to `--lora_dir`. In such a case the DoRA magnitudes will simply be ignored during inference and you may receive wrong output. - -Proceed to execute the engine as you would a normal LoRA engine: - -``` bash -python ../run.py --engine_dir $ENGINE_DIR --tokenizer_dir Meta-Llama-3-8B --lora_task_uids 0 --max_output_len 32 --input_text ... -``` - -## Usage with Triton Server -Using DoRA over Triton is the same as using LoRA, but before using [hf_lora_convert.py](../hf_lora_convert.py), make sure you call [normalize_weights.py](./normalize_weights.py) and use the resulting normalized adapter. diff --git a/examples/dora/normalize_weights.py b/examples/dora/normalize_weights.py deleted file mode 100755 index 40d5ae172431..000000000000 --- a/examples/dora/normalize_weights.py +++ /dev/null @@ -1,289 +0,0 @@ -#! /usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -""" -This script applies preprocessing to the DoRA magnitude vector to speed up inference. -DoRA applies columnwise normalization and scaling to the LoRA output. -By applying the normalization to the scaling vector, we can skip calculating the normalization vector at inference time. -""" -import abc -import enum -import json -from pathlib import Path - -import numpy as np -import safetensors.torch as st -import torch - -StateDict = dict[str, torch.Tensor] - -PEFT_MODULE_PREFIX = "base_model.model." -PEFT_MODULE_SUFFIXES = [".lora_A.weight", ".lora_B.weight"] -DORA_VECTOR_SUFFIXES = [ - ".lora_magnitude_vector", # HF peft - ".weight_m_wdecomp.weight" # NVLabs -] - - -class HFWeightsReader(abc.ABC): - - @abc.abstractmethod - def __init__(self, model_dir: str) -> None: - ... - - @abc.abstractmethod - def get_weight(self, weight_name: str) -> torch.Tensor: - ... - - @abc.abstractmethod - def get_all(self) -> StateDict: - ... - - -class HFSafeTensorsReader(HFWeightsReader): - - def __init__(self, model_dir: str) -> None: - self.model_dir = Path(model_dir) - - self._fds = [ - st.safe_open(f, framework="torch") - for f in self.model_dir.glob("*.safetensors") - ] - - self._weight_to_fd = {} - for f in self._fds: - for k in f.keys(): - self._weight_to_fd[k] = f - - def get_weight(self, weight_name: str) -> torch.Tensor: - return self._weight_to_fd[weight_name].get_tensor(weight_name) - - def get_all(self) -> StateDict: - return {k: self.get_weight(k) for k in self._weight_to_fd.keys()} - - -class HFBinWeightsReader(HFWeightsReader): - - def __init__(self, model_dir: str) -> None: - self.model_dir = Path(model_dir) - - self._weights = {} - - for f in self.model_dir.glob("*.bin"): - self._weights.update( - torch.load(f, weights_only=True, mmap=True, map_location="cpu")) - - def get_weight(self, weight_name: str) -> torch.Tensor: - weight_name = f"{weight_name}.weight" - return self._weights[weight_name] - - def get_all(self) -> StateDict: - return self._weights - - -class WeightsFormat(enum.Enum): - BINARY = enum.auto() - SAFETENSORS = enum.auto() - UNKNOWN = enum.auto() - - def __str__(self) -> str: - if self == self.BINARY: - return "bin" - elif self == self.SAFETENSORS: - return "safetensors" - return "unknown" - - -def deduce_weights_format(model_dir: str) -> WeightsFormat: - model_dir_p = Path(model_dir) - - if any(model_dir_p.glob("*.safetensors")): - return WeightsFormat.SAFETENSORS - elif any(model_dir_p.glob("*.bin")): - return WeightsFormat.BINARY - return WeightsFormat.UNKNOWN - - -def get_weights_reader(model_dir: str) -> HFWeightsReader: - model_dir_p = Path(model_dir) - - if not model_dir_p.is_dir(): - raise ValueError( - f"{model_dir} is not a valid model directory: not found") - - weights_format = deduce_weights_format(model_dir) - - if weights_format == WeightsFormat.SAFETENSORS: - return HFSafeTensorsReader(model_dir) - elif weights_format == WeightsFormat.BINARY: - return HFBinWeightsReader(model_dir) - else: - raise ValueError( - f"{model_dir} does not contain .safetensors or .bin weights") - - -def normalize_hf_peft_module_name(module_name: str) -> str: - """ - Remove parts of the module name in the peft adapter to derive the module name of the base model. - """ - - if module_name.startswith(PEFT_MODULE_PREFIX): - module_name = module_name[len(PEFT_MODULE_PREFIX):] - - for suffix in PEFT_MODULE_SUFFIXES + DORA_VECTOR_SUFFIXES: - if module_name.endswith(suffix): - module_name = module_name[:-len(suffix)] - - return module_name - - -def get_peft_module_names(base_module_name: str) -> tuple[str, ...]: - """ - Convert the name of a base module to the names of its LoRA A and LoRA B weights. - """ - return tuple([ - f"{PEFT_MODULE_PREFIX}{base_module_name}{suffix}" - for suffix in PEFT_MODULE_SUFFIXES - ]) - - -def get_dora_magnitude_names(base_module_name: str) -> tuple[str, ...]: - """ - Convert the name of a base module to the potential names of its DoRA magnitude vectors. - """ - return tuple([ - f"{PEFT_MODULE_PREFIX}{base_module_name}{suffix}" - for suffix in DORA_VECTOR_SUFFIXES - ]) - - -def normalize_dora_vector(W: torch.Tensor, A: torch.Tensor, B: torch.Tensor, - mag: torch.Tensor, scale: float) -> torch.Tensor: - return mag / torch.linalg.norm(W + scale * B @ A, dim=1).to(W.dtype) - - -def normalize_dora_scales(lora_sd: StateDict, - weights_reader: HFWeightsReader, - alpha: float, - use_rslora: bool, - strip: bool = False) -> StateDict: - out_sd = {} - - while lora_sd: - # take some lora weight name - module_name = next(iter(lora_sd.keys())) - base_module_name = normalize_hf_peft_module_name(module_name) - A_name, B_name = get_peft_module_names(base_module_name) - magnitude_names = get_dora_magnitude_names(base_module_name) - - if module_name not in [A_name, B_name] + list(magnitude_names): - raise ValueError(f"Encountered unknown weight: {module_name}") - - # get lora weights - A = lora_sd.pop(A_name) - B = lora_sd.pop(B_name) - for name in magnitude_names: - if name in lora_sd: - mag_name = name - mag = lora_sd.pop(mag_name).view(-1) - break - else: - mag_name = "" - mag = None - - out_sd[A_name] = A.contiguous() - out_sd[B_name] = B.contiguous() - - if mag is not None and not strip: - # get base weight and normalize - W = weights_reader.get_weight(base_module_name + ".weight") - - adapter_size = A.size(0) - - if use_rslora: - scale = alpha / np.sqrt(adapter_size) - else: - scale = alpha / adapter_size - - mag = normalize_dora_vector(W, A, B, mag, scale) - out_sd[mag_name] = mag.contiguous() - - return out_sd - - -def save_state_dict(out_file: str, sd: StateDict) -> None: - out_path = Path(out_file) - - if out_path.suffix == ".safetensors": - st.save_file(sd, out_path) - elif out_path.suffix == ".bin": - torch.save(sd, out_path) - else: - raise ValueError(f"Unregornized weights format: {out_path.suffix}") - - -def normalize_peft_ckpt(model_dir: str, - base_model_dir: str, - out_dir: str, - strip: bool = False) -> None: - out_path = Path(out_dir) - out_path.mkdir(parents=True, exist_ok=True) - - weights_reader = get_weights_reader(base_model_dir) - lora_sd = get_weights_reader(model_dir).get_all() - - adapter_config_path = Path(f"{model_dir}/adapter_config.json") - with adapter_config_path.open() as f: - adapter_config = json.load(f) - - alpha = adapter_config["lora_alpha"] - use_rslora = adapter_config.get("use_rslora", False) - - new_sd = normalize_dora_scales(lora_sd, weights_reader, alpha, use_rslora, - strip) - - with (out_path / "adapter_config.json").open("w") as f: - json.dump(adapter_config, f) - - weights_format = deduce_weights_format(model_dir) - save_state_dict( - (out_path / f"adapter_model.{str(weights_format)}").as_posix(), new_sd) - - -if __name__ == "__main__": - import argparse - parser = argparse.ArgumentParser() - parser.add_argument( - '--out-dir', - '-o', - type=Path, - help='path to output adapter weights with normalized DoRA vectors', - required=True) - parser.add_argument('--in-dir', - '-i', - type=Path, - help='path to input lora checkpoint file', - required=True) - parser.add_argument("--base-model", - "-b", - help="Path to base model", - required=True) - parser.add_argument("--strip", - action="store_true", - help="remove DoRA vectors entirely") - - args = parser.parse_args() - - normalize_peft_ckpt(args.in_dir, args.base_model, args.out_dir, args.strip) diff --git a/examples/eval_long_context.py b/examples/eval_long_context.py deleted file mode 100644 index 90b7ef2dd270..000000000000 --- a/examples/eval_long_context.py +++ /dev/null @@ -1,331 +0,0 @@ -# MIT License - -# Copyright (c) 2023 OpenBMB - -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: - -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. - -# SPDX-FileCopyrightText: Copyright (c) 2022-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. - -# reference: https://github.com/OpenBMB/InfiniteBench/blob/main/src/eval_yarn_mistral.py - -import argparse -import ast -import json -from pathlib import Path - -import torch -from infinitebench.compute_scores import compute_scores -from infinitebench.eval_utils import (DATA_NAME_TO_MAX_NEW_TOKENS, - create_prompt, dump_jsonl, get_answer, - load_data) -from utils import (DEFAULT_HF_MODEL_DIRS, DEFAULT_PROMPT_TEMPLATES, - add_common_args, load_tokenizer, read_model_name) - -import tensorrt_llm -import tensorrt_llm.profiler as profiler -from tensorrt_llm.logger import logger -from tensorrt_llm.runtime import PYTHON_BINDINGS, ModelRunner - -if PYTHON_BINDINGS: - from tensorrt_llm.runtime import ModelRunnerCpp - -MAX_POSITION_ID = 128 * 1024 # Determined by the model -TRUNCATE_LEN = 128 * 1024 - - -def parse_arguments(args=None): - parser = argparse.ArgumentParser() - parser.add_argument('--batch_size', type=int, default=1) - parser.add_argument('--max_input_length', type=int, default=923) - parser.add_argument('--output_log_probs_npy', - type=str, - help='Numpy file where the log_probs are stored', - default=None) - - parser.add_argument('--output_cum_log_probs_npy', - type=str, - help='Numpy file where the cum_log_probs are stored', - default=None) - - parser.add_argument( - "--task", - type=str, - choices=['passkey', 'kv_retrieval'], - required=True, - help= - "Which task to use. Note that \"all\" can only be used in `compute_scores.py`.", # noqa - ) - parser.add_argument('--data_dir', - type=str, - default='./', - help="The directory of data.") - parser.add_argument("--output_dir", - type=str, - default=None, - help="Where to dump the prediction results.") # noqa - parser.add_argument( - "--start_idx", - type=int, - default=0, - help= - "The index of the first example to infer on. This is used if you want to evaluate on a (contiguous) subset of the data." - ) # noqa - parser.add_argument( - "--stop_idx", - type=int, - help= - "The index of the last example to infer on. This is used if you want to evaluate on a (contiguous) subset of the data. Defaults to the length of dataset." - ) # noqa - parser.add_argument('--tensorrt_llm_accuracy_threshold', - type=float, - default=99) - parser = add_common_args(parser) - - return parser.parse_args(args=args) - - -def parse_input(tokenizer, - input_text=None, - prompt_template=None, - add_special_tokens=True, - max_input_length=923, - pad_id=None, - num_prepend_vtokens=[], - model_name=None, - model_version=None): - if pad_id is None: - pad_id = tokenizer.pad_token_id - - batch_input_ids = [] - for curr_text in input_text: - if prompt_template is not None: - curr_text = prompt_template.format(input_text=curr_text) - input_ids = tokenizer.encode(curr_text, - add_special_tokens=add_special_tokens, - truncation=True, - max_length=max_input_length) - batch_input_ids.append(input_ids) - - if num_prepend_vtokens: - assert len(num_prepend_vtokens) == len(batch_input_ids) - base_vocab_size = tokenizer.vocab_size - len( - tokenizer.special_tokens_map.get('additional_special_tokens', [])) - for i, length in enumerate(num_prepend_vtokens): - batch_input_ids[i] = list( - range(base_vocab_size, - base_vocab_size + length)) + batch_input_ids[i] - - if 'GLM' in model_name and model_version == 'glm': - for ids in batch_input_ids: - ids.append(tokenizer.sop_token_id) - - batch_input_ids = [ - torch.tensor(x, dtype=torch.int32) for x in batch_input_ids - ] - return batch_input_ids - - -def main(args): - # model_name = "yarn-mistral" - runtime_rank = tensorrt_llm.mpi_rank() - logger.set_level(args.log_level) - - print(json.dumps(vars(args), indent=4)) - data_name = args.task - - # Model - max_tokens = DATA_NAME_TO_MAX_NEW_TOKENS[data_name] - - model_name, model_version = read_model_name(args.engine_dir) - if args.tokenizer_dir is None: - logger.warning( - "tokenizer_dir is not specified. Try to infer from model_name, but this may be incorrect." - ) - args.tokenizer_dir = DEFAULT_HF_MODEL_DIRS[model_name] - - tokenizer, pad_id, end_id = load_tokenizer( - tokenizer_dir=args.tokenizer_dir, - vocab_file=args.vocab_file, - model_name=model_name, - model_version=model_version, - tokenizer_type=args.tokenizer_type, - ) - - if not PYTHON_BINDINGS and not args.use_py_session: - logger.warning( - "Python bindings of C++ session is unavailable, fallback to Python session." - ) - args.use_py_session = True - if args.debug_mode and not args.use_py_session: - logger.warning( - "Debug mode is not supported in C++ session for now, fallback to Python session." - ) - args.use_py_session = True - runner_cls = ModelRunner if args.use_py_session else ModelRunnerCpp - runner_kwargs = dict( - engine_dir=args.engine_dir, - lora_dir=args.lora_dir, - rank=runtime_rank, - debug_mode=args.debug_mode, - lora_ckpt_source=args.lora_ckpt_source, - gpu_weights_percent=args.gpu_weights_percent, - ) - if args.medusa_choices is not None: - args.medusa_choices = ast.literal_eval(args.medusa_choices) - assert args.temperature == 1.0, "Medusa should use temperature == 1.0" - assert args.num_beams == 1, "Medusa should use num_beams == 1" - runner_kwargs.update(medusa_choices=args.medusa_choices) - if not args.use_py_session: - runner_kwargs.update( - max_batch_size=args.batch_size, - max_input_len=args.max_input_length, - max_output_len=max_tokens, - max_beam_width=args.num_beams, - max_attention_window_size=args.max_attention_window_size, - sink_token_length=args.sink_token_length, - max_tokens_in_paged_kv_cache=args.max_tokens_in_paged_kv_cache, - kv_cache_enable_block_reuse=args.kv_cache_enable_block_reuse, - kv_cache_free_gpu_memory_fraction=args. - kv_cache_free_gpu_memory_fraction, - enable_chunked_context=args.enable_chunked_context, - ) - runner = runner_cls.from_dir(**runner_kwargs) - - # Data - examples = load_data(data_name, data_dir=args.data_dir) - if args.stop_idx is None: - args.stop_idx = len(examples) - - output_path = None - if runtime_rank == 0: - if args.output_dir is not None: - result_dir = Path(args.output_dir, model_name) - result_dir.mkdir(exist_ok=True, parents=True) - - if args.stop_idx is None: - output_path = (result_dir / f"preds_{data_name}.jsonl") - else: - output_path = ( - result_dir / - f"preds_{data_name}_{args.start_idx}-{args.stop_idx}.jsonl" # noqa - ) - - prompt_template = None - if args.use_prompt_template and model_name in DEFAULT_PROMPT_TEMPLATES: - prompt_template = DEFAULT_PROMPT_TEMPLATES[model_name] - - if runtime_rank == 0: - preds = [] - logger.info("==== Evaluation ====") - logger.info(f"# examples: {len(examples)}") - logger.info(f"Start index: {args.start_idx}") - logger.info(f"Stop index: {args.stop_idx}") - logger.info(f"Max tokens: {max_tokens}") - assert args.batch_size == 1 - profiler.start('Evaluation') - for i in range(args.start_idx, args.stop_idx): - eg = examples[i] - input_text = [create_prompt(eg, data_name, args.data_dir)] - batch_input_ids = parse_input( - tokenizer=tokenizer, - input_text=input_text, - prompt_template=prompt_template, - add_special_tokens=args.add_special_tokens, - max_input_length=args.max_input_length, - pad_id=pad_id, - num_prepend_vtokens=args.num_prepend_vtokens, - model_name=model_name, - model_version=model_version) - input_lengths = [x.size(0) for x in batch_input_ids] - - if runtime_rank == 0: - logger.debug(f"====== Example {i} ======") - logger.debug(f"input_lengths: {input_lengths}") - logger.debug(f"input_text: {input_text}") - logger.debug(f"answer: {get_answer(eg, data_name)}") - outputs = runner.generate( - batch_input_ids, - max_new_tokens=max_tokens, - max_attention_window_size=args.max_attention_window_size, - sink_token_length=args.sink_token_length, - end_id=end_id, - pad_id=pad_id, - temperature=args.temperature, - top_k=args.top_k, - top_p=args.top_p, - num_beams=args.num_beams, - length_penalty=args.length_penalty, - early_stopping=args.early_stopping, - beam_width_array=args.beam_width_array, - repetition_penalty=args.repetition_penalty, - presence_penalty=args.presence_penalty, - frequency_penalty=args.frequency_penalty, - prompt_ignore_length=args.prompt_ignore_length, - # stop_words_list=stop_words_list, - # bad_words_list=bad_words_list, - output_cum_log_probs=(args.output_cum_log_probs_npy != None), - output_log_probs=(args.output_log_probs_npy != None), - lora_uids=args.lora_task_uids, - prompt_table=args.prompt_table_path, - prompt_tasks=args.prompt_tasks, - streaming=args.streaming, - output_sequence_lengths=True, - return_dict=True, - medusa_choices=args.medusa_choices) - torch.cuda.synchronize() - if runtime_rank == 0: - output_ids = outputs['output_ids'] - output_beams_list = [ - tokenizer.batch_decode(output_ids[batch_idx, :, - input_lengths[batch_idx]:], - skip_special_tokens=True) - for batch_idx in range(args.batch_size) - ] - - logger.debug(f"preds: {output_beams_list[0]}") - preds.append({ - "id": i, - "prediction": output_beams_list[0][0], - "ground_truth": get_answer(eg, data_name), - "input_lengths": input_lengths, - }) - if output_path is not None: - dump_jsonl(preds, output_path) - profiler.stop('Evaluation') - - if runtime_rank == 0: - logger.info( - f'Evaluation takes: {profiler.elapsed_time_in_sec("Evaluation")} sec.' - ) - logger.info("Compute the score") - acc = compute_scores(preds, args.task) * 100 - logger.info(f"{args.task} accuracy: {acc:.2f} ({len(preds)})") - - if args.tensorrt_llm_accuracy_threshold is not None: - assert acc >= args.tensorrt_llm_accuracy_threshold, f"acc ({acc}) < tensorrt_llm_accuracy_threshold ({args.tensorrt_llm_accuracy_threshold})" - - -if __name__ == "__main__": - args = parse_arguments() - main(args) diff --git a/examples/generate_checkpoint_config.py b/examples/generate_checkpoint_config.py deleted file mode 100644 index a11104eeba68..000000000000 --- a/examples/generate_checkpoint_config.py +++ /dev/null @@ -1,149 +0,0 @@ -import argparse -import json -import os - -from tensorrt_llm.quantization import KV_CACHE_QUANT_ALGO_LIST, QUANT_ALGO_LIST - - -def parse_arguments(): - parser = argparse.ArgumentParser() - - parser.add_argument( - '--output_path', - type=str, - default='config.json', - help='The path to save the TensorRT LLM checkpoint config.json file') - parser.add_argument('--architecture', type=str, default='GPTForCausalLM') - parser.add_argument('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument('--vocab_size', type=int, default=32000) - parser.add_argument('--max_position_embeddings', type=int, default=1024) - parser.add_argument('--hidden_size', type=int, default=768) - parser.add_argument('--intermediate_size', type=int, default=None) - parser.add_argument('--num_hidden_layers', type=int, default=12) - parser.add_argument('--num_attention_heads', type=int, default=12) - parser.add_argument('--num_key_value_heads', type=int, default=None) - parser.add_argument('--hidden_act', type=str, default='gelu') - parser.add_argument('--norm_epsilon', type=float, default=1e-5) - parser.add_argument('--position_embedding_type', - type=str, - default='learned_absolute') - parser.add_argument( - '--use_parallel_embedding', - action='store_true', - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - - parser.add_argument('--quant_algo', - type=str, - default=None, - choices=[None] + QUANT_ALGO_LIST) - parser.add_argument('--kv_cache_quant_algo', - type=str, - default=None, - choices=[None] + KV_CACHE_QUANT_ALGO_LIST) - parser.add_argument('--group_size', type=int, default=64) - parser.add_argument('--smoothquant_val', type=float, default=None) - parser.add_argument('--has_zero_point', default=False, action='store_true') - parser.add_argument('--pre_quant_scale', default=False, action='store_true') - parser.add_argument('--exclude_modules', nargs='+', default=None) - - parser.add_argument('--bias', default=False, action='store_true') - parser.add_argument('--apply_query_key_layer_scaling', - default=False, - action='store_true') - parser.add_argument('--rotary_pct', type=float, default=1.0) - parser.add_argument('--rotary_base', type=float, default=10000.0) - parser.add_argument('--rotary_scaling', nargs=2, type=str, default=None) - - args = parser.parse_args() - return args - - -if __name__ == '__main__': - args = parse_arguments() - world_size = args.tp_size * args.pp_size - - assert args.output_path.endswith('.json') - output_dir = os.path.dirname(args.output_path) - if output_dir and not os.path.exists(output_dir): - os.makedirs(output_dir) - - config = { - 'architecture': args.architecture, - 'dtype': args.dtype, - 'vocab_size': args.vocab_size, - 'max_position_embeddings': args.max_position_embeddings, - 'hidden_size': args.hidden_size, - 'intermediate_size': args.intermediate_size, - 'num_hidden_layers': args.num_hidden_layers, - 'num_attention_heads': args.num_attention_heads, - 'num_key_value_heads': args.num_key_value_heads, - 'hidden_act': args.hidden_act, - 'norm_epsilon': args.norm_epsilon, - 'position_embedding_type': args.position_embedding_type, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'quantization': { - 'quant_algo': args.quant_algo, - 'kv_cache_quant_algo': args.kv_cache_quant_algo, - 'exclude_modules': args.exclude_modules, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'bias': args.bias, - 'apply_query_key_layer_scaling': args.apply_query_key_layer_scaling, - 'rotary_pct': args.rotary_pct, - 'rotary_base': args.rotary_base, - 'rotary_scaling': args.rotary_scaling, - } - - if args.intermediate_size is None: - config['intermediate_size'] = args.hidden_size * 4 - if args.num_key_value_heads is None: - config['num_key_value_heads'] = args.num_attention_heads - - if args.quant_algo is not None: - if 'AWQ' in args.quant_algo or 'GPTQ' in args.quant_algo: - config['quantization'].update({ - 'group_size': - args.group_size, - 'has_zero_point': - args.has_zero_point, - 'pre_quant_scale': - args.pre_quant_scale, - }) - if 'SQ' in args.quant_algo: - config['quantization'].update({ - 'smoothquant_val': - args.smoothquant_val, - }) - - with open(args.output_path, 'w') as f: - json.dump(config, f, indent=4) diff --git a/examples/generate_xgrammar_tokenizer_info.py b/examples/generate_xgrammar_tokenizer_info.py deleted file mode 100644 index 67e05eb9c968..000000000000 --- a/examples/generate_xgrammar_tokenizer_info.py +++ /dev/null @@ -1,51 +0,0 @@ -#!/usr/bin/env python3 -# 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. - -import argparse -import json -import os -from pathlib import Path - -from transformers import AutoTokenizer - -from tensorrt_llm.llmapi.tokenizer import _xgrammar_tokenizer_info - - -def generate_xgrammar_tokenizer_info(args): - - tokenizer = AutoTokenizer.from_pretrained(str(args.model_dir)) - tokenizer_info = _xgrammar_tokenizer_info(tokenizer) - - os.makedirs(args.output_dir, exist_ok=True) - with open(str(args.output_dir / "xgrammar_tokenizer_info.json"), 'w') as f: - json.dump(tokenizer_info, f) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', - type=Path, - default=None, - required=True, - help="HF model directory") - parser.add_argument( - '--output_dir', - type=Path, - default=None, - required=True, - help="File path to save xgrammar's info. in json format") - args = parser.parse_args() - generate_xgrammar_tokenizer_info(args) diff --git a/examples/hf_lora_convert.py b/examples/hf_lora_convert.py deleted file mode 100755 index 019d78a48563..000000000000 --- a/examples/hf_lora_convert.py +++ /dev/null @@ -1,266 +0,0 @@ -#! /usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -import argparse -import datetime -import json -import logging -import re -from collections import defaultdict -from pathlib import Path - -import numpy as np -import torch - -from tensorrt_llm._utils import str_dtype_to_torch, torch_to_numpy -from tensorrt_llm.lora_manager import LoraManager -from tensorrt_llm.models.convert_utils import get_model_path, load_state_dict - -log_format = "%(asctime)s %(name)s [%(levelname)s] %(message)s" -logging.basicConfig(format=log_format) -LOGGER = logging.getLogger(__name__) - - -def save_val(val, dir, key, tp_num=None, write_npy=False): - ext = "npy" if write_npy else "bin" - suffix = ext if tp_num is None else f"{tp_num}.{ext}" - if write_npy: - np.save(dir / f"model.{key}.{suffix}", val) - else: - val.tofile(dir / f"model.{key}.{suffix}") - - -def get_all_lora_weights(lora_weights): - all_weights = defaultdict(lambda: defaultdict(dict)) - pattern = re.compile( - r'(.*\.layers\.([0-9]+)\.(self_attn|mlp)\.([a-z_]+))\.(?:lora_(?:(A|B)\.weight|(magnitude)_vector)|weight_(m_wdecomp).weight).*' - ) - moe_pattern = re.compile( - r'(.*\.layers\.([0-9]+)\.(block_sparse_moe)\.((experts)\.([0-9]+)\.|)([a-zA-Z0-9_]+))\.(?:lora_(?:(A|B)\.weight|(magnitude)_vector)|weight_(m_wdecomp).weight).*' - ) - for key, weights in lora_weights.items(): - m = pattern.match(key) - m_moe = moe_pattern.match(key) - if m: - layer_idx = int(m.group(2)) - hf_module = m.group(4) - inout = m.group(5) - dora_magnitude = m.group(6) or m.group(7) - - if inout: - inout = "in" if inout == "A" else "out" - all_weights[layer_idx][hf_module][inout] = weights - elif dora_magnitude: - LOGGER.warning( - "Detected DoRA magnitude vector, make sure it was preprocessed and normalized using the proper base model weights" - ) - all_weights[layer_idx][hf_module]["magnitude"] = weights.view( - -1) - - elif m_moe: - layer_idx = int(m_moe.group(2)) - hf_module = m_moe.group(7) - inout = m_moe.group(8) - dora_magnitude = m_moe.group(9) or m.group(10) - - if inout: - inout = "in" if inout == "A" else "out" - all_weights[layer_idx][hf_module][inout] = weights - elif dora_magnitude: - LOGGER.warning( - "Detected DoRA magnitude vector, make sure it was preprocessed and normalized using the proper base model weights" - ) - all_weights[layer_idx][hf_module]["magnitude"] = weights.view( - -1) - else: - print(f"no match {key}") - continue - return all_weights - - -def preprocess_lora_weights(lora_model): - # Swap weights of gate_up_proj - for key, value in lora_model.items(): - if "gate_up_proj.lora_B.weight" in key: - print("Swap {}".format(key)) - original_weights = value.contiguous().clone() - half_split = original_weights.shape[0] // 2 - first_half = original_weights[:half_split, :] - second_half = original_weights[half_split:, :] - value = torch.cat((second_half, first_half), dim=0) - lora_model[key] = value - return lora_model - - -hf_modules_to_trtllm_modules = { - "q_proj": "attn_q", - "v_proj": "attn_v", - "k_proj": "attn_k", - "qkv_proj": "attn_qkv", - "query_key_value": "attn_qkv", - "o_proj": "attn_dense", - "dense": "attn_dense", - "gate_proj": "mlp_h_to_4h", - "down_proj": "mlp_4h_to_h", - "up_proj": "mlp_gate", - "gate_up_proj": "mlp_h_to_4h", - "c_fc": "mlp_h_to_4h", - "c_proj": "mlp_4h_to_h", - "w1": "moe_h_to_4h", - "w2": "moe_4h_to_h", - "w3": "moe_gate", - "gate": "moe_router", -} # lora modules on llama -hf_modules_to_module_id = { - k: LoraManager.LORA_MODULE_IDS[v] - for k, v in hf_modules_to_trtllm_modules.items() -} - - -def convert_hf_model(model_dir, dtype, out_dir): - saved_dir = Path(out_dir) - saved_dir.mkdir(parents=True, exist_ok=True) - with open(f"{model_dir}/adapter_config.json", "r") as f: - config = json.load(f) - - alpha = config.get("lora_alpha") - use_rslora = config.get("use_rslora", False) - - lora_model = load_state_dict(get_model_path(model_dir, "adapter_model")) - lora_model = preprocess_lora_weights(lora_model) - all_weights = get_all_lora_weights(lora_model) - converted_weights = [] - converted_config = [] - - def derive_adapter_size(inout_weight: torch.Tensor) -> int: - assert len(inout_weight.shape) == 2 - dim0, dim1 = inout_weight.shape - # assume the hidden dim is the larger of the 2 - adapter_size = min(dim0, dim1) - return adapter_size - - def derive_weights_scale(adapter_size: int, alpha: float, - use_rslora: bool) -> float: - if use_rslora: - return alpha / np.sqrt(adapter_size) - return alpha / adapter_size - - for layer_idx, layer_weights in all_weights.items(): - for hf_module, module_weights in layer_weights.items(): - in_weights = module_weights['in'] - out_weights = module_weights['out'] - magnitude = module_weights.get("magnitude", None) - is_dora = magnitude is not None - - processed_weights = [] - - assert len(in_weights.shape) == 2 - assert len(out_weights.shape) == 2 - assert not is_dora or len(magnitude.shape) == 1 - - adapter_size = derive_adapter_size(in_weights) - assert adapter_size == derive_adapter_size( - out_weights), "adapter size of A mismatches adapter size of B" - scale = derive_weights_scale(adapter_size, alpha, use_rslora) - - for w, inout in ((in_weights, "in"), (out_weights, "out")): - dim0 = w.shape[0] - dim1 = w.shape[1] - # in_weights should have shape [adaper_size, hidden] - if dim1 < dim0 and inout == "in": - w = w.transpose(1, 0) - # out_weights should have shape [hidden, adapter_size] - elif dim0 < dim1 and inout == "out": - w = w.transpose(1, 0) - if inout == "out": - w = w * scale - w = w.contiguous().flatten().to(dtype=str_dtype_to_torch(dtype)) - processed_weights.append(w) - - if is_dora: - processed_weights.append(magnitude.contiguous().flatten().to( - dtype=str_dtype_to_torch(dtype))) - - processed_weights = torch.concatenate(processed_weights).flatten() - converted_weights.append(processed_weights) - converted_config.append([ - hf_modules_to_module_id[hf_module], layer_idx, adapter_size, - 1 if is_dora else 0 - ]) - max_row_size = 0 - for t in converted_weights: - max_row_size = max(max_row_size, t.shape[0]) - for i in range(len(converted_weights)): - converted_weights[i] = torch.nn.functional.pad( - converted_weights[i], - (0, max_row_size - converted_weights[i].shape[0])).unsqueeze(0) - converted_weights = torch_to_numpy( - torch.concatenate( - converted_weights, - dim=0).unsqueeze(0).to(dtype=str_dtype_to_torch(dtype)).cpu()) - converted_config = torch.tensor(converted_config, - dtype=torch.int32, - device='cpu').unsqueeze(0).numpy() - - save_val(converted_weights, - saved_dir, - "lora_weights", - tp_num=None, - write_npy=True) - save_val(converted_config, - saved_dir, - "lora_config", - tp_num=None, - write_npy=True) - - -def main(args): - start_time = datetime.datetime.now() - convert_hf_model(args.in_file, args.storage_type, args.out_dir) - - LOGGER.info("Spent %s (h:m:s) to convert the prompt model", - datetime.datetime.now() - start_time) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - '--out-dir', - '-o', - type=Path, - help='path to output embedding table file in the .npy format', - required=True) - parser.add_argument('--in-file', - '-i', - type=Path, - help='path to input lora checkpoint file', - required=True) - parser.add_argument("--verbose", - action="store_true", - help="Provide verbose messages") - parser.add_argument("--storage-type", - type=str, - default="float16", - choices=["float32", "float16", "bfloat16"]) - args = parser.parse_args() - - LOGGER.setLevel(logging.DEBUG if args.verbose else logging.INFO) - - print("\n=============== Argument ===============") - for key in vars(args): - print(f"{key}: {vars(args)[key]}") - print("========================================") - - main(args) diff --git a/examples/language_adapter/README.md b/examples/language_adapter/README.md deleted file mode 100755 index 8487c8ab42a0..000000000000 --- a/examples/language_adapter/README.md +++ /dev/null @@ -1,99 +0,0 @@ -# Language-Adapter - -This document shows how to build and run a model with Language-Adapter plugin in TensorRT LLM on NVIDIA GPUs. - -## Overview -The concept of Language Adapter during inference time was introduced in [MAD-X: An Adapter-Based Framework for Multi-Task Cross-Lingual Transfer -](https://arxiv.org/pdf/2005.00052): -> we can simply replace a language-specific adapter trained for English with a language-specific adapter trained for Quechua at inference time. - -The implementation is done with MOE plugin with static expert selection passed during runtime as a parameter in request. - -For instance, encoder-decoder model may leverage language adapter for language-specific translation tasks when each of the language-adapter is trained for a specific language, this language adapter plugin achieves the language switching within one session only by passing in the `language_task_uid` to the plugin. - -The model checkpoint here is not publicly available. Please leverage `layers/language_adapter.py` in your own model. - -### Engine Preparation (convert and build) -``` -MODEL_DIR="dummy_model" # model not publicly available -INFERENCE_PRECISION="float16" -TP_SIZE=1 -PP_SIZE=1 -WORLD_SIZE=1 -MODEL_TYPE=language_adapter -MODEL_NAME=$MODEL_TYPE -CKPT_DIR=/scratch/tmp/trt_models/${MODEL_NAME}/${WORLD_SIZE}-gpu/${INFERENCE_PRECISION} -ENGINE_DIR=/scratch/tmp/trt_engines/${MODEL_NAME}/${WORLD_SIZE}-gpu/${INFERENCE_PRECISION} - -max_beam=5 -max_batch=32 -max_input_len=1024 -max_output_len=1024 - -python ../enc_dec/convert_checkpoint.py --model_type ${MODEL_TYPE} \ - --model_dir ${MODEL_DIR} \ - --output_dir $CKPT_DIR \ - --tp_size ${TP_SIZE} \ - --pp_size ${PP_SIZE} \ - --dtype ${INFERENCE_PRECISION} \ - --workers 1 - -trtllm-build --checkpoint_dir $CKPT_DIR/encoder \ - --output_dir $ENGINE_DIR/encoder \ - --paged_kv_cache disable \ - --moe_plugin auto \ - --bert_attention_plugin ${INFERENCE_PRECISION} \ - --gpt_attention_plugin ${INFERENCE_PRECISION} \ - --gemm_plugin ${INFERENCE_PRECISION} \ - --remove_input_padding enable \ - --max_input_len ${max_input_len} \ - --max_beam_width ${max_beam} \ - --max_batch_size ${max_batch} - -trtllm-build --checkpoint_dir $CKPT_DIR/decoder \ - --output_dir $ENGINE_DIR/decoder \ - --paged_kv_cache enable \ - --moe_plugin auto \ - --bert_attention_plugin ${INFERENCE_PRECISION} \ - --gpt_attention_plugin ${INFERENCE_PRECISION} \ - --gemm_plugin ${INFERENCE_PRECISION} \ - --remove_input_padding enable \ - --max_input_len 1 \ - --max_beam_width ${max_beam} \ - --max_batch_size ${max_batch} \ - --max_seq_len ${max_output_len} -``` - -### CPP runtime -A list `language_task_uids` that includes the language_task_uid for each input prompt is required: -``` -# translate 2 sentence, 1 to France (language_task_uid=3) 1 to Spanish (language_task_uid=2). -# language_task_uids = [3, 2] - -TEXT="Where is the nearest restaurant? Wikipedia is a free online encyclopedia written and maintained by a community of volunteers (called Wikis) through open collaboration and the use of MediaWiki, a wiki-based editing system." - -python3 ../run.py --engine_dir $ENGINE_DIR --tokenizer_type "language_adapter" --max_input_length 512 --max_output_len 512 --num_beams 1 --input_file input_ids.npy --tokenizer_dir $MODEL_DIR --language_task_uids 3 2 - -# Input [Text 0]: "" -# Output [Text 0 Beam 0]: "Où se trouve le restaurant le plus proche ? Wikipédia est une encyclopédie en ligne gratuite écrite et maintenue par une communauté de bénévoles (appelés Wikis) grâce à une collaboration ouverte et à l'utilisation de MediaWiki, un système d'édition basé sur wiki." -# Input [Text 1]: "" -# Output [Text 1 Beam 0]: "¿Dónde está el restaurante más cercano? Wikipedia es una enciclopedia en línea gratuita escrita y mantenida por una comunidad de voluntarios (llamada Wikis) a través de la colaboración abierta y el uso de MediaWiki, un sistema de edición basado en wiki." - -``` - -### Python runtime -Currently Python runtime does not support beam_width > 1. - -For Python runtime, full routing information of length [num_tokens, 1] is required for both encoder and decoder, which stacks routing information for each token in a batch of requests. -``` -# language_adapter_routing = get_language_adapter_routings(language_task_uid, input_ids) - -TEXT="Where is the nearest restaurant? Wikipedia is a free online encyclopedia written and maintained by a community of volunteers (called Wikis) through open collaboration and the use of MediaWiki, a wiki-based editing system." - -python3 ../enc_dec/run.py --engine_dir $ENGINE_DIR --engine_name ${MODEL_NAME} --model_name $MODEL_DIR --max_new_token=64 --num_beams=1 - -# in the run.py, 2 input prompts and 2 language task uids are provided. The two task uid represent the language of the input prompts to be translated to. - -# TRT-LLM output text: ['¿Dónde está el restaurante más cercano? Wikipedia es una enciclopedia en línea gratuita escrita y mantenida por una comunidad de voluntarios (llamada Wikis) a través de la colaboración abierta y el uso de MediaWiki, un sistema de edición basado en wiki.', "Où se trouve le restaurant le plus proche ? Wikipédia est une encyclopédie en ligne gratuite é -crite et maintenue par une communauté de bénévoles (appelés Wikis) grâce à une collaboration ouverte et à l'utilisation de MediaWiki, un système d'édition basé sur wiki."] -``` diff --git a/examples/llm-eval/lm-eval-harness/README.md b/examples/llm-eval/lm-eval-harness/README.md deleted file mode 100644 index c3854654dd0d..000000000000 --- a/examples/llm-eval/lm-eval-harness/README.md +++ /dev/null @@ -1,67 +0,0 @@ -# Evaluation scripts for LLM tasks - -This folder includes code to use the [LM-Eval-Harness](https://github.com/EleutherAI/lm-evaluation-harness), a unified framework to test generative language models on a large number of different evaluation tasks. The supported eval tasks are [here](https://github.com/EleutherAI/lm-evaluation-harness/tree/main/lm_eval/tasks). - -The following instructions show how to evaluate TRT-LLM engines with the benchmark. - -## Instructions - -### TRT-LLM API - -Build the TRT-LLM engine using `trtllm-build`. - -Install the `lm_eval` package in the `requirements.txt` file in this folder. - -Run the evaluation script with the following command: - -```sh -python lm_eval_tensorrt_llm.py --model trt-llm \ - --model_args tokenizer=,model=,chunk_size= \ - --tasks -``` - -In the LM-Eval-Harness, model args are submitted as a comma-separated list of the form `arg=value`. The `trt-llm` model supports the following `model_args`: - -| Name | Description | Default Value | -|--------------------------|-------------------------------------------------------------------|----------------| -| tokenizer | directory containing the HF tokenizer. | | -| model | directory containing the TRTLLM engine or torch model. | | -| max_gen_toks | max number of tokens to generate (if not specified in gen_kwargs) | 256 | -| chunk_size | number of async requests to send at once to the engine | 200 | -| max_tokens_kv_cache | max tokens in paged KV cache | None | -| free_gpu_memory_fraction | KV cache free GPU memory fraction | 0.9 | -| trust_remote_code | trust remote code; use if necessary to set up the tokenizer | False | -| tp | tensor parallel size (for torch backend) | no. of workers | -| use_cuda_graph | enable CUDA graph | True | -| max_context_length | maximum context length for evaluation | None | -| moe_expert_parallel_size | expert parallel size for MoE models | None | -| moe_backend | backend for MoE models (e.g., "TRTLLM") | "TRTLLM" | - -### Torch backend - -Install the `lm_eval` package in the `requirements.txt` file in this folder. - -Run the evaluation script with the same command as above, but include `backend=torch` in the `model_args`. For example: - -```sh -python lm_eval_tensorrt_llm.py --model trt-llm \ - --model_args model=,backend=torch,chunk_size= \ - --tasks -``` - -### trtllm-serve - -Build the TRT-LLM engine using `trtllm-build` and deploy with `trtllm-serve`. - -Install the `lm_eval` package in the `requirements.txt` file in this folder. - -Run the evaluation script with the following command: - -```sh -python lm_eval_tensorrt_llm.py --model local-completions \ - --model_args base_url=http://${HOST_NAME}:8001/v1/completions,model=,tokenizer= \ - --tasks \ - --batch_size <#> -``` - -Because `trtllm-serve` is OpenAI API compatible, we can use the `local-completions` model built in to `lm_eval`, which supports [these model_args](https://github.com/EleutherAI/lm-evaluation-harness/blob/v0.4.7/lm_eval/models/openai_completions.py#L12). diff --git a/examples/llm-eval/lm-eval-harness/lm_eval_tensorrt_llm.py b/examples/llm-eval/lm-eval-harness/lm_eval_tensorrt_llm.py deleted file mode 100644 index 1738242267d9..000000000000 --- a/examples/llm-eval/lm-eval-harness/lm_eval_tensorrt_llm.py +++ /dev/null @@ -1,317 +0,0 @@ -# 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. -import gc -import json -import logging -import os -import signal -import threading -import time -from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple - -import torch -import torch.nn.functional as F -import transformers -from lm_eval.__main__ import cli_evaluate -from lm_eval.api.model import TemplateLM -from lm_eval.api.registry import register_model -from packaging.version import parse -from tqdm import tqdm - -import tensorrt_llm -from tensorrt_llm import LLM as TORCH_LLM -from tensorrt_llm._tensorrt_engine import LLM as TRT_LLM -from tensorrt_llm.bindings.executor import DecodingConfig -from tensorrt_llm.llmapi import KvCacheConfig as TRT_KvCacheConfig -from tensorrt_llm.llmapi import RequestOutput, SamplingParams -from tensorrt_llm.llmapi.llm_args import MoeConfig - -logger = logging.getLogger(__name__) - - -@register_model("trt-llm") -class TRTLLMEvalBase(TemplateLM): - - def __init__( - self, - model: str, - tokenizer: Optional[str] = None, - tp: int = 0, # tensor_parallel_size - max_gen_toks: int = 256, - chunk_size: int = 200, - max_tokens_kv_cache: Optional[int] = None, - free_gpu_memory_fraction: float = 0.9, - trust_remote_code: bool = False, - use_cuda_graph: bool = True, - backend: str = 'trt', - max_context_length: Optional[int] = None, - moe_expert_parallel_size: Optional[int] = None, - moe_backend: Optional[str] = "TRTLLM", - enable_chunked_prefill: bool = False, - max_num_tokens: Optional[int] = None, - **kwargs, - ): - # initialize TemplateLM, copied from TemplateAPI - super().__init__() - assert isinstance(model, str) - assert parse(tensorrt_llm.__version__) >= parse("0.15.0") - - self.max_gen_toks = max_gen_toks - self.chunk_size = chunk_size - self.backend = backend - self.max_context_length = max_context_length - self.moe_expert_parallel_size = moe_expert_parallel_size - self.moe_backend = moe_backend - trt_kv_cache_config = TRT_KvCacheConfig(enable_block_reuse=False) - trt_kv_cache_config.free_gpu_memory_fraction = free_gpu_memory_fraction - if max_tokens_kv_cache is not None: - trt_kv_cache_config.max_tokens = max_tokens_kv_cache - - if tokenizer is None: - # Assume the tokenizer is stored in the model_dir if not specified. - tokenizer = model - logger.info(f"Tokenizer: {tokenizer}") - self.tokenizer = transformers.AutoTokenizer.from_pretrained( - tokenizer, trust_remote_code=trust_remote_code) - - if self.tokenizer.pad_token_id is None: - self.tokenizer.pad_token_id = self.tokenizer.eos_token_id - - if self.backend == 'torch': - kwargs.pop('batch_size') - if tp < 1: - tp = torch.cuda.device_count() - - pytorch_config_params = { - 'cuda_graph_config': {} if use_cuda_graph else None, - "print_iter_log": False, - 'moe_config': MoeConfig(backend=self.moe_backend) - } - - # stop words not currently supported by torch backend - self.use_stop_words = False - - self.llm = TORCH_LLM( - model=model, - tensor_parallel_size=tp, - trust_remote_code=trust_remote_code, - enable_chunked_prefill=enable_chunked_prefill, - max_num_tokens=max_num_tokens, - **pytorch_config_params, - tokenizer=self.tokenizer, - kv_cache_config=trt_kv_cache_config, - moe_expert_parallel_size=self.moe_expert_parallel_size, - **kwargs) - logger.info("Loaded TRT-LLM Torch engine") - else: - with open(Path(model) / "config.json", "r") as engine_config_file: - engine_config = json.load(engine_config_file) - build_config = engine_config["build_config"] - world_size = (engine_config.get("pretrained_config", {}).get( - "mapping", {}).get("world_size", 1)) - if max_tokens_kv_cache is None: - max_tokens_kv_cache = build_config[ - "max_seq_len"] * build_config["max_batch_size"] - self.gather_context_logits = build_config.get( - "gather_context_logits", False) - - medusa_choices = kwargs[ - 'medusa_choices'] if 'medusa_choices' in kwargs else None - 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." - - self.llm = TRT_LLM(model=model, - tokenizer=self.tokenizer, - kv_cache_config=trt_kv_cache_config, - **kwargs) - self.max_length = build_config['max_seq_len'] - 1 - logger.info("Loaded TRT-LLM engine") - - @property - def eot_token_id(self) -> int: - return self.llm.tokenizer.eos_token_id - - def tok_encode(self, string, add_special_tokens=False, **kwargs): - return self.llm.tokenizer.encode(string, - add_special_tokens=add_special_tokens, - **kwargs) - - def _loglikelihood_tokens( - self, - requests: List[Any], - disable_tqdm: bool = False) -> List[Tuple[float, bool]]: - """Compute the log likelihood of the continuation given the context.""" - if self.backend == 'torch': - raise NotImplementedError( - 'Torch backend does not return context logits yet') - - num_r = len(requests) - desc = "Processing loglikelihood requests" - sampling_params = SamplingParams(max_tokens=1, - return_context_logits=True) - - # process requests - futures: Dict[int, RequestOutput] = {} - results = [] - for i, request in tqdm(enumerate(requests), - desc=desc, - total=num_r, - disable=disable_tqdm): - # asynchronously submit a chunk of requests ahead of time... - if i % self.chunk_size == 0: - for j in range(i, min(i + self.chunk_size, num_r)): - prompt_ids = requests[j][1] + requests[j][2] - futures[j] = self.llm.generate_async( - prompt_ids, sampling_params) - - # process the output of the request i - r_out: RequestOutput = futures.pop(i).result() - - # check continuation portion of the prompt - # NOTE: context_logits are offset by 1 since they predict future token - ctxlen = len(request[1]) - token_ids_cont = request[2] - logits_cont = r_out.context_logits[ctxlen - 1:-1] # [sl, vocab] - logprobs_cont = F.log_softmax(logits_cont, dim=-1) # [sl, vocab] - top_tokens_cont = logprobs_cont.argmax(dim=-1).tolist() # [sl] - - # compute logprob and check for greedy - logprob_sum = sum(logprobs_cont[list(range(len(logprobs_cont))), - token_ids_cont]).item() - is_greedy = top_tokens_cont == token_ids_cont - - results.append((logprob_sum, is_greedy)) - - # clear response - del r_out - - return results - - def loglikelihood_rolling(self, requests, disable_tqdm: bool = False): - raise NotImplementedError - - def generate_until(self, - requests: List[Any], - disable_tqdm: bool = False) -> List[str]: - # some book-keeping and parameters... - num_r = len(requests) - desc = "Processing generate requests" - - if self.max_context_length is not None: - """ - Create updated_requests to contain qualified requests with the context length <= max_context_length. - Unqualified requests cannot simply be dropped as lm-eval library requires the number of requests to be the same. - - Note: The final score will drop if disqualified requests exist. - """ - request_idx_to_replace = [] - qualified_requests = [] - updated_requests = [] - for i, request in enumerate(requests): - context, gen_kwargs = request.args - if len(self.tok_encode(context)) > self.max_context_length: - request_idx_to_replace.append(i) - else: - qualified_requests.append(request) - - assert len( - qualified_requests - ) > 1, "No requests with context length <= max_context_length. Cannot run the evaluation." - if len(request_idx_to_replace) > 0: - print( - f"Warning: {len(request_idx_to_replace)} requests with context length > max_context_length will be replaced. The final score will drop." - ) - - for i, request in enumerate(requests): - if i in request_idx_to_replace: - # Replace the requests with context length > max_context_length with the qualified requests - updated_requests.append( - qualified_requests[i % len(qualified_requests)]) - else: - updated_requests.append(request) - assert len( - updated_requests - ) == num_r, "Number of updated requests does not match the number of requests." - requests = updated_requests - - def _get_sp(gen_kwargs): - k_mapping = { - "temperature": "temperature", - "top_p": "top_p", - "max_gen_toks": "max_tokens", - "until": "stop", - } - kwargs_mapped = { - k_sp: gen_kwargs[k_gen] - for k_gen, k_sp in k_mapping.items() if k_gen in gen_kwargs - } - if "max_tokens" not in kwargs_mapped: - kwargs_mapped["max_tokens"] = self.max_gen_toks - return SamplingParams(**kwargs_mapped) - - # process requests - futures: Dict[int, RequestOutput] = {} - future_stop_words: Dict[int, RequestOutput] = {} - results = [] - for i, _ in tqdm(enumerate(requests), - desc=desc, - total=num_r, - disable=disable_tqdm): - # asynchronously submit a chunk of requests ahead of time... - if i % self.chunk_size == 0: - for j in range(i, min(i + self.chunk_size, num_r)): - context, gen_kwargs = requests[j].args - prompt_ids = self.tok_encode(context) - if self.max_context_length is not None: - assert len( - prompt_ids - ) <= self.max_context_length, f"Prompt length > {self.max_context_length}, {len(prompt_ids)}, should be filtered out." - kwargs_mapped = _get_sp(gen_kwargs) - futures[j] = self.llm.generate_async( - prompt_ids, kwargs_mapped) - del kwargs_mapped - future_stop_words[j] = gen_kwargs["until"] - - # process the output of the request i - r_out: RequestOutput = futures.pop(i).result() - stop_words = future_stop_words.pop(i) - txt = r_out.outputs[0].text - if stop_words: - for word in stop_words: - word_index = txt.find(word) - if word_index >= 0: - txt = txt[:word_index] - results.append(txt) - - return results - - -if __name__ == "__main__": - cli_evaluate() - # Force clean up the LLM instance and void hanging. - gc.collect() - - # Force terminate in case gc.collect() is not enough. - def _terminate(): - time.sleep(10) - os.kill(os.getpid(), signal.SIGTERM) - - termination_thread = threading.Thread(target=_terminate, daemon=True) - termination_thread.start() diff --git a/examples/llm-eval/lm-eval-harness/requirements.txt b/examples/llm-eval/lm-eval-harness/requirements.txt deleted file mode 100644 index 44a3383f612d..000000000000 --- a/examples/llm-eval/lm-eval-harness/requirements.txt +++ /dev/null @@ -1 +0,0 @@ -lm_eval[api]==0.4.7 diff --git a/examples/mmlu.py b/examples/mmlu.py deleted file mode 100644 index 82bb9a7ec9ad..000000000000 --- a/examples/mmlu.py +++ /dev/null @@ -1,479 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2020 Dan Hendrycks -# SPDX-FileCopyrightText: Copyright (c) 2023 Deep Cognition and Language Research (DeCLaRe) Lab -# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 and MIT -# -# 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. -"""Adapted from https://github.com/declare-lab/instruct-eval -Helper script to compare TRTLLM and HF models on the MMLU dataset. -Example usage: - mkdir data; wget https://people.eecs.berkeley.edu/~hendrycks/data.tar -O data/mmlu.tar - tar -xf data/mmlu.tar -C data && mv data/data data/mmlu - - python mmlu.py --hf_model_dir --engine_dir --test_trt_llm - python mmlu.py --hf_model_dir --engine_dir --test_hf -""" - -import argparse -import os -import random - -import numpy as np -import pandas as pd -import torch -import torch.nn as nn -from tqdm import tqdm -from transformers import (AutoConfig, AutoModel, AutoModelForCausalLM, - AutoModelForSeq2SeqLM, AutoTokenizer, - GenerationConfig) -from utils import (add_common_args, load_tokenizer, prepare_enc_dec_inputs, - read_is_enc_dec, read_model_name) - -import tensorrt_llm -from tensorrt_llm.runtime import PYTHON_BINDINGS, ModelRunner - -if PYTHON_BINDINGS: - from tensorrt_llm.runtime import ModelRunnerCpp - -os.environ["TOKENIZERS_PARALLELISM"] = "false" - -DTYPE_STR_MAPPING = { - "fp32": torch.float32, - "fp16": torch.float16, - "bf16": torch.bfloat16, - "float32": torch.float32, - "float16": torch.float16, - "bfloat16": torch.bfloat16, -} -RAND_SEED = 1234 - - -def get_choices(): - return ["A", "B", "C", "D"] - - -def get_subcategories(): - return { - "abstract_algebra": ["math"], - "anatomy": ["health"], - "astronomy": ["physics"], - "business_ethics": ["business"], - "clinical_knowledge": ["health"], - "college_biology": ["biology"], - "college_chemistry": ["chemistry"], - "college_computer_science": ["computer science"], - "college_mathematics": ["math"], - "college_medicine": ["health"], - "college_physics": ["physics"], - "computer_security": ["computer science"], - "conceptual_physics": ["physics"], - "econometrics": ["economics"], - "electrical_engineering": ["engineering"], - "elementary_mathematics": ["math"], - "formal_logic": ["philosophy"], - "global_facts": ["other"], - "high_school_biology": ["biology"], - "high_school_chemistry": ["chemistry"], - "high_school_computer_science": ["computer science"], - "high_school_european_history": ["history"], - "high_school_geography": ["geography"], - "high_school_government_and_politics": ["politics"], - "high_school_macroeconomics": ["economics"], - "high_school_mathematics": ["math"], - "high_school_microeconomics": ["economics"], - "high_school_physics": ["physics"], - "high_school_psychology": ["psychology"], - "high_school_statistics": ["math"], - "high_school_us_history": ["history"], - "high_school_world_history": ["history"], - "human_aging": ["health"], - "human_sexuality": ["culture"], - "international_law": ["law"], - "jurisprudence": ["law"], - "logical_fallacies": ["philosophy"], - "machine_learning": ["computer science"], - "management": ["business"], - "marketing": ["business"], - "medical_genetics": ["health"], - "miscellaneous": ["other"], - "moral_disputes": ["philosophy"], - "moral_scenarios": ["philosophy"], - "nutrition": ["health"], - "philosophy": ["philosophy"], - "prehistory": ["history"], - "professional_accounting": ["other"], - "professional_law": ["law"], - "professional_medicine": ["health"], - "professional_psychology": ["psychology"], - "public_relations": ["politics"], - "security_studies": ["politics"], - "sociology": ["culture"], - "us_foreign_policy": ["politics"], - "virology": ["health"], - "world_religions": ["philosophy"], - } - - -def get_categories(): - return { - "STEM": [ - "physics", - "chemistry", - "biology", - "computer science", - "math", - "engineering", - ], - "humanities": ["history", "philosophy", "law"], - "social sciences": [ - "politics", - "culture", - "economics", - "geography", - "psychology", - ], - "other (business, health, misc.)": ["other", "business", "health"], - } - - -def format_subject(subject): - line = subject.split("_") - s = "" - for entry in line: - s += " " + entry - return s - - -def format_example(df, idx, include_answer=True): - prompt = df.iloc[idx, 0] - k = df.shape[1] - 2 - for j in range(k): - prompt += "\n{}. {}".format(get_choices()[j], df.iloc[idx, j + 1]) - prompt += "\nAnswer:" - if include_answer: - prompt += " {}\n\n".format(df.iloc[idx, k + 1]) - return prompt - - -def gen_prompt(train_df, subject, k=-1): - prompt = "The following are multiple choice questions (with answers) about {}.\n\n".format( - format_subject(subject)) - if k == -1: - k = train_df.shape[0] - for i in range(k): - prompt += format_example(train_df, i) - return prompt - - -def evaluate(args, subject, pipeline, dev_df, test_df): - rank = tensorrt_llm.mpi_rank() - cors = [] - all_probs = [] - for i in range(test_df.shape[0]): - if i >= args.max_ite: - break - # get prompt and make sure it fits - k = args.ntrain - prompt_end = format_example(test_df, i, include_answer=False) - train_prompt = gen_prompt(dev_df, subject, k) - prompt = train_prompt + prompt_end - - while not pipeline.check_valid_length(prompt) and k > 0: - k -= 1 - train_prompt = gen_prompt(dev_df, subject, k) - prompt = train_prompt + prompt_end - - label = test_df.iloc[i, test_df.shape[1] - 1] - pred = pipeline(prompt) - - if rank == 0: - probs = [0 for _ in get_choices()] - cor = pred.strip().startswith(label) - cors.append(cor) - all_probs.append(probs) - - if rank == 0: - acc = np.mean(cors) - cors = np.array(cors) - - all_probs = np.array(all_probs) - print("Average accuracy {:.3f} - {}".format(acc, subject)) - - return cors, acc, all_probs - else: - return None, 0, None - - -def get_tokenizer(ckpt_path, max_seq_len): - print(f"Initializing tokenizer from {ckpt_path}") - tokenizer = AutoTokenizer.from_pretrained( - ckpt_path, - model_max_length=max_seq_len, - padding_side="left", - trust_remote_code=True, - ) - tokenizer.pad_token = tokenizer.eos_token - - return tokenizer - - -class Pipeline: - - def __init__(self, tokenizer, model, model_name, pad_id, end_id, - max_attention_window_size, is_enc_dec, hf_model_dir, - engine_dir): - self.tokenizer = tokenizer - self.model = model - self.model_name = model_name - self.pad_id = pad_id - self.end_id = end_id - self.max_attention_window_size = max_attention_window_size - self.output_len = 2 - self.is_enc_dec = is_enc_dec - self.decoder_start_token_id = None - self.engine_dir = engine_dir - if self.is_enc_dec: - self.decoder_start_token_id = AutoConfig.from_pretrained( - hf_model_dir).decoder_start_token_id - - def __call__(self, prompt): - rank = tensorrt_llm.mpi_rank() - # Run the model in batch size 1 and beam size 1 - inputs = self.tokenizer.encode(prompt, return_tensors="pt").squeeze(0) - batch_input_ids = [inputs] - - # For multi-choice tasks like MMLU, we don't need to adjust following parameters - output_len = self.output_len - top_k = 1 - top_p = 0.0 - - input_lengths = [x.size(0) for x in batch_input_ids] - - with torch.no_grad(): - if isinstance(self.model, nn.Module): - # Left padding for HF - max_length = max(input_lengths) - paddings = [ - torch.ones(max_length - l, dtype=torch.int32) * self.pad_id - for l in input_lengths - ] - batch_input_ids = [ - torch.cat([pad, x]) - for x, pad in zip(batch_input_ids, paddings) - ] - batch_input_ids = torch.stack(batch_input_ids) - batch_input_ids = batch_input_ids.cuda() - if self.is_enc_dec: - batch_decoder_input_ids = torch.IntTensor( - [[self.decoder_start_token_id]]).to('cuda') - batch_decoder_input_ids = batch_decoder_input_ids.repeat( - (batch_input_ids.shape[0], 1)) - - with torch.no_grad(): - # Use default temperature and top_k - outputs = self.model.generate( - batch_input_ids, - max_new_tokens=output_len, - top_k=top_k, - decoder_input_ids=batch_decoder_input_ids - if self.is_enc_dec else None) - if not self.is_enc_dec: - output_ids = outputs[0, input_lengths[0]:] - else: - output_ids = outputs[0] - - elif isinstance(self.model, ModelRunnerCpp) or isinstance( - self.model, ModelRunner): - if self.is_enc_dec: - encoder_input_ids, encoder_input_features, encoder_output_lengths, decoder_input_ids = prepare_enc_dec_inputs( - batch_input_ids, self.model_name, self.engine_dir, None) - - outputs = self.model.generate( - batch_input_ids=decoder_input_ids - if self.is_enc_dec else batch_input_ids, - encoder_input_ids=encoder_input_ids - if self.is_enc_dec else None, - encoder_input_features=encoder_input_features - if self.is_enc_dec else None, - encoder_output_lengths=encoder_output_lengths - if self.is_enc_dec else None, - max_new_tokens=output_len, - max_attention_window_size=self.max_attention_window_size, - end_id=self.end_id, - pad_id=self.pad_id, - top_k=top_k, - top_p=top_p, - ) - torch.cuda.synchronize() - if rank == 0: - if not self.is_enc_dec: - output_ids = outputs[0, 0, input_lengths[0]:] - else: - output_ids = outputs[0, 0] - if rank == 0: - return self.tokenizer.decode(output_ids, skip_special_tokens=True) - else: - return None - - def check_valid_length(self, prompt): - if isinstance(self.model, nn.Module): - return True - input_len = len(self.tokenizer.encode(prompt)) - return input_len <= self.model.max_input_len and input_len + self.output_len <= self.model.max_seq_len - - -def parse_args(): - parser = argparse.ArgumentParser() - parser.add_argument( - "--data_dir", - type=str, - default="data/mmlu", - help=("Path to the data directory. If not available, " - "download https://people.eecs.berkeley.edu/~hendrycks/data.tar"), - ) - parser.add_argument("--ntrain", type=int, default=5) - parser.add_argument("--max_input_length", type=int, default=2048) - parser.add_argument("--test_trt_llm", action="store_true") - parser.add_argument("--test_hf", action="store_true") - parser.add_argument('--check_accuracy', action='store_true') - parser.add_argument('--accuracy_threshold', type=float, default=30) - parser.add_argument('--max_ite', type=int, default=10000000) - parser = add_common_args(parser) - - args = parser.parse_args() - - return args - - -def main(): - args = parse_args() - if args.tokenizer_dir is None: - args.tokenizer_dir = args.hf_model_dir - random.seed(RAND_SEED) - np.random.seed(RAND_SEED) - runtime_rank = tensorrt_llm.mpi_rank() - - os.path.dirname(os.path.abspath(__file__)) - data_fullpath = os.path.join(args.data_dir, "test") - - subjects = sorted([ - f.split("_test.csv")[0] for f in os.listdir(data_fullpath) - if "_test.csv" in f - ]) - - all_cors = [] - subcat_cors = { - subcat: [] - for subcat_lists in get_subcategories().values() - for subcat in subcat_lists - } - cat_cors = {cat: [] for cat in get_categories()} - - # different handling if encoder-decoder models - is_enc_dec = read_is_enc_dec( - args.engine_dir if not args.test_hf else args.hf_model_dir, - args.test_hf) - - model_name, model_version = read_model_name( - (args.engine_dir if not is_enc_dec else os.path.join( - args.engine_dir, 'encoder')) - if not args.test_hf else args.hf_model_dir, args.test_hf) - - tokenizer, pad_id, end_id = load_tokenizer( - tokenizer_dir=args.tokenizer_dir, - vocab_file=args.vocab_file, - model_name=model_name, - model_version=model_version, - ) - - if args.test_trt_llm: - assert not args.test_hf, "Cannot test both TRT-LLM and HF" - runner_cls = ModelRunner if not PYTHON_BINDINGS else ModelRunnerCpp - runner_kwargs = {} - if PYTHON_BINDINGS: - runner_kwargs.update(max_beam_width=1) - runner_kwargs.update( - is_enc_dec=is_enc_dec, - max_tokens_in_paged_kv_cache=args.max_tokens_in_paged_kv_cache, - kv_cache_enable_block_reuse=args.kv_cache_enable_block_reuse, - kv_cache_free_gpu_memory_fraction=args. - kv_cache_free_gpu_memory_fraction, - cross_kv_cache_fraction=args.cross_kv_cache_fraction - if is_enc_dec else None, - enable_chunked_context=args.enable_chunked_context, - multi_block_mode=args.multi_block_mode) - model = runner_cls.from_dir(engine_dir=args.engine_dir, - rank=runtime_rank, - **runner_kwargs) - else: - assert args.test_hf, "Must test either TRT-LLM or HF" - if 'GLM' in model_name and model_version == 'glm': - auto_model_cls = AutoModelForSeq2SeqLM - elif 'GLM' in model_name and model_version == 'chatglm': - auto_model_cls = AutoModel - elif is_enc_dec: - auto_model_cls = AutoModelForSeq2SeqLM - else: - auto_model_cls = AutoModelForCausalLM - model = auto_model_cls.from_pretrained( - args.hf_model_dir, - trust_remote_code=True, - dtype=DTYPE_STR_MAPPING[args.hf_data_type], - device_map="auto" if args.hf_device_map_auto else None, - ) - if not args.hf_device_map_auto: - model.cuda() - if model_name == "qwen": - model.generation_config = GenerationConfig.from_pretrained( - args.hf_model_dir, trust_remote_code=True) - - pipeline = Pipeline(tokenizer, model, model_name, pad_id, end_id, - args.max_attention_window_size, is_enc_dec, - args.hf_model_dir, args.engine_dir) - - for subject in tqdm(subjects): - dev_df = pd.read_csv(os.path.join(args.data_dir, "dev", - subject + "_dev.csv"), - header=None)[:args.ntrain] - test_df = pd.read_csv(os.path.join(args.data_dir, "test", - subject + "_test.csv"), - header=None) - - cors, acc, probs = evaluate(args, subject, pipeline, dev_df, test_df) - subcats = get_subcategories()[subject] - for subcat in subcats: - subcat_cors[subcat].append(cors) - for key in get_categories().keys(): - if subcat in get_categories()[key]: - cat_cors[key].append(cors) - all_cors.append(cors) - - if runtime_rank == 0: - for subcat in subcat_cors: - acc = np.mean(np.concatenate(subcat_cors[subcat])) * 100 - print(f"Average accuracy {acc:.2f} - {subcat}") - - for cat in cat_cors: - acc = np.mean(np.concatenate(cat_cors[cat])) * 100 - print(f"Average accuracy {acc:.2f} - {cat}") - - weighted_acc = np.mean(np.concatenate(all_cors)) * 100 - print(f"MMLU weighted average accuracy: {weighted_acc:.2f}") - - if args.check_accuracy: - assert weighted_acc >= args.accuracy_threshold, f"Expected accuracy >= {args.accuracy_threshold} while got {weighted_acc}" - return weighted_acc - - -if __name__ == "__main__": - main() diff --git a/examples/models/contrib/arctic/README.md b/examples/models/contrib/arctic/README.md deleted file mode 100644 index d346a3644517..000000000000 --- a/examples/models/contrib/arctic/README.md +++ /dev/null @@ -1,94 +0,0 @@ -# Arctic - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a [Arctic](https://huggingface.co/Snowflake/snowflake-arctic-instruct) model in TensorRT-LLM. - -The TensorRT LLM Arctic implementation is based on the LLaMA model, with Mixture of Experts (MoE) enabled. The implementation can -be found in [llama/model.py](../../../../tensorrt_llm/models/llama/model.py). -See the LLaMA example [`examples/models/core/llama`](../../core/llama) for details. - -- [Arctic](#arctic) - - [Download model checkpoints](#download-model-checkpoints) - - [TensorRT LLM workflow](#tensorrt-llm-workflow) - - [Apply FP8 PTQ](#apply-fp8-ptq) - - [Build TensorRT engine](#build-tensorrt-engine) - - [Run Engine](#run-engine) - - [OOTB](#ootb) - -## Download model checkpoints - -First, download the HuggingFace BF16 checkpoints of Arctic model. - -**CAVEAT: this model is a pretty large Mixture-of-Experts (MoE) model, which has nearly 500B parameters and requires around 900GB disk space for storage. Please make sure you have enough space before proceeding.** - -```bash -HF_MODEL="arctic" -git clone https://huggingface.co/Snowflake/snowflake-arctic-instruct tmp/hf_checkpoints/${HF_MODEL} - -``` - -## TensorRT LLM workflow -Next, we use the general quantization script `quantize.py` to convert the checkpoints in FP8, and build the model with `trtllm-build` on multi-GPUs. In the example below, we use Tensor Parallelism (TP) across 8 GPUs. - -**Note: for such large model, it is deemed necessary to apply Post-Training Quantization (PTQ) methods on the model weights to deploy it on a cluster node, e.g., 8xH100 GPUs. In this example, we demonstrate the FP8 quantization workflow, which is supported on Hopper-and-next GPU architectures. For instructions of other PTQ methods other than FP8, please refer to the LLaMA or Mixtral examples.** - - -Set environment variables and necessary directory: - -```bash -PREC_RAW="bfloat16" -PREC_QUANT="fp8" -TP=8 -ENGINE="${HF_MODEL}_${PREC_QUANT}_tp${TP}" - -mkdir -p tmp/trt_engines -``` - -### Apply FP8 PTQ - -Notes: -- currently quantize.py does not support for Expert Parallelism (EP) mode yet. User should use `../../core/llama/convert_checkpoint.py` and specify `--moe_ep_size 1` instead, if needed. -- TensorRT LLM uses static quantization methods, which is expected to be faster at runtime as compared to dynamic quantization methods. This comes at a cost of an offline calibration step during quantization. `batch_size` and `calib_size` can be adjusted to shorten the calibration time. Please refer to ../quantization/README.md for explanation. -- **due to the large model size and the calibration step (which has to load the HuggingFace model and run forward passes), it is likely that you will need more number of GPUs during quantization step than the number of GPUs for engine building and final deployment. For example, using 16xH100 or 8xH200 for quantization & 8xH100 for deployment.** - -```bash -python ../../../quantization/quantize.py --model_dir tmp/hf_checkpoints/${HF_MODEL} \ - --dtype ${PREC_RAW} \ - --qformat ${PREC_QUANT} \ - --kv_cache_dtype ${PREC_QUANT} \ - --output_dir tmp/tllm_checkpoints/${ENGINE} \ - --batch_size 1 \ - --calib_size 128 \ - --tp_size ${TP} |& tee tmp/trt_engines/${ENGINE}_quantize.log - -``` - -### Build TensorRT engine -```bash -# Enable fp8 context fmha to get further acceleration by setting `--use_fp8_context_fmha enable` -# Use --workers to enable parallel build -trtllm-build --checkpoint_dir ./tmp/tllm_checkpoints/${ENGINE} \ - --output_dir ./tmp/trt_engines/${ENGINE} \ - --gpt_attention_plugin ${PREC_RAW} \ - --gemm_plugin ${PREC_RAW} \ - --workers ${TP} |& tee tmp/trt_engines/${ENGINE}_build.log -``` - -### Run Engine -Test your engine with the [run.py](../../../run.py) script: - -```bash -mpirun -n ${TP} --allow-run-as-root python ../../../run.py --engine_dir ./tmp/trt_engines/${ENGINE} --tokenizer_dir tmp/hf_checkpoints/${HF_MODEL} --max_output_len 20 --input_text "The future of AI is" |& tee tmp/trt_engines/${ENGINE}_run.log -``` - -For more examples see [`examples/models/core/llama/README.md`](../../core/llama/README.md) - - -### OOTB - -Arctic supports OOTB operation without the plugin, however this comes at a significant performance cost. Users should prefer using the plugin path whenever possible. diff --git a/examples/models/contrib/blip2/README.md b/examples/models/contrib/blip2/README.md deleted file mode 100644 index 0e80ed4fd81b..000000000000 --- a/examples/models/contrib/blip2/README.md +++ /dev/null @@ -1,4 +0,0 @@ -This example has been moved to [`../multimodal`](../../../multimodal) and merged with other multimodal examples. -Please follow [`../multimodal/README.md`](../../../multimodal/README.md) for new instructions. - -**NOTICE:** This folder will be removed in v1.0 release. diff --git a/examples/models/contrib/chatglm-6b/README.md b/examples/models/contrib/chatglm-6b/README.md deleted file mode 100644 index 00a8c2424665..000000000000 --- a/examples/models/contrib/chatglm-6b/README.md +++ /dev/null @@ -1,112 +0,0 @@ -# ChatGLM - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document explains how to build the [ChatGLM-6B](https://huggingface.co/THUDM/chatglm-6b) models using TensorRT LLM and run on a single GPU, a single node with multiple GPUs or multiple nodes with multiple GPUs. - -- [ChatGLM](#chatglm) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Model comparison](#model-comparison) - - [Tokenizer and special tokens comparison](#tokenizer-and-special-tokens-comparison) - - [Usage](#usage) - - [1. Download repo and weights from HuggingFace Transformers](#1-download-repo-and-weights-from-huggingface-transformers) - - [2. Convert weights from HF Transformers to TensorRT LLM format](#2-convert-weights-from-hf-transformers-to-tensorrt-llm-format) - - [3. Build TensorRT engine(s)](#3-build-tensorrt-engines) - - [Enable plugins](#enable-plugins) - - [In-flight batching](#in-flight-batching) - - [4. Run inference](#4-run-inference) - - [Single node, single GPU](#single-node-single-gpu) - - [Single node, multi GPU](#single-node-multi-gpu) - - [5. Run summarization task](#5-run-summarization-task) - - [Weight Only quantization](#weight-only-quantization) - - [Smooth Quantization (SQ)](#smooth-quantization-sq) - - [Activation-aware Weight Quantization (AWQ)](#activation-aware-weight-quantization-awq) - - [FP8 Quantization](#fp8-quantization) - - [Benchmark](#benchmark) - - -## Overview - -The TensorRT LLM ChatGLM implementation can be found in [`tensorrt_llm/models/chatglm/model.py`](../../tensorrt_llm/models/chatglm/model.py). -The TensorRT LLM ChatGLM example code is located in [`examples/models/contrib/chatglm-6b`](./). There is one main file: - -* [`examples/models/core/glm-4-9b/convert_checkpoint.py`](../../../glm-4-9b/convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - -| Model Name | FP16 | FMHA | WO | SQ | AWQ | FP8 | TP | PP | ST | C++ | benchmark | IFB | -| :--------------: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :-------: | :---: | -| chatglm_6b | Y | | Y | | | | Y | | Y | Y | Y | Y | -| glm_10b | Y | Y | Y | | Y | Y | Y | | Y | Y | Y | Y | - -* Model Name: the name of the model, the same as the name on HuggingFace -* FMHA: Fused MultiHead Attention (see introduction below) -* WO: Weight Only Quantization (int8 / int4) -* SQ: Smooth Quantization (int8) -* AWQ: Activation Aware Weight Quantization (int4) -* FP8: FP8 Quantization -* TP: Tensor Parallel -* PP: Pipeline Parallel -* ST: Strongly Typed -* C++: C++ Runtime -* benchmark: benchmark by python / C++ Runtime -* IFB: In-flight Batching (see introduction below) - -## Model comparison - -| Name | nL | nAH | nKH | nHW | nH | nF | nMSL | nV | bP2D | bBQKV | bBDense | Comments | -| :--------------: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :----: | :---: | :---: | :-----: | :----------------------------------------------------------------- | -| chatglm_6b | 28 | 32 | 32 | 128 | 4096 | 16384 | 2048 | 130528 | Y | Y | Y | | -| glm_10b | 48 | 64 | 32 | 64 | 4096 | 16384 | 1024 | 50304 | Y | Y | Y | | - -* nL: number of layers -* nAH: number of attention heads -* nKH: number of kv heads (less than nAH if multi_query_attention is used) -* nHW: head width -* nH: hidden size -* nF: FFN hidden size -* nMSL: max sequence length (input + output) -* nV: vocabulary size -* bP2D: use position_encoding_2d (Y: Yes, N: No) -* bBQKV: use bias for QKV multiplication in self-attention -* bBDense: use bias for Dense multiplication in self-attention - -## Tokenizer and special tokens comparison - -| Name | Tokenizer | bos | eos | pad | cls | startofpiece | endofpiece | mask | smask | gmask | -| :--------------: | :--------------: | :----: | :----: | :---: | :---: | :----------: | :--------: | :----: | :---: | :----: | -| chatglm_6b | ChatGLMTokenizer | 130004 | 130005 | 3 | | 130004 | 130005 | 130000 | | 130001 | -| glm_10b | GLMGPT2Tokenizer | 50257 | 50256 | 50256 | 50259 | 50257 | 50258 | 50260 | 50264 | 50263 | - -## Usage - -The next section describe how to build the engine and run the inference demo. - -### 1. Download repo and weights from HuggingFace Transformers - -```bash -pip install -r requirements.txt -apt-get update -apt-get install git-lfs -rm -rf chatglm* - -# clone one or more models we want to build -git clone https://huggingface.co/THUDM/chatglm-6b chatglm_6b -git clone https://huggingface.co/THUDM/glm-10b glm_10b - -# replace tokenization file if using transformers-4.36.1 for model ChatGLM-6B (this might be needless in the future) -cp chatglm_6b/tokenization_chatglm.py chatglm_6b/tokenization_chatglm.py-backup -cp tokenization_chatglm.py chatglm_6b -``` - -For more example codes, please refer to the [examples/models/core/glm-4-9b/README.md](../../../glm-4-9b/README.md). diff --git a/examples/models/contrib/chatglm-6b/requirements.txt b/examples/models/contrib/chatglm-6b/requirements.txt deleted file mode 100644 index cdc65bf2bb38..000000000000 --- a/examples/models/contrib/chatglm-6b/requirements.txt +++ /dev/null @@ -1,8 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -protobuf -rouge_score -sentencepiece -tiktoken diff --git a/examples/models/contrib/chatglm-6b/tokenization_chatglm.py b/examples/models/contrib/chatglm-6b/tokenization_chatglm.py deleted file mode 100755 index 8ae124909e98..000000000000 --- a/examples/models/contrib/chatglm-6b/tokenization_chatglm.py +++ /dev/null @@ -1,467 +0,0 @@ -"""Tokenization classes for ChatGLM.""" -import os -from typing import Dict, List, Optional, Union - -import numpy as np -import sentencepiece as spm -from transformers.tokenization_utils import PreTrainedTokenizer -from transformers.tokenization_utils_base import BatchEncoding, EncodedInput -from transformers.utils import PaddingStrategy, logging - -logger = logging.get_logger(__name__) - -PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = { - "THUDM/chatglm-6b": 2048, -} - - -class TextTokenizer: - - def __init__(self, model_path): - self.sp = spm.SentencePieceProcessor() - self.sp.Load(model_path) - self.num_tokens = self.sp.vocab_size() - - def encode(self, text): - return self.sp.EncodeAsIds(text) - - def decode(self, ids: List[int]): - return self.sp.DecodeIds(ids) - - def tokenize(self, text): - return self.sp.EncodeAsPieces(text) - - def convert_tokens_to_string(self, tokens): - return self.sp.DecodePieces(tokens) - - def convert_tokens_to_ids(self, tokens): - return [self.sp.PieceToId(token) for token in tokens] - - def convert_token_to_id(self, token): - return self.sp.PieceToId(token) - - def convert_id_to_token(self, idx): - return self.sp.IdToPiece(idx) - - def __len__(self): - return self.num_tokens - - -class SPTokenizer: - - def __init__( - self, - vocab_file, - num_image_tokens=20000, - max_blank_length=80, - byte_fallback=True, - ): - assert vocab_file is not None - self.vocab_file = vocab_file - self.num_image_tokens = num_image_tokens - self.special_tokens = [ - "[MASK]", "[gMASK]", "[sMASK]", "", "", "", - "", "" - ] - self.max_blank_length = max_blank_length - self.byte_fallback = byte_fallback - self.text_tokenizer = TextTokenizer(vocab_file) - - def _get_text_tokenizer(self): - return self.text_tokenizer - - @staticmethod - def get_blank_token(length: int): - assert length >= 2 - return f"<|blank_{length}|>" - - @staticmethod - def get_tab_token(): - return f"<|tab|>" - - @property - def num_text_tokens(self): - return self.text_tokenizer.num_tokens - - @property - def num_tokens(self): - return self.num_image_tokens + self.num_text_tokens - - @staticmethod - def _encode_whitespaces(text: str, max_len: int = 80): - text = text.replace("\t", SPTokenizer.get_tab_token()) - for i in range(max_len, 1, -1): - text = text.replace(" " * i, SPTokenizer.get_blank_token(i)) - return text - - def _preprocess(self, text: str, linebreak=True, whitespaces=True): - if linebreak: - text = text.replace("\n", "") - if whitespaces: - text = self._encode_whitespaces(text, max_len=self.max_blank_length) - return text - - def encode(self, - text: str, - linebreak=True, - whitespaces=True, - add_dummy_prefix=True) -> List[int]: - """ - @param text: Text to encode. - @param linebreak: Whether to encode newline (\n) in text. - @param whitespaces: Whether to encode multiple whitespaces or tab in text, useful for source code encoding. - @param special_tokens: Whether to encode special token ([MASK], [gMASK], etc.) in text. - @param add_dummy_prefix: Whether to add dummy blank space in the beginning. - """ - text = self._preprocess(text, linebreak, whitespaces) - if not add_dummy_prefix: - text = "" + text - tmp = self._get_text_tokenizer().encode(text) - tokens = [x + self.num_image_tokens for x in tmp] - return tokens if add_dummy_prefix else tokens[2:] - - def postprocess(self, text): - text = text.replace("", "\n") - text = text.replace(SPTokenizer.get_tab_token(), "\t") - for i in range(2, self.max_blank_length + 1): - text = text.replace(self.get_blank_token(i), " " * i) - return text - - def decode(self, text_ids: List[int]) -> str: - ids = [int(_id) - self.num_image_tokens for _id in text_ids] - ids = [_id for _id in ids if _id >= 0] - text = self._get_text_tokenizer().decode(ids) - text = self.postprocess(text) - return text - - def decode_tokens(self, tokens: List[str]) -> str: - text = self._get_text_tokenizer().convert_tokens_to_string(tokens) - text = self.postprocess(text) - return text - - def tokenize(self, - text: str, - linebreak=True, - whitespaces=True, - add_dummy_prefix=True) -> List[str]: - """ - @param text: Text to encode. - @param linebreak: Whether to encode newline (\n) in text. - @param whitespaces: Whether to encode multiple whitespaces or tab in text, useful for source code encoding. - @param special_tokens: Whether to encode special token ([MASK], [gMASK], etc.) in text. - @param add_dummy_prefix: Whether to add dummy blank space in the beginning. - """ - text = self._preprocess(text, linebreak, whitespaces) - if not add_dummy_prefix: - text = "" + text - tokens = self._get_text_tokenizer().tokenize(text) - return tokens if add_dummy_prefix else tokens[2:] - - def __getitem__(self, x: Union[int, str]): - if isinstance(x, int): - if x < self.num_image_tokens: - return "".format(x) - else: - return self.text_tokenizer.convert_id_to_token( - x - self.num_image_tokens) - elif isinstance(x, str): - if x.startswith("") and x[7:-1].isdigit(): - return int(x[7:-1]) - else: - return self.text_tokenizer.convert_token_to_id( - x) + self.num_image_tokens - else: - raise ValueError("The key should be str or int.") - - -class ChatGLMTokenizer(PreTrainedTokenizer): - """ - Construct a ChatGLM tokenizer. Based on byte-level Byte-Pair-Encoding. - - Args: - vocab_file (`str`): - Path to the vocabulary file. - """ - - vocab_files_names = {"vocab_file": "ice_text.model"} - max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES - model_input_names = ["input_ids", "attention_mask", "position_ids"] - - def __init__(self, - vocab_file, - do_lower_case=False, - remove_space=False, - bos_token='', - eos_token='', - end_token='', - mask_token='[MASK]', - gmask_token='[gMASK]', - padding_side="left", - pad_token="", - unk_token="", - num_image_tokens=20000, - **kwargs) -> None: # Fix for new transformers - - self.do_lower_case = do_lower_case - self.remove_space = remove_space - self.vocab_file = vocab_file - - self.bos_token = bos_token - self.eos_token = eos_token - self.end_token = end_token - self.mask_token = mask_token - self.gmask_token = gmask_token - - self.sp_tokenizer = SPTokenizer(vocab_file, - num_image_tokens=num_image_tokens) - - super().__init__( - do_lower_case=do_lower_case, # Fix for new transformers - remove_space=remove_space, - padding_side=padding_side, - bos_token=bos_token, - eos_token=eos_token, - end_token=end_token, - mask_token=mask_token, - gmask_token=gmask_token, - pad_token=pad_token, - unk_token=unk_token, - num_image_tokens=num_image_tokens, - **kwargs) - """ Initialization """ - - @property - def gmask_token_id(self) -> Optional[int]: - if self.gmask_token is None: - return None - return self.convert_tokens_to_ids(self.gmask_token) - - @property - def end_token_id(self) -> Optional[int]: - """ - `Optional[int]`: Id of the end of context token in the vocabulary. Returns `None` if the token has not been - set. - """ - if self.end_token is None: - return None - return self.convert_tokens_to_ids(self.end_token) - - @property - def vocab_size(self): - """ Returns vocab size """ - return self.sp_tokenizer.num_tokens - - def get_vocab(self): - """ Returns vocab as a dict """ - vocab = { - self._convert_id_to_token(i): i - for i in range(self.vocab_size) - } - vocab.update(self.added_tokens_encoder) - return vocab - - def preprocess_text(self, inputs): - if self.remove_space: - outputs = " ".join(inputs.strip().split()) - else: - outputs = inputs - - if self.do_lower_case: - outputs = outputs.lower() - - return outputs - - def _tokenize(self, text, **kwargs): - """ Returns a tokenized string. """ - text = self.preprocess_text(text) - - seq = self.sp_tokenizer.tokenize(text) - - return seq - - def convert_tokens_to_string(self, tokens: List[str]) -> str: - return self.sp_tokenizer.decode_tokens(tokens) - - def _decode(self, token_ids: Union[int, List[int]], **kwargs) -> str: - if isinstance(token_ids, int): - token_ids = [token_ids] - if len(token_ids) == 0: - return "" - if self.pad_token_id in token_ids: # remove pad - token_ids = list(filter((self.pad_token_id).__ne__, token_ids)) - return super()._decode(token_ids, **kwargs) - - def _convert_token_to_id(self, token): - """ Converts a token (str) in an id using the vocab. """ - return self.sp_tokenizer[token] - - def _convert_id_to_token(self, index): - """Converts an index (integer) in a token (str) using the vocab.""" - return self.sp_tokenizer[index] - - def save_vocabulary(self, save_directory, filename_prefix=None): - """ - Save the vocabulary and special tokens file to a directory. - - Args: - save_directory (`str`): - The directory in which to save the vocabulary. - filename_prefix (`str`, *optional*): - An optional prefix to add to the named of the saved files. - - Returns: - `Tuple(str)`: Paths to the files saved. - """ - if os.path.isdir(save_directory): - vocab_file = os.path.join(save_directory, - self.vocab_files_names["vocab_file"]) - else: - vocab_file = save_directory - - with open(self.vocab_file, 'rb') as fin: - proto_str = fin.read() - - with open(vocab_file, "wb") as writer: - writer.write(proto_str) - - return (vocab_file, ) - - def build_inputs_with_special_tokens( - self, - token_ids_0: List[int], - token_ids_1: Optional[List[int]] = None) -> List[int]: - """ - Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and - adding special tokens. A BERT sequence has the following format: - - - single sequence: `[CLS] X [SEP]` - - pair of sequences: `[CLS] A [SEP] B [SEP]` - - Args: - token_ids_0 (`List[int]`): - List of IDs to which the special tokens will be added. - token_ids_1 (`List[int]`, *optional*): - Optional second list of IDs for sequence pairs. - - Returns: - `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens. - """ - gmask_id = self.sp_tokenizer[self.gmask_token] - eos_id = self.sp_tokenizer[self.eos_token] - token_ids_0 = token_ids_0 + [ - gmask_id, self.sp_tokenizer[self.bos_token] - ] - if token_ids_1 is not None: - token_ids_0 = token_ids_0 + token_ids_1 + [eos_id] - return token_ids_0 - - def _pad( - self, - encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding], - max_length: Optional[int] = None, - padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, - pad_to_multiple_of: Optional[int] = None, - return_attention_mask: Optional[bool] = None, - padding_side: str = "left", # Fix for new transformers - ) -> dict: - """ - Pad encoded inputs (on left/right and up to predefined length or max length in the batch) - - Args: - encoded_inputs: - Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`). - max_length: maximum length of the returned list and optionally padding length (see below). - Will truncate by taking into account the special tokens. - padding_strategy: PaddingStrategy to use for padding. - - - PaddingStrategy.LONGEST Pad to the longest sequence in the batch - - PaddingStrategy.MAX_LENGTH: Pad to the max length (default) - - PaddingStrategy.DO_NOT_PAD: Do not pad - The tokenizer padding sides are defined in self.padding_side: - - - 'left': pads on the left of the sequences - - 'right': pads on the right of the sequences - pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value. - This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability - `>= 7.5` (Volta). - return_attention_mask: - (optional) Set to False to avoid returning attention mask (default: set to model specifics) - """ - # Load from model defaults - bos_token_id = self.sp_tokenizer[self.bos_token] - mask_token_id = self.sp_tokenizer[self.mask_token] - gmask_token_id = self.sp_tokenizer[self.gmask_token] - assert self.padding_side == "left" - - required_input = encoded_inputs[self.model_input_names[0]] - seq_length = len(required_input) - - if padding_strategy == PaddingStrategy.LONGEST: - max_length = len(required_input) - - if max_length is not None and pad_to_multiple_of is not None and ( - max_length % pad_to_multiple_of != 0): - max_length = ( - (max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of - - needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len( - required_input) != max_length - - # Initialize attention mask if not present. - if max_length is not None: - if "attention_mask" not in encoded_inputs: - if bos_token_id in required_input: - context_length = required_input.index(bos_token_id) - else: - context_length = seq_length - attention_mask = np.ones((1, seq_length, seq_length)) - attention_mask = np.tril(attention_mask) - attention_mask[:, :, :context_length] = 1 - attention_mask = np.bool_(attention_mask < 0.5) - encoded_inputs["attention_mask"] = attention_mask - - if "position_ids" not in encoded_inputs: - if bos_token_id in required_input: - context_length = required_input.index(bos_token_id) - else: - context_length = seq_length - position_ids = np.arange(seq_length, dtype=np.int64) - mask_token = mask_token_id if mask_token_id in required_input else gmask_token_id - if mask_token in required_input: - mask_position = required_input.index(mask_token) - position_ids[context_length:] = mask_position - block_position_ids = np.concatenate([ - np.zeros(context_length, dtype=np.int64), - np.arange(1, - seq_length - context_length + 1, - dtype=np.int64) - ]) - encoded_inputs["position_ids"] = np.stack( - [position_ids, block_position_ids], axis=0) - - if needs_to_be_padded: - difference = max_length - len(required_input) - - if "attention_mask" in encoded_inputs: - encoded_inputs["attention_mask"] = np.pad( - encoded_inputs["attention_mask"], - pad_width=[(0, 0), (difference, 0), (difference, 0)], - mode='constant', - constant_values=True) - if "token_type_ids" in encoded_inputs: - encoded_inputs["token_type_ids"] = [ - self.pad_token_type_id - ] * difference + encoded_inputs["token_type_ids"] - if "special_tokens_mask" in encoded_inputs: - encoded_inputs["special_tokens_mask"] = [ - 1 - ] * difference + encoded_inputs["special_tokens_mask"] - if "position_ids" in encoded_inputs: - encoded_inputs["position_ids"] = np.pad( - encoded_inputs["position_ids"], - pad_width=[(0, 0), (difference, 0)]) - encoded_inputs[self.model_input_names[ - 0]] = [self.pad_token_id] * difference + required_input - - return encoded_inputs diff --git a/examples/models/contrib/chatglm2-6b/README.md b/examples/models/contrib/chatglm2-6b/README.md deleted file mode 100644 index c25818e3f95c..000000000000 --- a/examples/models/contrib/chatglm2-6b/README.md +++ /dev/null @@ -1,108 +0,0 @@ -# ChatGLM - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document explains how to build the [ChatGLM2-6B](https://huggingface.co/THUDM/chatglm2-6b), [ChatGLM2-6B-32k](https://huggingface.co/THUDM/chatglm2-6b-32k) models using TensorRT LLM and run on a single GPU, a single node with multiple GPUs or multiple nodes with multiple GPUs. - -- [ChatGLM](#chatglm) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Model comparison](#model-comparison) - - [Tokenizer and special tokens comparison](#tokenizer-and-special-tokens-comparison) - - [Usage](#usage) - - [1. Download repo and weights from HuggingFace Transformers](#1-download-repo-and-weights-from-huggingface-transformers) - - [2. Convert weights from HF Transformers to TensorRT LLM format](#2-convert-weights-from-hf-transformers-to-tensorrt-llm-format) - - [3. Build TensorRT engine(s)](#3-build-tensorrt-engines) - - [Enable plugins](#enable-plugins) - - [In-flight batching](#in-flight-batching) - - [4. Run inference](#4-run-inference) - - [Single node, single GPU](#single-node-single-gpu) - - [Single node, multi GPU](#single-node-multi-gpu) - - [5. Run summarization task](#5-run-summarization-task) - - [Weight Only quantization](#weight-only-quantization) - - [Smooth Quantization (SQ)](#smooth-quantization-sq) - - [Activation-aware Weight Quantization (AWQ)](#activation-aware-weight-quantization-awq) - - [FP8 Quantization](#fp8-quantization) - - [Benchmark](#benchmark) - - -## Overview - -The TensorRT LLM ChatGLM implementation can be found in [`tensorrt_llm/models/chatglm/model.py`](../../tensorrt_llm/models/chatglm/model.py). -The TensorRT LLM ChatGLM example code is located in [`examples/models/contrib/chatglm2-6b`](./). There is one main file: - -* [`examples/models/core/glm-4-9b/convert_checkpoint.py`](../../../glm-4-9b/convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - -| Model Name | FP16 | FMHA | WO | SQ | AWQ | FP8 | TP | PP | ST | C++ | benchmark | IFB | -| :--------------: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :-------: | :---: | -| chatglm2_6b | Y | Y | Y | Y | Y | Y | Y | | Y | Y | Y | Y | -| chatglm2_6b_32k | Y | Y | Y | | Y | Y | Y | | Y | Y | Y | Y | - -* Model Name: the name of the model, the same as the name on HuggingFace -* FMHA: Fused MultiHead Attention (see introduction below) -* WO: Weight Only Quantization (int8 / int4) -* SQ: Smooth Quantization (int8) -* AWQ: Activation Aware Weight Quantization (int4) -* FP8: FP8 Quantization -* TP: Tensor Parallel -* PP: Pipeline Parallel -* ST: Strongly Typed -* C++: C++ Runtime -* benchmark: benchmark by python / C++ Runtime -* IFB: In-flight Batching (see introduction below) - -## Model comparison - -| Name | nL | nAH | nKH | nHW | nH | nF | nMSL | nV | bP2D | bBQKV | bBDense | Comments | -| :--------------: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :----: | :---: | :---: | :-----: | :----------------------------------------------------------------- | -| chatglm2_6b | 28 | 32 | 2 | 128 | 4096 | 13696 | 32768 | 65024 | N | Y | N | Multi_query_attention, RMSNorm rather than LayerNorm in chatglm_6b | -| chatglm2_6b_32k | 28 | 32 | 2 | 128 | 4096 | 13696 | 32768 | 65024 | N | Y | N | RoPE base=160000 rather than 10000 in chatglm2_6b | - -* nL: number of layers -* nAH: number of attention heads -* nKH: number of kv heads (less than nAH if multi_query_attention is used) -* nHW: head width -* nH: hidden size -* nF: FFN hidden size -* nMSL: max sequence length (input + output) -* nV: vocabulary size -* bP2D: use position_encoding_2d (Y: Yes, N: No) -* bBQKV: use bias for QKV multiplication in self-attention -* bBDense: use bias for Dense multiplication in self-attention - -## Tokenizer and special tokens comparison - -| Name | Tokenizer | bos | eos | pad | cls | startofpiece | endofpiece | mask | smask | gmask | -| :--------------: | :--------------: | :----: | :----: | :---: | :---: | :----------: | :--------: | :----: | :---: | :----: | -| chatglm2_6b | ChatGLMTokenizer | 1 | 2 | 0 | | | | | | | -| chatglm2_6b_32k | ChatGLMTokenizer | 1 | 2 | 0 | | | | | | | - -## Usage - -The next section describe how to build the engine and run the inference demo. - -### 1. Download repo and weights from HuggingFace Transformers - -```bash -pip install -r requirements.txt -apt-get update -apt-get install git-lfs -rm -rf chatglm* - -# clone one or more models we want to build -git clone https://huggingface.co/THUDM/chatglm2-6b chatglm2_6b -git clone https://huggingface.co/THUDM/chatglm2-6b-32k chatglm2_6b_32k -``` - -For more example codes, please refer to the [examples/models/core/glm-4-9b/README.md](../../../glm-4-9b/README.md). diff --git a/examples/models/contrib/chatglm2-6b/requirements.txt b/examples/models/contrib/chatglm2-6b/requirements.txt deleted file mode 100644 index cdc65bf2bb38..000000000000 --- a/examples/models/contrib/chatglm2-6b/requirements.txt +++ /dev/null @@ -1,8 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -protobuf -rouge_score -sentencepiece -tiktoken diff --git a/examples/models/contrib/chatglm2-6b/tokenization_chatglm.py b/examples/models/contrib/chatglm2-6b/tokenization_chatglm.py deleted file mode 100644 index dea5fdc7c7dc..000000000000 --- a/examples/models/contrib/chatglm2-6b/tokenization_chatglm.py +++ /dev/null @@ -1,282 +0,0 @@ -import os -from typing import Dict, List, Optional, Union - -from sentencepiece import SentencePieceProcessor -from transformers import PreTrainedTokenizer -from transformers.tokenization_utils_base import BatchEncoding, EncodedInput -from transformers.utils import PaddingStrategy - - -class SPTokenizer: - - def __init__(self, model_path: str): - # reload tokenizer - assert os.path.isfile(model_path), model_path - self.sp_model = SentencePieceProcessor(model_file=model_path) - - # BOS / EOS token IDs - self.n_words: int = self.sp_model.vocab_size() - self.bos_id: int = self.sp_model.bos_id() - self.eos_id: int = self.sp_model.eos_id() - self.pad_id: int = self.sp_model.unk_id() - assert self.sp_model.vocab_size() == self.sp_model.get_piece_size() - - special_tokens = ["[MASK]", "[gMASK]", "[sMASK]", "sop", "eop"] - self.special_tokens = {} - self.index_special_tokens = {} - for token in special_tokens: - self.special_tokens[token] = self.n_words - self.index_special_tokens[self.n_words] = token - self.n_words += 1 - - def tokenize(self, s: str): - return self.sp_model.EncodeAsPieces(s) - - def encode(self, s: str, bos: bool = False, eos: bool = False) -> List[int]: - assert type(s) is str - t = self.sp_model.encode(s) - if bos: - t = [self.bos_id] + t - if eos: - t = t + [self.eos_id] - return t - - def decode(self, t: List[int]) -> str: - return self.sp_model.decode(t) - - def decode_tokens(self, tokens: List[str]) -> str: - text = self.sp_model.DecodePieces(tokens) - return text - - def convert_token_to_id(self, token): - """ Converts a token (str) in an id using the vocab. """ - if token in self.special_tokens: - return self.special_tokens[token] - return self.sp_model.PieceToId(token) - - def convert_id_to_token(self, index): - """Converts an index (integer) in a token (str) using the vocab.""" - if index in self.index_special_tokens or index in [ - self.eos_id, self.bos_id, self.pad_id - ] or index < 0: - return "" - return self.sp_model.IdToPiece(index) - - -class ChatGLMTokenizer(PreTrainedTokenizer): - vocab_files_names = {"vocab_file": "tokenizer.model"} - - model_input_names = ["input_ids", "attention_mask", "position_ids"] - - def __init__(self, - vocab_file, - padding_side="left", - clean_up_tokenization_spaces=False, - **kwargs): - self.name = "GLMTokenizer" - - self.vocab_file = vocab_file - self.tokenizer = SPTokenizer(vocab_file) - self.special_tokens = { - "": self.tokenizer.bos_id, - "": self.tokenizer.eos_id, - "": self.tokenizer.pad_id - } - super().__init__( - padding_side=padding_side, - clean_up_tokenization_spaces=clean_up_tokenization_spaces, - **kwargs) - - def get_command(self, token): - if token in self.special_tokens: - return self.special_tokens[token] - assert token in self.tokenizer.special_tokens, f"{token} is not a special token for {self.name}" - return self.tokenizer.special_tokens[token] - - @property - def unk_token(self) -> str: - return "" - - @property - def pad_token(self) -> str: - return "" - - @property - def pad_token_id(self): - return self.get_command("") - - @property - def eos_token(self) -> str: - return "" - - @property - def eos_token_id(self): - return self.get_command("") - - @property - def vocab_size(self): - return self.tokenizer.n_words - - def get_vocab(self): - """ Returns vocab as a dict """ - vocab = { - self._convert_id_to_token(i): i - for i in range(self.vocab_size) - } - vocab.update(self.added_tokens_encoder) - return vocab - - def _tokenize(self, text, **kwargs): - return self.tokenizer.tokenize(text) - - def _convert_token_to_id(self, token): - """ Converts a token (str) in an id using the vocab. """ - return self.tokenizer.convert_token_to_id(token) - - def _convert_id_to_token(self, index): - """Converts an index (integer) in a token (str) using the vocab.""" - return self.tokenizer.convert_id_to_token(index) - - def convert_tokens_to_string(self, tokens: List[str]) -> str: - return self.tokenizer.decode_tokens(tokens) - - def save_vocabulary(self, save_directory, filename_prefix=None): - """ - Save the vocabulary and special tokens file to a directory. - - Args: - save_directory (`str`): - The directory in which to save the vocabulary. - filename_prefix (`str`, *optional*): - An optional prefix to add to the named of the saved files. - - Returns: - `Tuple(str)`: Paths to the files saved. - """ - if os.path.isdir(save_directory): - vocab_file = os.path.join(save_directory, - self.vocab_files_names["vocab_file"]) - else: - vocab_file = save_directory - - with open(self.vocab_file, 'rb') as fin: - proto_str = fin.read() - - with open(vocab_file, "wb") as writer: - writer.write(proto_str) - - return (vocab_file, ) - - def get_prefix_tokens(self): - prefix_tokens = [self.get_command("[gMASK]"), self.get_command("sop")] - return prefix_tokens - - def build_prompt(self, query, history=None): - if history is None: - history = [] - prompt = "" - for i, (old_query, response) in enumerate(history): - prompt += "[Round {}]\n\n问：{}\n\n答：{}\n\n".format( - i + 1, old_query, response) - prompt += "[Round {}]\n\n问：{}\n\n答：".format(len(history) + 1, query) - return prompt - - def build_inputs_with_special_tokens( - self, - token_ids_0: List[int], - token_ids_1: Optional[List[int]] = None) -> List[int]: - """ - Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and - adding special tokens. A BERT sequence has the following format: - - - single sequence: `[CLS] X [SEP]` - - pair of sequences: `[CLS] A [SEP] B [SEP]` - - Args: - token_ids_0 (`List[int]`): - List of IDs to which the special tokens will be added. - token_ids_1 (`List[int]`, *optional*): - Optional second list of IDs for sequence pairs. - - Returns: - `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens. - """ - prefix_tokens = self.get_prefix_tokens() - token_ids_0 = prefix_tokens + token_ids_0 - if token_ids_1 is not None: - token_ids_0 = token_ids_0 + token_ids_1 + [ - self.get_command("") - ] - return token_ids_0 - - def _pad( - self, - encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding], - max_length: Optional[int] = None, - padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, - pad_to_multiple_of: Optional[int] = None, - return_attention_mask: Optional[bool] = None, - padding_side: str = "left", # Fix for new transformers - ) -> dict: - """ - Pad encoded inputs (on left/right and up to predefined length or max length in the batch) - - Args: - encoded_inputs: - Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`). - max_length: maximum length of the returned list and optionally padding length (see below). - Will truncate by taking into account the special tokens. - padding_strategy: PaddingStrategy to use for padding. - - - PaddingStrategy.LONGEST Pad to the longest sequence in the batch - - PaddingStrategy.MAX_LENGTH: Pad to the max length (default) - - PaddingStrategy.DO_NOT_PAD: Do not pad - The tokenizer padding sides are defined in self.padding_side: - - - 'left': pads on the left of the sequences - - 'right': pads on the right of the sequences - pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value. - This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability - `>= 7.5` (Volta). - return_attention_mask: - (optional) Set to False to avoid returning attention mask (default: set to model specifics) - """ - # Load from model defaults - assert self.padding_side == "left" - - required_input = encoded_inputs[self.model_input_names[0]] - seq_length = len(required_input) - - if padding_strategy == PaddingStrategy.LONGEST: - max_length = len(required_input) - - if max_length is not None and pad_to_multiple_of is not None and ( - max_length % pad_to_multiple_of != 0): - max_length = ( - (max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of - - needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len( - required_input) != max_length - - # Initialize attention mask if not present. - if "attention_mask" not in encoded_inputs: - encoded_inputs["attention_mask"] = [1] * seq_length - - if "position_ids" not in encoded_inputs: - encoded_inputs["position_ids"] = list(range(seq_length)) - - if needs_to_be_padded: - difference = max_length - len(required_input) - - if "attention_mask" in encoded_inputs: - encoded_inputs["attention_mask"] = [ - 0 - ] * difference + encoded_inputs["attention_mask"] - if "position_ids" in encoded_inputs: - encoded_inputs["position_ids"] = [ - 0 - ] * difference + encoded_inputs["position_ids"] - encoded_inputs[self.model_input_names[ - 0]] = [self.pad_token_id] * difference + required_input - - return encoded_inputs diff --git a/examples/models/contrib/chatglm3-6b-32k/README.md b/examples/models/contrib/chatglm3-6b-32k/README.md deleted file mode 100644 index 0636cba34f9c..000000000000 --- a/examples/models/contrib/chatglm3-6b-32k/README.md +++ /dev/null @@ -1,112 +0,0 @@ -# ChatGLM - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document explains how to build the [ChatGLM3-6B](https://huggingface.co/THUDM/chatglm3-6b), [ChatGLM3-6B-Base](https://huggingface.co/THUDM/chatglm3-6b-base), [ChatGLM3-6B-32k](https://huggingface.co/THUDM/chatglm3-6b-32k) models using TensorRT LLM and run on a single GPU, a single node with multiple GPUs or multiple nodes with multiple GPUs. - -- [ChatGLM](#chatglm) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Model comparison](#model-comparison) - - [Tokenizer and special tokens comparison](#tokenizer-and-special-tokens-comparison) - - [Usage](#usage) - - [1. Download repo and weights from HuggingFace Transformers](#1-download-repo-and-weights-from-huggingface-transformers) - - [2. Convert weights from HF Transformers to TensorRT LLM format](#2-convert-weights-from-hf-transformers-to-tensorrt-llm-format) - - [3. Build TensorRT engine(s)](#3-build-tensorrt-engines) - - [Enable plugins](#enable-plugins) - - [In-flight batching](#in-flight-batching) - - [4. Run inference](#4-run-inference) - - [Single node, single GPU](#single-node-single-gpu) - - [Single node, multi GPU](#single-node-multi-gpu) - - [5. Run summarization task](#5-run-summarization-task) - - [Weight Only quantization](#weight-only-quantization) - - [Smooth Quantization (SQ)](#smooth-quantization-sq) - - [Activation-aware Weight Quantization (AWQ)](#activation-aware-weight-quantization-awq) - - [FP8 Quantization](#fp8-quantization) - - [Benchmark](#benchmark) - - -## Overview - -The TensorRT LLM ChatGLM implementation can be found in [`tensorrt_llm/models/chatglm/model.py`](../../tensorrt_llm/models/chatglm/model.py). -The TensorRT LLM ChatGLM example code is located in [`examples/models/contrib/chatglm3-6b-32k`](./). There is one main file: - -* [`examples/models/core/glm-4-9b/convert_checkpoint.py`](../../../glm-4-9b/convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - -| Model Name | FP16 | FMHA | WO | SQ | AWQ | FP8 | TP | PP | ST | C++ | benchmark | IFB | -| :--------------: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :-------: | :---: | -| chatglm3_6b | Y | Y | Y | Y | Y | Y | Y | | Y | Y | Y | Y | -| chatglm3_6b_base | Y | Y | Y | Y | Y | Y | Y | | Y | Y | Y | Y | -| chatglm3_6b_32k | Y | Y | Y | Y | Y | Y | Y | | Y | Y | Y | Y | - -* Model Name: the name of the model, the same as the name on HuggingFace -* FMHA: Fused MultiHead Attention (see introduction below) -* WO: Weight Only Quantization (int8 / int4) -* SQ: Smooth Quantization (int8) -* AWQ: Activation Aware Weight Quantization (int4) -* FP8: FP8 Quantization -* TP: Tensor Parallel -* PP: Pipeline Parallel -* ST: Strongly Typed -* C++: C++ Runtime -* benchmark: benchmark by python / C++ Runtime -* IFB: In-flight Batching (see introduction below) - -## Model comparison - -| Name | nL | nAH | nKH | nHW | nH | nF | nMSL | nV | bP2D | bBQKV | bBDense | Comments | -| :--------------: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :----: | :---: | :---: | :-----: | :----------------------------------------------------------------- | -| chatglm3_6b | 28 | 32 | 2 | 128 | 4096 | 13696 | 8192 | 65024 | N | Y | N | Different in preprocess and postprocess than chatglm2_6b | -| chatglm3_6b_base | 28 | 32 | 2 | 128 | 4096 | 13696 | 32768 | 65024 | N | Y | N | | -| chatglm3_6b_32k | 28 | 32 | 2 | 128 | 4096 | 13696 | 32768 | 65024 | N | Y | N | RoPE base=500000 rather than 10000 in chatglm3_6b | - -* nL: number of layers -* nAH: number of attention heads -* nKH: number of kv heads (less than nAH if multi_query_attention is used) -* nHW: head width -* nH: hidden size -* nF: FFN hidden size -* nMSL: max sequence length (input + output) -* nV: vocabulary size -* bP2D: use position_encoding_2d (Y: Yes, N: No) -* bBQKV: use bias for QKV multiplication in self-attention -* bBDense: use bias for Dense multiplication in self-attention - -## Tokenizer and special tokens comparison - -| Name | Tokenizer | bos | eos | pad | cls | startofpiece | endofpiece | mask | smask | gmask | -| :--------------: | :--------------: | :----: | :----: | :---: | :---: | :----------: | :--------: | :----: | :---: | :----: | -| chatglm3_6b | ChatGLMTokenizer | 1 | 2 | 0 | | | | 130000 | | | -| chatglm3_6b_base | ChatGLMTokenizer | 1 | 2 | 0 | | | | 130000 | | | -| chatglm3_6b_32k | ChatGLMTokenizer | 1 | 2 | 0 | | | | 130000 | | | - -## Usage - -The next section describe how to build the engine and run the inference demo. - -### 1. Download repo and weights from HuggingFace Transformers - -```bash -pip install -r requirements.txt -apt-get update -apt-get install git-lfs -rm -rf chatglm* - -# clone one or more models we want to build -git clone https://huggingface.co/THUDM/chatglm3-6b chatglm3_6b -git clone https://huggingface.co/THUDM/chatglm3-6b-base chatglm3_6b_base -git clone https://huggingface.co/THUDM/chatglm3-6b-32k chatglm3_6b_32k -``` - -For more example codes, please refer to the [examples/models/core/glm-4-9b/README.md](../../../glm-4-9b/README.md). diff --git a/examples/models/contrib/chatglm3-6b-32k/requirements.txt b/examples/models/contrib/chatglm3-6b-32k/requirements.txt deleted file mode 100644 index cdc65bf2bb38..000000000000 --- a/examples/models/contrib/chatglm3-6b-32k/requirements.txt +++ /dev/null @@ -1,8 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -protobuf -rouge_score -sentencepiece -tiktoken diff --git a/examples/models/contrib/chatglm3-6b-32k/tokenization_chatglm.py b/examples/models/contrib/chatglm3-6b-32k/tokenization_chatglm.py deleted file mode 100644 index 67a3d52442a6..000000000000 --- a/examples/models/contrib/chatglm3-6b-32k/tokenization_chatglm.py +++ /dev/null @@ -1,313 +0,0 @@ -import json -import os -from typing import Dict, List, Optional, Union - -from sentencepiece import SentencePieceProcessor -from transformers import PreTrainedTokenizer -from transformers.tokenization_utils_base import BatchEncoding, EncodedInput -from transformers.utils import PaddingStrategy - - -class SPTokenizer: - - def __init__(self, model_path: str): - # reload tokenizer - assert os.path.isfile(model_path), model_path - self.sp_model = SentencePieceProcessor(model_file=model_path) - - # BOS / EOS token IDs - self.n_words: int = self.sp_model.vocab_size() - self.bos_id: int = self.sp_model.bos_id() - self.eos_id: int = self.sp_model.eos_id() - self.pad_id: int = self.sp_model.unk_id() - assert self.sp_model.vocab_size() == self.sp_model.get_piece_size() - - special_tokens = [ - "[MASK]", "[gMASK]", "[sMASK]", "sop", "eop", "<|system|>", - "<|user|>", "<|assistant|>", "<|observation|>" - ] - self.special_tokens = {} - self.index_special_tokens = {} - for token in special_tokens: - self.special_tokens[token] = self.n_words - self.index_special_tokens[self.n_words] = token - self.n_words += 1 - - def tokenize(self, s: str): - return self.sp_model.EncodeAsPieces(s) - - def encode(self, s: str, bos: bool = False, eos: bool = False) -> List[int]: - assert type(s) is str - t = self.sp_model.encode(s) - if bos: - t = [self.bos_id] + t - if eos: - t = t + [self.eos_id] - return t - - def decode(self, t: List[int]) -> str: - text, buffer = "", [] - for token in t: - if token in self.index_special_tokens: - if buffer: - text += self.sp_model.decode(buffer) - buffer = [] - text += self.index_special_tokens[token] - else: - buffer.append(token) - if buffer: - text += self.sp_model.decode(buffer) - return text - - def decode_tokens(self, tokens: List[str]) -> str: - text = self.sp_model.DecodePieces(tokens) - return text - - def convert_token_to_id(self, token): - """ Converts a token (str) in an id using the vocab. """ - if token in self.special_tokens: - return self.special_tokens[token] - return self.sp_model.PieceToId(token) - - def convert_id_to_token(self, index): - """Converts an index (integer) in a token (str) using the vocab.""" - if index in self.index_special_tokens: - return self.index_special_tokens[index] - if index in [self.eos_id, self.bos_id, self.pad_id] or index < 0: - return "" - return self.sp_model.IdToPiece(index) - - -class ChatGLMTokenizer(PreTrainedTokenizer): - vocab_files_names = {"vocab_file": "tokenizer.model"} - - model_input_names = ["input_ids", "attention_mask", "position_ids"] - - def __init__(self, - vocab_file, - padding_side="left", - clean_up_tokenization_spaces=False, - **kwargs): - self.name = "GLMTokenizer" - - self.vocab_file = vocab_file - self.tokenizer = SPTokenizer(vocab_file) - self.special_tokens = { - "": self.tokenizer.bos_id, - "": self.tokenizer.eos_id, - "": self.tokenizer.pad_id - } - super().__init__( - padding_side=padding_side, - clean_up_tokenization_spaces=clean_up_tokenization_spaces, - **kwargs) - - def get_command(self, token): - if token in self.special_tokens: - return self.special_tokens[token] - assert token in self.tokenizer.special_tokens, f"{token} is not a special token for {self.name}" - return self.tokenizer.special_tokens[token] - - @property - def unk_token(self) -> str: - return "" - - @property - def pad_token(self) -> str: - return "" - - @property - def pad_token_id(self): - return self.get_command("") - - @property - def eos_token(self) -> str: - return "" - - @property - def eos_token_id(self): - return self.get_command("") - - @property - def vocab_size(self): - return self.tokenizer.n_words - - def get_vocab(self): - """ Returns vocab as a dict """ - vocab = { - self._convert_id_to_token(i): i - for i in range(self.vocab_size) - } - vocab.update(self.added_tokens_encoder) - return vocab - - def _tokenize(self, text, **kwargs): - return self.tokenizer.tokenize(text) - - def _convert_token_to_id(self, token): - """ Converts a token (str) in an id using the vocab. """ - return self.tokenizer.convert_token_to_id(token) - - def _convert_id_to_token(self, index): - """Converts an index (integer) in a token (str) using the vocab.""" - return self.tokenizer.convert_id_to_token(index) - - def convert_tokens_to_string(self, tokens: List[str]) -> str: - return self.tokenizer.decode_tokens(tokens) - - def save_vocabulary(self, save_directory, filename_prefix=None): - """ - Save the vocabulary and special tokens file to a directory. - - Args: - save_directory (`str`): - The directory in which to save the vocabulary. - filename_prefix (`str`, *optional*): - An optional prefix to add to the named of the saved files. - - Returns: - `Tuple(str)`: Paths to the files saved. - """ - if os.path.isdir(save_directory): - vocab_file = os.path.join(save_directory, - self.vocab_files_names["vocab_file"]) - else: - vocab_file = save_directory - - with open(self.vocab_file, 'rb') as fin: - proto_bytes = fin.read() - - with open(vocab_file, "wb") as writer: - writer.write(proto_bytes) - - return (vocab_file, ) - - def get_prefix_tokens(self): - prefix_tokens = [self.get_command("[gMASK]"), self.get_command("sop")] - return prefix_tokens - - def build_single_message(self, role, metadata, message): - assert role in ["system", "user", "assistant", "observation"], role - role_tokens = [self.get_command(f"<|{role}|>") - ] + self.tokenizer.encode(f"{metadata}\n") - message_tokens = self.tokenizer.encode(message) - tokens = role_tokens + message_tokens - return tokens - - def build_chat_input(self, query, history=None, role="user"): - if history is None: - history = [] - input_ids = [] - for item in history: - content = item["content"] - if item["role"] == "system" and "tools" in item: - content = content + "\n" + json.dumps( - item["tools"], indent=4, ensure_ascii=False) - input_ids.extend( - self.build_single_message(item["role"], - item.get("metadata", ""), content)) - input_ids.extend(self.build_single_message(role, "", query)) - input_ids.extend([self.get_command("<|assistant|>")]) - return self.batch_encode_plus([input_ids], - return_tensors="pt", - is_split_into_words=True) - - def build_inputs_with_special_tokens( - self, - token_ids_0: List[int], - token_ids_1: Optional[List[int]] = None) -> List[int]: - """ - Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and - adding special tokens. A BERT sequence has the following format: - - - single sequence: `[CLS] X [SEP]` - - pair of sequences: `[CLS] A [SEP] B [SEP]` - - Args: - token_ids_0 (`List[int]`): - List of IDs to which the special tokens will be added. - token_ids_1 (`List[int]`, *optional*): - Optional second list of IDs for sequence pairs. - - Returns: - `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens. - """ - prefix_tokens = self.get_prefix_tokens() - token_ids_0 = prefix_tokens + token_ids_0 - if token_ids_1 is not None: - token_ids_0 = token_ids_0 + token_ids_1 + [ - self.get_command("") - ] - return token_ids_0 - - def _pad( - self, - encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding], - max_length: Optional[int] = None, - padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, - pad_to_multiple_of: Optional[int] = None, - return_attention_mask: Optional[bool] = None, - padding_side: str = "left", # Fix for new transformers - ) -> dict: - """ - Pad encoded inputs (on left/right and up to predefined length or max length in the batch) - - Args: - encoded_inputs: - Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`). - max_length: maximum length of the returned list and optionally padding length (see below). - Will truncate by taking into account the special tokens. - padding_strategy: PaddingStrategy to use for padding. - - - PaddingStrategy.LONGEST Pad to the longest sequence in the batch - - PaddingStrategy.MAX_LENGTH: Pad to the max length (default) - - PaddingStrategy.DO_NOT_PAD: Do not pad - The tokenizer padding sides are defined in self.padding_side: - - - 'left': pads on the left of the sequences - - 'right': pads on the right of the sequences - pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value. - This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability - `>= 7.5` (Volta). - return_attention_mask: - (optional) Set to False to avoid returning attention mask (default: set to model specifics) - """ - # Load from model defaults - assert self.padding_side == "left" - - required_input = encoded_inputs[self.model_input_names[0]] - seq_length = len(required_input) - - if padding_strategy == PaddingStrategy.LONGEST: - max_length = len(required_input) - - if max_length is not None and pad_to_multiple_of is not None and ( - max_length % pad_to_multiple_of != 0): - max_length = ( - (max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of - - needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len( - required_input) != max_length - - # Initialize attention mask if not present. - if "attention_mask" not in encoded_inputs: - encoded_inputs["attention_mask"] = [1] * seq_length - - if "position_ids" not in encoded_inputs: - encoded_inputs["position_ids"] = list(range(seq_length)) - - if needs_to_be_padded: - difference = max_length - len(required_input) - - if "attention_mask" in encoded_inputs: - encoded_inputs["attention_mask"] = [ - 0 - ] * difference + encoded_inputs["attention_mask"] - if "position_ids" in encoded_inputs: - encoded_inputs["position_ids"] = [ - 0 - ] * difference + encoded_inputs["position_ids"] - encoded_inputs[self.model_input_names[ - 0]] = [self.pad_token_id] * difference + required_input - - return encoded_inputs diff --git a/examples/models/contrib/internlm/.gitignore b/examples/models/contrib/internlm/.gitignore deleted file mode 100644 index 7ce339719a3e..000000000000 --- a/examples/models/contrib/internlm/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -internlm* -tokenizer.model diff --git a/examples/models/contrib/internlm/README.md b/examples/models/contrib/internlm/README.md deleted file mode 100644 index f44180f01aad..000000000000 --- a/examples/models/contrib/internlm/README.md +++ /dev/null @@ -1,320 +0,0 @@ -# InternLM - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run InternLM 7B / 20B models in TensorRT LLM on both single GPU, single node multi-GPU and multi-node multi-GPU. - -- [InternLM](#internlm) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [Build TensorRT engine(s)](#build-tensorrt-engines) - - [INT8 weight only + INT8 KV cache](#int8-weight-only--int8-kv-cache) - - [SmoothQuant](#smoothquant) - - [Run](#run) - - [Summarization using the InternLM model](#summarization-using-the-internlm-model) - -## Overview - -The TensorRT LLM InternLM implementation is based on the LLaMA model. The implementation can -be found in [tensorrt_llm/models/llama/model.py](../../../../tensorrt_llm/models/llama/model.py). -The TensorRT LLM InternLM example code lies in [`examples/models/contrib/internlm`](./): - -* [`convert_checkpoint.py`](../../core/llama/convert_checkpoint.py) converts the Huggingface Model of InternLM into TensorRT LLM checkpoint. -* [`convert_checkpoint.py`] to to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 / BF16 - * INT8 & INT4 Weight-Only - * Smooth Quant - * INT8 KV Cache - * Tensor Parallel & Pipeline Parallel - -## Usage - -The TensorRT LLM InternLM example code locates at [examples/models/contrib/internlm](./). It takes HF weights as input, and builds the corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Build TensorRT engine(s) - -Please install required packages first: - -```bash -pip install -r requirements.txt -``` - -TensorRT LLM InternLM builds TensorRT engine(s) from HF checkpoint. If no checkpoint directory is specified, TensorRT LLM will build engine(s) with dummy weights. - -InternLM has released several checkpoints of different size or capabilities under https://huggingface.co/internlm. Users can pick any one repository and follow instructions to prepare the checkpoint. - -Below examples use [internlm-chat-7b](https://huggingface.co/internlm/internlm-chat-7b) and [internlm-chat-20b](https://huggingface.co/internlm/internlm-chat-20b) and assume these repositories are cloned or linked under this directory, for example `./internlm-chat-7b/`. - -Normally `trtllm-build` only requires single GPU, but if you've already got all the GPUs needed for inference, you could enable parallel building to make the engine building process faster by adding `--workers` argument. Please note that currently `--workers` feature only supports single node. - -Here're some examples: - -```bash -# Build a single-GPU float16 engine from HF weights. -# gpt_attention_plugin is necessary in InternLM. -# Try use_gemm_plugin to prevent accuracy issue. -cd examples/models/core/llama - -# Convert the InternLM 7B model using a single GPU and FP16. -python convert_checkpoint.py --model_dir ./internlm-chat-7b/ \ - --dtype float16 \ - --output_dir ./internlm-chat-7b/trt_engines/fp16/1-gpu/ -# Note: setting `--dtype bfloat16` to use bfloat16 precision. - -# BUild the InternLM 7B model using a single GPU -trtllm-build --checkpoint_dir ./internlm-chat-7b/trt_engines/fp16/1-gpu/ \ - --output_dir ./engine_outputs \ - --gemm_plugin float16 - -# Convert the InternLM 7B model using a single GPU and apply INT8 weight-only quantization.. -python convert_checkpoint.py --model_dir ./internlm-chat-7b/ \ - --dtype float16 \ - --output_dir ./internlm-chat-7b/trt_engines/int8/1-gpu/ \ - --use_weight_only \ - --weight_only_precision int8 - -trtllm-build --checkpoint_dir ./internlm-chat-7b/trt_engines/int8/1-gpu/ \ - --output_dir ./engine_outputs \ - --gemm_plugin float16 - -# Note: setting `--weight_only_precision int4` to use INT4 weight-only quantization - -# Build InternLM 7B using 2-way tensor parallelism. -python convert_checkpoint.py --model_dir ./internlm-chat-7b/ \ - --dtype float16 \ - --output_dir ./internlm-chat-7b/trt_engines/fp16/2-gpu/ \ - --tp_size 2 - -trtllm-build --checkpoint_dir ./internlm-chat-7b/trt_engines/fp16/2-gpu/ \ - --output_dir ./engine_outputs \ - --gemm_plugin float16 - -# Build InternLM 20B using 2-way tensor parallelism. -python convert_checkpoint.py --model_dir ./internlm-chat-20b/ \ - --dtype bfloat16 \ - --output_dir ./internlm-chat-20b/trt_engines/bf16/2-gpu/ \ - --tp_size 2 --workers 2 - -trtllm-build --checkpoint_dir ./internlm-chat-7b/trt_engines/bf16/2-gpu/ \ - --output_dir ./engine_outputs \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 -``` - -#### INT8 weight only + INT8 KV cache - -For INT8 KV cache, [`convert_checkpoint.py`](./convert_checkpoint.py) features a -`--int8_kv_cache` option. Setting `--int8_kv_cache` will calibrate the model, -and then export the scaling factors needed for INT8 KV cache inference. - - -Example: - -```bash -cd examples/models/core/llama - -# For 7B models -python convert_checkpoint.py --model_dir ./internlm-chat-7b \ - --output_dir ./internlm-chat-7b/smooth_internlm/int8_kv_cache/ \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int8 \ - --int8_kv_cache - -# Build 7B model with both INT8 weight-only and INT8 KV cache enabled -trtllm-build --checkpoint_dir ./internlm-chat-7b/smooth_internlm/int8_kv_cache/ \ - --output_dir ./engine_outputs \ - --gemm_plugin float16 \ -``` - - -```bash -cd examples/models/core/llama - -# For 20B models -python convert_checkpoint.py --model_dir ./internlm-chat-20b \ - --output_dir ./internlm-chat-20b/smooth_internlm/int8_kv_cache/ \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int8 \ - --int8_kv_cache - -# Build 20B model with both INT8 weight-only and INT8 KV cache enabled -trtllm-build --checkpoint_dir ./internlm-chat-20b/smooth_internlm/int8_kv_cache/ \ - --output_dir ./engine_outputs \ - --gemm_plugin float16 \ -``` - - -Test with `../../../run.py` or `../../../summarize.py`: - -```bash -python ../../../run.py --max_output_len=120 \ - --input_text 'Tell me about yourself.' \ - --tokenizer_dir ./internlm-chat-7b/ \ - --engine_dir ./internlm-chat-7b/trt_engines/int8_kv_cache_weight_only/1-gpu - -python ../../../run.py --max_output_len=120 \ - --input_text 'Tell me about yourself.' \ - --tokenizer_dir ./internlm-chat-20b/ \ - --engine_dir ./internlm-chat-20b/trt_engines/int8_kv_cache_weight_only/1-gpu - -python ../../../summarize.py --test_trt_llm --test_hf \ - --hf_model_dir ./internlm-chat-7b \ - --data_type fp16 \ - --engine_dir ./internlm-chat-7b/trt_engines/int8_kv_cache_weight_only/1-gpu - -python ../../../summarize.py --test_trt_llm --test_hf \ - --hf_model_dir ./internlm-chat-20b \ - --data_type fp16 \ - --engine_dir ./internlm-chat-20b/trt_engines/int8_kv_cache_weight_only/1-gpu -``` - -#### SmoothQuant - -Unlike the FP16 build where the HF weights are processed and loaded into the TensorRT LLM directly, the SmoothQuant needs to load INT8 weights which should be pre-processed before building an engine. - -Example: -```bash -cd examples/models/core/llama - -# For 7B models -python convert_checkpoint.py --model_dir ./internlm-chat-7b --output_dir ./internlm-chat-7b/smooth_internlm/sq0.5/ --dtype float16 --smoothquant 0.5 -# Build the engine -trtllm-build --checkpoint_dir ./internlm-chat-7b/smooth_internlm/sq0.5/ \ - --output_dir ./engine_outputs \ - --gemm_plugin float16 - -# For 20B models -cd examples/models/core/llama - -python convert_checkpoint.py --model_dir ./internlm-chat-20b --output_dir ./internlm-chat-20b/smooth_internlm/sq0.5/ --dtype float16 --smoothquant 0.5 -trtllm-build --checkpoint_dir ./internlm-chat-20b/smooth_internlm/sq0.5/ \ - --output_dir ./engine_outputs \ - --gemm_plugin float16 -``` - -[`convert_checkpoint.py`](./convert_checkpoint.py) add new options for the support of INT8 inference of SmoothQuant models. - -`--smoothquant` is the starting point of INT8 inference. By default, it -will run the model in the _per-tensor_ mode. - -Then, you can add any combination of `--per-token` and `--per-channel` to get the corresponding behaviors. - -Examples of build invocations: - -```bash -# Build model for SmoothQuant in the _per_token_ + _per_channel_ mode -cd examples/models/core/llama - -# 7B model -python convert_checkpoint.py --model_dir ./internlm-chat-7b --output_dir ./internlm-chat-7b/smooth_internlm/sq0.5/ --dtype float16 --smoothquant 0.5 --per_channel --per_token - -# 20B model -python convert_checkpoint.py --model_dir ./internlm-chat-20b --output_dir ./internlm-chat-20b/smooth_internlm/sq0.5/ --dtype float16 --smoothquant 0.5 --per_channel --per_token -``` - - -Test with `../../../run.py` or `../../../summarize.py`: - -```bash -python ../../../run.py --max_output_len=120 \ - --input_text 'Tell me about yourself.' \ - --tokenizer_dir ./internlm-chat-7b/ \ - --engine_dir ./internlm-chat-7b/smooth_internlm/sq0.5/ - -python ../../../run.py --max_output_len=120 \ - --input_text 'Tell me about yourself.' \ - --tokenizer_dir ./internlm-chat-20b/ \ - --engine_dir ./internlm-chat-20b/smooth_internlm/sq0.5/ - -python ../../../summarize.py --test_trt_llm --test_hf \ - --hf_model_dir ./internlm-chat-7b \ - --data_type fp16 \ - --engine_dir ./internlm-chat-7b/smooth_internlm/sq0.5/ - -python ../../../summarize.py --test_trt_llm --test_hf \ - --hf_model_dir ./internlm-chat-20b \ - --data_type fp16 \ - --engine_dir ./internlm-chat-20b/smooth_internlm/sq0.5/ -``` - -### Run - -To run a TensorRT LLM InternLM model using the engines generated by `trtllm-build` - -```bash -# InternLM 7B with fp16 -python ../../../run.py --max_output_len=120 \ - --input_text 'Tell me about yourself.' \ - --tokenizer_dir ./internlm-chat-7b/ \ - --engine_dir=./internlm-chat-7b/trt_engines/fp16/1-gpu/ - -# InternLM 7B with bf16 -python ../../../run.py --max_output_len=120 \ - --input_text 'Tell me about yourself.' \ - --tokenizer_dir ./internlm-chat-7b/ \ - --engine_dir=./internlm-chat-7b/trt_engines/bf16/1-gpu/ - -# InternLM 7B with int8 weight only quantization -python ../../../run.py --max_output_len=120 \ - --input_text 'Tell me about yourself.' \ - --tokenizer_dir ./internlm-chat-7b/ \ - --engine_dir=./internlm-chat-7b/trt_engines/weight_only/1-gpu/ - -# InternLM 7B with fp16 and tensor parallelism -mpirun -n 2 --allow-run-as-root \ - python ../../../run.py --max_output_len=120 \ - --input_text 'Tell me about yourself.' \ - --tokenizer_dir ./internlm-chat-7b/ \ - --engine_dir=./internlm-chat-7b/trt_engines/fp16/2-gpu/ - -# InternLM 20B with fp16 and tensor parallelism and pipeline parallelism -mpirun -n 4 --allow-run-as-root \ - python ../../../run.py --max_output_len=120 \ - --input_text 'Tell me about yourself.' \ - --tokenizer_dir ./internlm-chat-7b/ \ - --engine_dir=./internlm-chat-7b/trt_engines/bf16/4-gpu/ -``` - -### Summarization using the InternLM model - -```bash -# Run summarization using the InternLM 7B model in FP16. -python ../../../summarize.py --test_trt_llm --test_hf \ - --hf_model_dir ./internlm-chat-7b/ \ - --data_type fp16 \ - --engine_dir ./engine_outputs - -# Run summarization using the InternLM 7B model quantized to INT8. -python ../../../summarize.py --test_trt_llm --test_hf \ - --hf_model_dir ./internlm-chat-7b/ \ - --data_type fp16 \ - --engine_dir ./engine_outputs - -# Run summarization using the InternLM 7B model in FP16 using two GPUs. -mpirun -n 2 --allow-run-as-root \ - python ../../../summarize.py --test_trt_llm --test_hf \ - --hf_model_dir ./internlm-chat-7b/ \ - --data_type fp16 \ - --engine_dir ./internlm-chat-7b/trt_engines/fp16/2-gpu/ - -# Run summarization using the InternLM 20B model in BF16 using 4 GPUs. -mpirun -n 4 --allow-run-as-root \ - python ../../../summarize.py --test_trt_llm --test_hf \ - --hf_model_dir ./internlm-chat-20b/ \ - --data_type bf16 \ - --engine_dir ./internlm-chat-20b/trt_engines/bf16/4-gpu/ -``` diff --git a/examples/models/contrib/internlm/requirements.txt b/examples/models/contrib/internlm/requirements.txt deleted file mode 100644 index d9354a133c65..000000000000 --- a/examples/models/contrib/internlm/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -rouge_score -sentencepiece>=0.1.99 -evaluate diff --git a/examples/models/contrib/jais/README.md b/examples/models/contrib/jais/README.md deleted file mode 100644 index d15cfd116962..000000000000 --- a/examples/models/contrib/jais/README.md +++ /dev/null @@ -1,125 +0,0 @@ -# Jais - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document elaborates how to build Jais model to runnable engines on multi-GPU node and perform a summarization task using these engines. - -Currently it has been tested on -- [Jais-13b-chat](https://huggingface.co/core42/jais-13b-chat) -- [Jais-30b-chat-v3](https://huggingface.co/core42/jais-30b-chat-v3) - - -- [Jais](#jais) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [Build TensorRT engine(s)](#build-tensorrt-engines) - - [Run inference](#run) - -## Overview - -The TensorRT LLM support for Jais is based on the GPT model, the implementation can be found in [tensorrt_llm/models/gpt/model.py](../../../../tensorrt_llm/models/gpt/model.py). Jais model resembles GPT very much except it uses alibi embedding, embedding scale, swiglu, and logits scale, we therefore reuse the [GPT example code](../../../gpt) for Jais, - -* [`convert_checkpoint.py`](../../../gpt/convert_checkpoint.py) to convert the Jais model into TensorRT LLM checkpoint format. - -In addition, there are two shared files in the parent folder [`examples`](../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix -The tested configurations are: - * FP16 - * FP8 - * Inflight Batching - * Tensor Parallel - -## Usage - -This section gives a whole process where we convert HF models, build TensorRT LLM engines and ultimately perform summarization. - -### Build TensorRT engine(s) - -Run the following commands and TRT-LLM will first transforms a HF model into its own checkpoint format, then builds a TRT engine based on the checkpoint - -```bash -# single gpu, dtype float16 for jais-13b-chat -python3 ../../../gpt/convert_checkpoint.py --model_dir core42/jais-13b-chat \ - --dtype float16 \ - --output_dir jais-13b-chat/trt_ckpt/fp16/1-gpu - -# 2-way tensor parallelism for jais-30b-chat-v3 -python3 ../../../gpt/convert_checkpoint.py --model_dir core42/jais-30b-chat-v3 \ - --dtype float16 \ - --tp_size 2 \ - --output_dir jais-30b-chat-v3/trt_ckpt/fp16/2-gpu -``` - -```bash -# Build a single-GPU float16 engine from TensorRT LLM checkpoint for jais-13b-chat -# Enable the special TensorRT LLM GPT Attention plugin (--gpt_attention_plugin) to increase runtime performance. -# It is recommend to use --remove_input_padding along with --gpt_attention_plugin for better performance -trtllm-build --checkpoint_dir jais-13b-chat/trt_ckpt/fp16/1-gpu \ - --gpt_attention_plugin float16 \ - --remove_input_padding enable \ - --output_dir jais-13b-chat/trt_engines/fp16/1-gpu - -# Build 2-way tensor parallelism engines from TensorRT LLM checkpoint for jais-30b-chat-v3 -trtllm-build --checkpoint_dir jais-30b-chat-v3/trt_ckpt/fp16/2-gpu \ - --gpt_attention_plugin float16 \ - --remove_input_padding enable \ - --output_dir jais-30b-chat-v3/trt_engines/fp16/2-gpu -``` - - -### Run - -The [`../../../run.py`](../../../run.py) script can be used to run inference with the built engine(s). - -```bash -python3 ../../../run.py --engine_dir jais-13b-chat/trt_engines/fp16/1-gpu \ - --tokenizer_dir core42/jais-13b-chat \ - --max_output_len 10 -``` - -If the engines are run successfully, you will see output like: -``` -...... -Input [Text 0]: "Born in north-east France, Soyer trained as a" -Output [Text 0 Beam 0]: " chef in Paris before moving to England in 1816" -``` - -```bash -python3 ../../../run.py --engine_dir jais-13b-chat/trt_engines/fp16/1-gpu \ - --tokenizer_dir core42/jais-13b-chat \ - --max_output_len 8 \ - --input_text "ولد في 1304 ميلادياً ابن بطوطه, لقد ذهب" -``` - -If the engines are run successfully, you will see output like: -``` -..... -Input [Text 0]: "ولد في 1304 ميلادياً ابن بطوطه, لقد ذهب" -Output [Text 0 Beam 0]: " في جميع أنحاء العالم المعروف في ذلك الوقت" -``` - - -To run a 2 TP model you can do the following -```bash -mpirun -np 2 \ - python3 ../../../run.py --engine_dir jais-30b-chat-v3/trt_engines/fp16/2-gpu \ - --tokenizer_dir core42/jais-30b-chat-v3 \ - --max_output_len 30 -``` - -If the engines are run successfully, you will see output like: -``` -Input [Text 0]: "Born in north-east France, Soyer trained as a" -Output [Text 0 Beam 0]: " chef, working in a series of high-end establishments. - -Soyer's career took him to work in a number of establishments across Europe," -``` diff --git a/examples/models/contrib/jais/requirements.txt b/examples/models/contrib/jais/requirements.txt deleted file mode 100644 index 592e01e5ba6d..000000000000 --- a/examples/models/contrib/jais/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score -SentencePiece>=0.1.99 diff --git a/examples/models/contrib/sdxl/README.md b/examples/models/contrib/sdxl/README.md deleted file mode 100644 index 05d874f5f27a..000000000000 --- a/examples/models/contrib/sdxl/README.md +++ /dev/null @@ -1,42 +0,0 @@ -# Stable Diffusion XL - -This document showcases how to build and run the [Stable Diffusion XL (SDXL)](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) model on multiple GPUs using TensorRT-LLM. The community-contributed SDXL example in TRT-LLM is intended solely to showcase distributed inference for high-resolution use cases. For an optimized single-GPU setup in Stable Diffusion inference, please refer to the [TensorRT DemoDiffusion example](https://github.com/NVIDIA/TensorRT/tree/main/demo/Diffusion). - -The design of distributed parallel inference comes from the CVPR 2024 paper [DistriFusion](https://github.com/mit-han-lab/distrifuser) from [MIT HAN Lab](https://hanlab.mit.edu/). To simplify the implementation, all communications in this example are handled synchronously. - -## Usage - -### 1. Build TensorRT Engine - -```bash -# 1 gpu -python build_sdxl_unet.py --size 1024 - -# 2 gpus -mpirun -n 2 --allow-run-as-root python build_sdxl_unet.py --size 1024 -``` - -### 2. Generate images using the engine - - -```bash -# 1 gpu -python run_sdxl.py --size 1024 --prompt "flowers, rabbit" - -# 2 gpus -mpirun -n 2 --allow-run-as-root python run_sdxl.py --size 1024 --prompt "flowers, rabbit" -``` - -## Latency Benchmark -This benchmark is provided as reference points and should not be considered as the peak inference speed that can be delivered by TensorRT-LLM. - -| Framework | Resolution | n_gpu | A100 latency (s) | A100 speedup | H100 latency (s) | H100 speedup | -|:---------:|:----------:|:-----:|:---------------:|:-------------:|:---------------:|:-------------:| -| Torch | 1024x1024 | 1 | 6.280 | 1 | 5.820 | 1 | -| TRT-LLM | 1024x1024 | 2 | 2.803 | **2.24x** | 1.719 | **3.39x** | -| TRT-LLM | 1024x1024 | 4 | 2.962 | **2.12x** | 2.592 | **2.25x** | -| Torch | 2048x2048 | 1 | 27.865 | 1 | 18.330 | 1 | -| TRT-LLM | 2048x2048 | 2 | 13.152 | **2.12x** | 7.943 | **2.31x** | -| TRT-LLM | 2048x2048 | 4 | 9.781 | **2.85x** | 7.596 | **2.41x** | - -torch v2.5.0. TRT-LLM v0.15.0.dev2024102900, `--num-warmup-runs=5; --avg-runs=20`. All communications are synchronous. diff --git a/examples/models/contrib/sdxl/build_sdxl_unet.py b/examples/models/contrib/sdxl/build_sdxl_unet.py deleted file mode 100755 index e2893859f731..000000000000 --- a/examples/models/contrib/sdxl/build_sdxl_unet.py +++ /dev/null @@ -1,148 +0,0 @@ -import argparse -import os - -import tensorrt as trt -import torch -from diffusers import DiffusionPipeline - -import tensorrt_llm -from tensorrt_llm.builder import Builder -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.unet.pp.unet_pp import DistriUNetPP -from tensorrt_llm.models.unet.unet_2d_condition import UNet2DConditionModel -from tensorrt_llm.models.unet.weights import load_from_hf_unet -from tensorrt_llm.network import net_guard - -parser = argparse.ArgumentParser(description='build the UNet TensorRT engine.') -parser.add_argument('--model_dir', - type=str, - default='stabilityai/stable-diffusion-xl-base-1.0') -parser.add_argument('--size', type=int, default=1024, help='image size') -parser.add_argument('--output_dir', - type=str, - default=None, - help='output directory') - -args = parser.parse_args() - -model_dir = args.model_dir -size = args.size -sample_size = size // 8 - -world_size = tensorrt_llm.mpi_world_size() -rank = tensorrt_llm.mpi_rank() -output_dir = f'sdxl_s{size}_w{world_size}' if args.output_dir is None else args.output_dir -if rank == 0 and not os.path.exists(output_dir): - os.makedirs(output_dir) - -device_per_batch = world_size // 2 if world_size > 1 else 1 -batch_group = 2 if world_size > 1 else 1 - -# Use tp_size to indicate the size of patch parallelism -# Use pp_size to indicate the size of batch parallelism -mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=device_per_batch, - pp_size=batch_group) - -torch.cuda.set_device(tensorrt_llm.mpi_rank()) - -tensorrt_llm.logger.set_level('verbose') -builder = Builder() -builder_config = builder.create_builder_config( - name='UNet2DConditionModel', - precision='float16', - timing_cache='model.cache', - profiling_verbosity='detailed', - tensor_parallel=world_size, - precision_constraints= - None, # do not use obey or the precision error will be too large -) - -pipeline = DiffusionPipeline.from_pretrained(model_dir, - torch_dtype=torch.float16) -model = UNet2DConditionModel( - sample_size=sample_size, - in_channels=4, - out_channels=4, - center_input_sample=False, - flip_sin_to_cos=True, - freq_shift=0, - down_block_types=("DownBlock2D", "CrossAttnDownBlock2D", - "CrossAttnDownBlock2D"), - up_block_types=("CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D"), - block_out_channels=(320, 640, 1280), - layers_per_block=2, - downsample_padding=1, - mid_block_scale_factor=1.0, - act_fn="silu", - norm_num_groups=32, - norm_eps=1e-5, - cross_attention_dim=2048, - attention_head_dim=[5, 10, 20], - addition_embed_type="text_time", - addition_time_embed_dim=256, - projection_class_embeddings_input_dim=2816, - transformer_layers_per_block=[1, 2, 10], - use_linear_projection=True, - dtype=trt.float16, -) - -load_from_hf_unet(pipeline.unet, model) -model = DistriUNetPP(model, mapping) - -# Module -> Network -network = builder.create_network() -network.plugin_config.to_legacy_setting() -if mapping.world_size > 1: - network.plugin_config.set_nccl_plugin('float16') - -with net_guard(network): - # Prepare - network.set_named_parameters(model.named_parameters()) - - # Forward - sample = tensorrt_llm.Tensor( - name='sample', - dtype=trt.float16, - shape=[2, 4, sample_size, sample_size], - ) - timesteps = tensorrt_llm.Tensor( - name='timesteps', - dtype=trt.float16, - shape=[ - 1, - ], - ) - encoder_hidden_states = tensorrt_llm.Tensor( - name='encoder_hidden_states', - dtype=trt.float16, - shape=[2, 77, 2048], - ) - text_embeds = tensorrt_llm.Tensor( - name='text_embeds', - dtype=trt.float16, - shape=[2, 1280], - ) - time_ids = tensorrt_llm.Tensor( - name='time_ids', - dtype=trt.float16, - shape=[2, 6], - ) - - output = model(sample, timesteps, encoder_hidden_states, text_embeds, - time_ids) - - # Mark outputs - output_dtype = trt.float16 - output.mark_output('pred', output_dtype) - -# Network -> Engine -engine = builder.build_engine(network, builder_config) -assert engine is not None, 'Failed to build engine.' - -engine_name = f'sdxl_unet_s{size}_w{world_size}_r{rank}.engine' -engine_path = os.path.join(output_dir, engine_name) -with open(engine_path, 'wb') as f: - f.write(engine) -builder.save_config(builder_config, os.path.join(output_dir, 'config.json')) diff --git a/examples/models/contrib/sdxl/pipeline_stable_diffusion_xl.py b/examples/models/contrib/sdxl/pipeline_stable_diffusion_xl.py deleted file mode 100755 index 018f3fad783f..000000000000 --- a/examples/models/contrib/sdxl/pipeline_stable_diffusion_xl.py +++ /dev/null @@ -1,1365 +0,0 @@ -# Copyright 2023 The HuggingFace Team. All rights reserved. -# -# 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. - -import inspect -import json -import os -from typing import Any, Callable, Dict, List, Optional, Tuple, Union - -import tensorrt as trt -import torch -from diffusers.image_processor import PipelineImageInput, VaeImageProcessor -from diffusers.loaders import (FromSingleFileMixin, IPAdapterMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin) -from diffusers.models import AutoencoderKL, UNet2DConditionModel -from diffusers.models.attention_processor import (AttnProcessor2_0, - LoRAAttnProcessor2_0, - LoRAXFormersAttnProcessor, - XFormersAttnProcessor) -from diffusers.models.lora import adjust_lora_scale_text_encoder -from diffusers.pipelines.pipeline_utils import DiffusionPipeline -from diffusers.pipelines.stable_diffusion_xl.pipeline_output import \ - StableDiffusionXLPipelineOutput -from diffusers.schedulers import KarrasDiffusionSchedulers -from diffusers.utils import (USE_PEFT_BACKEND, deprecate, - is_invisible_watermark_available, - is_torch_xla_available, logging, - replace_example_docstring, scale_lora_layers, - unscale_lora_layers) -from diffusers.utils.torch_utils import randn_tensor -from transformers import (CLIPImageProcessor, CLIPTextModel, - CLIPTextModelWithProjection, CLIPTokenizer, - CLIPVisionModelWithProjection) - -import tensorrt_llm -from tensorrt_llm.runtime import Session, TensorInfo - -if is_invisible_watermark_available(): - from diffusers.pipelines.stable_diffusion_xl.watermark import \ - StableDiffusionXLWatermarker - -if is_torch_xla_available(): - import torch_xla.core.xla_model as xm - - XLA_AVAILABLE = True -else: - XLA_AVAILABLE = False - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> import torch - >>> from diffusers import StableDiffusionXLPipeline - - >>> pipe = StableDiffusionXLPipeline.from_pretrained( - ... "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 - ... ) - >>> pipe = pipe.to("cuda") - - >>> prompt = "a photo of an astronaut riding a horse on mars" - >>> image = pipe(prompt).images[0] - ``` -""" - - -# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg -def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): - """ - Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and - Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4 - """ - std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), - keepdim=True) - std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) - # rescale the results from guidance (fixes overexposure) - noise_pred_rescaled = noise_cfg * (std_text / std_cfg) - # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images - noise_cfg = guidance_rescale * noise_pred_rescaled + \ - (1 - guidance_rescale) * noise_cfg - return noise_cfg - - -# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps -def retrieve_timesteps( - scheduler, - num_inference_steps: Optional[int] = None, - device: Optional[Union[str, torch.device]] = None, - timesteps: Optional[List[int]] = None, - **kwargs, -): - """ - Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles - custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. - - Args: - scheduler (`SchedulerMixin`): - The scheduler to get timesteps from. - num_inference_steps (`int`): - The number of diffusion steps used when generating samples with a pre-trained model. If used, - `timesteps` must be `None`. - device (`str` or `torch.device`, *optional*): - The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`List[int]`, *optional*): - Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default - timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps` - must be `None`. - - Returns: - `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the - second element is the number of inference steps. - """ - if timesteps is not None: - accepts_timesteps = "timesteps" in set( - inspect.signature(scheduler.set_timesteps).parameters.keys()) - if not accepts_timesteps: - raise ValueError( - f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" - f" timestep schedules. Please check whether you are using the correct scheduler." - ) - scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) - timesteps = scheduler.timesteps - num_inference_steps = len(timesteps) - else: - scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) - timesteps = scheduler.timesteps - return timesteps, num_inference_steps - - -class StableDiffusionXLPipeline( - DiffusionPipeline, - FromSingleFileMixin, - StableDiffusionXLLoraLoaderMixin, - TextualInversionLoaderMixin, - IPAdapterMixin, -): - r""" - Pipeline for text-to-image generation using Stable Diffusion XL. - - This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the - library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) - - In addition the pipeline inherits the following loading methods: - - *LoRA*: [`loaders.StableDiffusionXLLoraLoaderMixin.load_lora_weights`] - - *Ckpt*: [`loaders.FromSingleFileMixin.from_single_file`] - - as well as the following saving methods: - - *LoRA*: [`loaders.StableDiffusionXLLoraLoaderMixin.save_lora_weights`] - - Args: - vae ([`AutoencoderKL`]): - Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. - text_encoder ([`CLIPTextModel`]): - Frozen text-encoder. Stable Diffusion XL uses the text portion of - [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically - the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant. - text_encoder_2 ([` CLIPTextModelWithProjection`]): - Second frozen text-encoder. Stable Diffusion XL uses the text and pool portion of - [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection), - specifically the - [laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k) - variant. - tokenizer (`CLIPTokenizer`): - Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). - tokenizer_2 (`CLIPTokenizer`): - Second Tokenizer of class - [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer). - unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents. - scheduler ([`SchedulerMixin`]): - A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of - [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. - force_zeros_for_empty_prompt (`bool`, *optional*, defaults to `"True"`): - Whether the negative prompt embeddings shall be forced to always be set to 0. Also see the config of - `stabilityai/stable-diffusion-xl-base-1-0`. - add_watermarker (`bool`, *optional*): - Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to - watermark output images. If not defined, it will default to True if the package is installed, otherwise no - watermarker will be used. - """ - - model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae" - _optional_components = [ - "tokenizer", - "tokenizer_2", - "text_encoder", - "text_encoder_2", - "image_encoder", - "feature_extractor", - ] - _callback_tensor_inputs = [ - "latents", - "prompt_embeds", - "negative_prompt_embeds", - "add_text_embeds", - "add_time_ids", - "negative_pooled_prompt_embeds", - "negative_add_time_ids", - ] - - def __init__( - self, - vae: AutoencoderKL, - text_encoder: CLIPTextModel, - text_encoder_2: CLIPTextModelWithProjection, - tokenizer: CLIPTokenizer, - tokenizer_2: CLIPTokenizer, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - image_encoder: CLIPVisionModelWithProjection = None, - feature_extractor: CLIPImageProcessor = None, - force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, - ): - super().__init__() - - self.register_modules( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - unet=unet, - scheduler=scheduler, - image_encoder=image_encoder, - feature_extractor=feature_extractor, - ) - self.register_to_config( - force_zeros_for_empty_prompt=force_zeros_for_empty_prompt) - self.vae_scale_factor = 2**(len(self.vae.config.block_out_channels) - 1) - self.image_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor) - - self.default_sample_size = self.unet.config.sample_size - - add_watermarker = add_watermarker if add_watermarker is not None else is_invisible_watermark_available( - ) - - if add_watermarker: - self.watermark = StableDiffusionXLWatermarker() - else: - self.watermark = None - - self.execution_device = torch.device('cpu') - self.engine = {} - - def to( - self, - torch_device: Optional[Union[str, torch.device]] = None, - torch_dtype: Optional[torch.dtype] = None, - silence_dtype_warnings: bool = False, - ): - super().to(torch_device) - if isinstance(torch_device, str): - torch_device = torch.device(torch_device) - self.execution_device = torch_device - return self - - def prepare(self, path, size): - self.unet.cpu() - torch.cuda.empty_cache() - - def trt_dtype_to_torch(dtype): - if dtype == trt.float16: - return torch.float16 - elif dtype == trt.float32: - return torch.float32 - elif dtype == trt.int32: - return torch.int32 - else: - raise TypeError("%s is not supported" % dtype) - - config_path = os.path.join(path, 'config.json') - with open(config_path, 'r') as f: - config = json.load(f) - config['builder_config']['precision'] - world_size = config['builder_config']['tensor_parallel'] - - runtime_world_size = tensorrt_llm.mpi_world_size() - assert world_size == runtime_world_size, f'Engine world size ({world_size}) != Runtime world size ({runtime_world_size})' - runtime_rank = tensorrt_llm.mpi_rank() if world_size > 1 else 0 - torch.cuda.set_device(runtime_rank) - - serialize_file = f'sdxl_unet_s{size}_w{world_size}_r{runtime_rank}.engine' - serialize_path = os.path.join(path, serialize_file) - self.stream = torch.cuda.current_stream().cuda_stream - print(f'Loading engine from {serialize_path}') - with open(serialize_path, 'rb') as f: - engine_buffer = f.read() - print(f'Creating session from engine') - self.session = Session.from_serialized_engine(engine_buffer) - - output_info = self.session.infer_shapes([ - TensorInfo('sample', trt.DataType.HALF, - [2, 4, size // 8, size // 8]), - TensorInfo('timesteps', trt.DataType.HALF, [ - 1, - ]), - TensorInfo('encoder_hidden_states', trt.DataType.HALF, - [2, 77, 2048]), - TensorInfo('text_embeds', trt.DataType.HALF, [2, 1280]), - TensorInfo('time_ids', trt.DataType.HALF, [2, 6]), - ]) - self.outputs = { - t.name: - torch.empty(tuple(t.shape), - dtype=trt_dtype_to_torch(t.dtype), - device='cuda') - for t in output_info - } - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_vae_slicing - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.vae.enable_slicing() - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_vae_slicing - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_vae_tiling - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - self.vae.enable_tiling() - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_vae_tiling - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - - def encode_prompt( - self, - prompt: str, - prompt_2: Optional[str] = None, - device: Optional[torch.device] = None, - num_images_per_prompt: int = 1, - do_classifier_free_guidance: bool = True, - negative_prompt: Optional[str] = None, - negative_prompt_2: Optional[str] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - lora_scale: Optional[float] = None, - clip_skip: Optional[int] = None, - ): - r""" - Encodes the prompt into text encoder hidden states. - - Args: - prompt (`str` or `List[str]`, *optional*): - prompt to be encoded - prompt_2 (`str` or `List[str]`, *optional*): - The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is - used in both text-encoders - device: (`torch.device`): - torch device - num_images_per_prompt (`int`): - number of images that should be generated per prompt - do_classifier_free_guidance (`bool`): - whether to use classifier free guidance or not - negative_prompt (`str` or `List[str]`, *optional*): - The prompt or prompts not to guide the image generation. If not defined, one has to pass - `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is - less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): - The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and - `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders - prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not - provided, text embeddings will be generated from `prompt` input argument. - negative_prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input - argument. - pooled_prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. - If not provided, pooled text embeddings will be generated from `prompt` input argument. - negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` - input argument. - lora_scale (`float`, *optional*): - A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded. - clip_skip (`int`, *optional*): - Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that - the output of the pre-final layer will be used for computing the prompt embeddings. - """ - device = self.execution_device - - # set lora scale so that monkey patched LoRA - # function of text encoder can correctly access it - if lora_scale is not None and isinstance( - self, StableDiffusionXLLoraLoaderMixin): - self._lora_scale = lora_scale - - # dynamically adjust the LoRA scale - if self.text_encoder is not None: - if not USE_PEFT_BACKEND: - adjust_lora_scale_text_encoder(self.text_encoder, - lora_scale) - else: - scale_lora_layers(self.text_encoder, lora_scale) - - if self.text_encoder_2 is not None: - if not USE_PEFT_BACKEND: - adjust_lora_scale_text_encoder(self.text_encoder_2, - lora_scale) - else: - scale_lora_layers(self.text_encoder_2, lora_scale) - - prompt = [prompt] if isinstance(prompt, str) else prompt - - if prompt is not None: - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - # Define tokenizers and text encoders - tokenizers = [self.tokenizer, self.tokenizer_2 - ] if self.tokenizer is not None else [self.tokenizer_2] - text_encoders = ([ - self.text_encoder, self.text_encoder_2 - ] if self.text_encoder is not None else [self.text_encoder_2]) - - if prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - # textual inversion: procecss multi-vector tokens if necessary - prompt_embeds_list = [] - prompts = [prompt, prompt_2] - for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, - text_encoders): - if isinstance(self, TextualInversionLoaderMixin): - prompt = self.maybe_convert_prompt(prompt, tokenizer) - - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=tokenizer.model_max_length, - truncation=True, - return_tensors="pt", - ) - - text_input_ids = text_inputs.input_ids - untruncated_ids = tokenizer(prompt, - padding="longest", - return_tensors="pt").input_ids - - if untruncated_ids.shape[ - -1] >= text_input_ids.shape[-1] and not torch.equal( - text_input_ids, untruncated_ids): - removed_text = tokenizer.batch_decode( - untruncated_ids[:, tokenizer.model_max_length - 1:-1]) - logger.warning( - "The following part of your input was truncated because CLIP can only handle sequences up to" - f" {tokenizer.model_max_length} tokens: {removed_text}") - - prompt_embeds = text_encoder(text_input_ids.to(device), - output_hidden_states=True) - - # We are only ALWAYS interested in the pooled output of the final text encoder - pooled_prompt_embeds = prompt_embeds[0] - if clip_skip is None: - prompt_embeds = prompt_embeds.hidden_states[-2] - else: - # "2" because SDXL always indexes from the penultimate layer. - prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + - 2)] - - prompt_embeds_list.append(prompt_embeds) - - prompt_embeds = torch.concat(prompt_embeds_list, dim=-1) - - # get unconditional embeddings for classifier free guidance - zero_out_negative_prompt = negative_prompt is None and self.config.force_zeros_for_empty_prompt - if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt: - negative_prompt_embeds = torch.zeros_like(prompt_embeds) - negative_pooled_prompt_embeds = torch.zeros_like( - pooled_prompt_embeds) - elif do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt or "" - negative_prompt_2 = negative_prompt_2 or negative_prompt - - # normalize str to list - negative_prompt = batch_size * \ - [negative_prompt] if isinstance( - negative_prompt, str) else negative_prompt - negative_prompt_2 = (batch_size * [negative_prompt_2] if isinstance( - negative_prompt_2, str) else negative_prompt_2) - - uncond_tokens: List[str] - if prompt is not None and type(prompt) is not type(negative_prompt): - raise TypeError( - f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" - f" {type(prompt)}.") - elif batch_size != len(negative_prompt): - raise ValueError( - f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" - f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" - " the batch size of `prompt`.") - else: - uncond_tokens = [negative_prompt, negative_prompt_2] - - negative_prompt_embeds_list = [] - for negative_prompt, tokenizer, text_encoder in zip( - uncond_tokens, tokenizers, text_encoders): - if isinstance(self, TextualInversionLoaderMixin): - negative_prompt = self.maybe_convert_prompt( - negative_prompt, tokenizer) - - max_length = prompt_embeds.shape[1] - uncond_input = tokenizer( - negative_prompt, - padding="max_length", - max_length=max_length, - truncation=True, - return_tensors="pt", - ) - - negative_prompt_embeds = text_encoder( - uncond_input.input_ids.to(device), - output_hidden_states=True, - ) - # We are only ALWAYS interested in the pooled output of the final text encoder - negative_pooled_prompt_embeds = negative_prompt_embeds[0] - negative_prompt_embeds = negative_prompt_embeds.hidden_states[ - -2] - - negative_prompt_embeds_list.append(negative_prompt_embeds) - - negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, - dim=-1) - - if self.text_encoder_2 is not None: - prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, - device=device) - else: - prompt_embeds = prompt_embeds.to(dtype=self.unet.dtype, - device=device) - - bs_embed, seq_len, _ = prompt_embeds.shape - # duplicate text embeddings for each generation per prompt, using mps friendly method - prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) - prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, - seq_len, -1) - - if do_classifier_free_guidance: - # duplicate unconditional embeddings for each generation per prompt, using mps friendly method - seq_len = negative_prompt_embeds.shape[1] - - if self.text_encoder_2 is not None: - negative_prompt_embeds = negative_prompt_embeds.to( - dtype=self.text_encoder_2.dtype, device=device) - else: - negative_prompt_embeds = negative_prompt_embeds.to( - dtype=self.unet.dtype, device=device) - - negative_prompt_embeds = negative_prompt_embeds.repeat( - 1, num_images_per_prompt, 1) - negative_prompt_embeds = negative_prompt_embeds.view( - batch_size * num_images_per_prompt, seq_len, -1) - - pooled_prompt_embeds = pooled_prompt_embeds.repeat( - 1, num_images_per_prompt).view(bs_embed * num_images_per_prompt, -1) - if do_classifier_free_guidance: - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat( - 1, num_images_per_prompt).view(bs_embed * num_images_per_prompt, - -1) - - if self.text_encoder is not None: - if isinstance( - self, - StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND: - # Retrieve the original scale by scaling back the LoRA layers - unscale_lora_layers(self.text_encoder, lora_scale) - - if self.text_encoder_2 is not None: - if isinstance( - self, - StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND: - # Retrieve the original scale by scaling back the LoRA layers - unscale_lora_layers(self.text_encoder_2, lora_scale) - - return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image - def encode_image(self, image, device, num_images_per_prompt): - dtype = next(self.image_encoder.parameters()).dtype - - if not isinstance(image, torch.Tensor): - image = self.feature_extractor(image, - return_tensors="pt").pixel_values - - image = image.to(device=device, dtype=dtype) - image_embeds = self.image_encoder(image).image_embeds - image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, - dim=0) - - uncond_image_embeds = torch.zeros_like(image_embeds) - return image_embeds, uncond_image_embeds - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs - def prepare_extra_step_kwargs(self, generator, eta): - # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature - # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. - # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 - # and should be between [0, 1] - - accepts_eta = "eta" in set( - inspect.signature(self.scheduler.step).parameters.keys()) - extra_step_kwargs = {} - if accepts_eta: - extra_step_kwargs["eta"] = eta - - # check if the scheduler accepts generator - accepts_generator = "generator" in set( - inspect.signature(self.scheduler.step).parameters.keys()) - if accepts_generator: - extra_step_kwargs["generator"] = generator - return extra_step_kwargs - - def check_inputs( - self, - prompt, - prompt_2, - height, - width, - callback_steps, - negative_prompt=None, - negative_prompt_2=None, - prompt_embeds=None, - negative_prompt_embeds=None, - pooled_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - ): - if height % 8 != 0 or width % 8 != 0: - raise ValueError( - f"`height` and `width` have to be divisible by 8 but are {height} and {width}." - ) - - if callback_steps is not None and (not isinstance(callback_steps, int) - or callback_steps <= 0): - raise ValueError( - f"`callback_steps` has to be a positive integer but is {callback_steps} of type" - f" {type(callback_steps)}.") - - if callback_on_step_end_tensor_inputs is not None and not all( - k in self._callback_tensor_inputs - for k in callback_on_step_end_tensor_inputs): - raise ValueError( - f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" - ) - - if prompt is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" - " only forward one of the two.") - elif prompt_2 is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to" - " only forward one of the two.") - elif prompt is None and prompt_embeds is None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." - ) - elif prompt is not None and (not isinstance(prompt, str) - and not isinstance(prompt, list)): - raise ValueError( - f"`prompt` has to be of type `str` or `list` but is {type(prompt)}" - ) - elif prompt_2 is not None and (not isinstance(prompt_2, str) - and not isinstance(prompt_2, list)): - raise ValueError( - f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}" - ) - - if negative_prompt is not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - elif negative_prompt_2 is not None and negative_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:" - f" {negative_prompt_embeds}. Please make sure to only forward one of the two." - ) - - if prompt_embeds is not None and negative_prompt_embeds is not None: - if prompt_embeds.shape != negative_prompt_embeds.shape: - raise ValueError( - "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" - f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" - f" {negative_prompt_embeds.shape}.") - - if prompt_embeds is not None and pooled_prompt_embeds is None: - raise ValueError( - "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`." - ) - - if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None: - raise ValueError( - "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`." - ) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents - def prepare_latents(self, - batch_size, - num_channels_latents, - height, - width, - dtype, - device, - generator, - latents=None): - shape = (batch_size, num_channels_latents, - height // self.vae_scale_factor, - width // self.vae_scale_factor) - if isinstance(generator, list) and len(generator) != batch_size: - raise ValueError( - f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" - f" size of {batch_size}. Make sure the batch size matches the length of the generators." - ) - - if latents is None: - latents = randn_tensor(shape, - generator=generator, - device=device, - dtype=dtype) - else: - latents = latents.to(device) - - # scale the initial noise by the standard deviation required by the scheduler - latents = latents * self.scheduler.init_noise_sigma - return latents - - def _get_add_time_ids(self, - original_size, - crops_coords_top_left, - target_size, - dtype, - text_encoder_projection_dim=None): - add_time_ids = list(original_size + crops_coords_top_left + target_size) - - passed_add_embed_dim = ( - self.unet.config.addition_time_embed_dim * len(add_time_ids) + - text_encoder_projection_dim) - expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features - - if expected_add_embed_dim != passed_add_embed_dim: - raise ValueError( - f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`." - ) - - add_time_ids = torch.tensor([add_time_ids], dtype=dtype) - return add_time_ids - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae - def upcast_vae(self): - dtype = self.vae.dtype - self.vae.to(dtype=torch.float32) - use_torch_2_0_or_xformers = isinstance( - self.vae.decoder.mid_block.attentions[0].processor, - ( - AttnProcessor2_0, - XFormersAttnProcessor, - LoRAXFormersAttnProcessor, - LoRAAttnProcessor2_0, - ), - ) - # if xformers or torch_2_0 is used attention block does not need - # to be in float32 which can save lots of memory - if use_torch_2_0_or_xformers: - self.vae.post_quant_conv.to(dtype) - self.vae.decoder.conv_in.to(dtype) - self.vae.decoder.mid_block.to(dtype) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_freeu - def enable_freeu(self, s1: float, s2: float, b1: float, b2: float): - r"""Enables the FreeU mechanism as in https://arxiv.org/abs/2309.11497. - - The suffixes after the scaling factors represent the stages where they are being applied. - - Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of the values - that are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL. - - Args: - s1 (`float`): - Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to - mitigate "oversmoothing effect" in the enhanced denoising process. - s2 (`float`): - Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to - mitigate "oversmoothing effect" in the enhanced denoising process. - b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features. - b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features. - """ - if not hasattr(self, "unet"): - raise ValueError("The pipeline must have `unet` for using FreeU.") - self.unet.enable_freeu(s1=s1, s2=s2, b1=b1, b2=b2) - - # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_freeu - def disable_freeu(self): - """Disables the FreeU mechanism if enabled.""" - self.unet.disable_freeu() - - # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding - def get_guidance_scale_embedding(self, - w, - embedding_dim=512, - dtype=torch.float32): - """ - See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 - - Args: - timesteps (`torch.Tensor`): - generate embedding vectors at these timesteps - embedding_dim (`int`, *optional*, defaults to 512): - dimension of the embeddings to generate - dtype: - data type of the generated embeddings - - Returns: - `torch.FloatTensor`: Embedding vectors with shape `(len(timesteps), embedding_dim)` - """ - assert len(w.shape) == 1 - w = w * 1000.0 - - half_dim = embedding_dim // 2 - emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) - emb = w.to(dtype)[:, None] * emb[None, :] - emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) - if embedding_dim % 2 == 1: # zero pad - emb = torch.nn.functional.pad(emb, (0, 1)) - assert emb.shape == (w.shape[0], embedding_dim) - return emb - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - @property - def clip_skip(self): - return self._clip_skip - - # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) - # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` - # corresponds to doing no classifier free guidance. - @property - def do_classifier_free_guidance(self): - return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None - - @property - def cross_attention_kwargs(self): - return self._cross_attention_kwargs - - @property - def denoising_end(self): - return self._denoising_end - - @property - def num_timesteps(self): - return self._num_timesteps - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - timesteps: List[int] = None, - denoising_end: Optional[float] = None, - guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: float = 0.0, - generator: Optional[Union[torch.Generator, - List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - output_type: Optional[str] = "pil", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - guidance_rescale: float = 0.0, - original_size: Optional[Tuple[int, int]] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Optional[Tuple[int, int]] = None, - negative_original_size: Optional[Tuple[int, int]] = None, - negative_crops_coords_top_left: Tuple[int, int] = (0, 0), - negative_target_size: Optional[Tuple[int, int]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], - **kwargs, - ): - r""" - Function invoked when calling the pipeline for generation. - - Args: - prompt (`str` or `List[str]`, *optional*): - The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. - instead. - prompt_2 (`str` or `List[str]`, *optional*): - The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is - used in both text-encoders - height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): - The height in pixels of the generated image. This is set to 1024 by default for the best results. - Anything below 512 pixels won't work well for - [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) - and checkpoints that are not specifically fine-tuned on low resolutions. - width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): - The width in pixels of the generated image. This is set to 1024 by default for the best results. - Anything below 512 pixels won't work well for - [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) - and checkpoints that are not specifically fine-tuned on low resolutions. - num_inference_steps (`int`, *optional*, defaults to 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - timesteps (`List[int]`, *optional*): - Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument - in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is - passed will be used. Must be in descending order. - denoising_end (`float`, *optional*): - When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be - completed before it is intentionally prematurely terminated. As a result, the returned sample will - still retain a substantial amount of noise as determined by the discrete timesteps selected by the - scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a - "Mixture of Denoisers" multi-pipeline setup, as elaborated in [**Refining the Image - Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output) - guidance_scale (`float`, *optional*, defaults to 5.0): - Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). - `guidance_scale` is defined as `w` of equation 2. of [Imagen - Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > - 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, - usually at the expense of lower image quality. - negative_prompt (`str` or `List[str]`, *optional*): - The prompt or prompts not to guide the image generation. If not defined, one has to pass - `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is - less than `1`). - negative_prompt_2 (`str` or `List[str]`, *optional*): - The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and - `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - eta (`float`, *optional*, defaults to 0.0): - Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to - [`schedulers.DDIMScheduler`], will be ignored for others. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): - One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) - to make generation deterministic. - latents (`torch.FloatTensor`, *optional*): - Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image - generation. Can be used to tweak the same generation with different prompts. If not provided, a latents - tensor will ge generated by sampling using the supplied random `generator`. - prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not - provided, text embeddings will be generated from `prompt` input argument. - negative_prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input - argument. - pooled_prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. - If not provided, pooled text embeddings will be generated from `prompt` input argument. - negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` - input argument. - ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between - [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead - of a plain tuple. - cross_attention_kwargs (`dict`, *optional*): - A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under - `self.processor` in - [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - guidance_rescale (`float`, *optional*, defaults to 0.0): - Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are - Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of - [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). - Guidance rescale factor should fix overexposure when using zero terminal SNR. - original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): - If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled. - `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as - explained in section 2.2 of - [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): - `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position - `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting - `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of - [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): - For most cases, `target_size` should be set to the desired height and width of the generated image. If - not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in - section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). - negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): - To negatively condition the generation process based on a specific image resolution. Part of SDXL's - micro-conditioning as explained in section 2.2 of - [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more - information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)): - To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's - micro-conditioning as explained in section 2.2 of - [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more - information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)): - To negatively condition the generation process based on a target image resolution. It should be as same - as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of - [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more - information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208. - callback_on_step_end (`Callable`, *optional*): - A function that calls at the end of each denoising steps during the inference. The function is called - with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, - callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by - `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): - The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list - will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the - `._callback_tensor_inputs` attribute of your pipeline class. - - Examples: - - Returns: - [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`: - [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a - `tuple`. When returning a tuple, the first element is a list with the generated images. - """ - - callback = kwargs.pop("callback", None) - callback_steps = kwargs.pop("callback_steps", None) - - if callback is not None: - deprecate( - "callback", - "1.0.0", - "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - if callback_steps is not None: - deprecate( - "callback_steps", - "1.0.0", - "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", - ) - - # 0. Default height and width to unet - height = height or self.default_sample_size * self.vae_scale_factor - width = width or self.default_sample_size * self.vae_scale_factor - - original_size = original_size or (height, width) - target_size = target_size or (height, width) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - prompt_2, - height, - width, - callback_steps, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - self._clip_skip = clip_skip - self._cross_attention_kwargs = cross_attention_kwargs - self._denoising_end = denoising_end - - # 2. Define call parameters - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - device = self.execution_device - - # 3. Encode input prompt - lora_scale = (self.cross_attention_kwargs.get("scale", None) - if self.cross_attention_kwargs is not None else None) - - ( - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - device=device, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - lora_scale=lora_scale, - clip_skip=self.clip_skip, - ) - - # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, num_inference_steps, device, timesteps) - - # 5. Prepare latent variables - num_channels_latents = self.unet.config.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Prepare added time ids & embeddings - add_text_embeds = pooled_prompt_embeds - if self.text_encoder_2 is None: - text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1]) - else: - text_encoder_projection_dim = self.text_encoder_2.config.projection_dim - - add_time_ids = self._get_add_time_ids( - original_size, - crops_coords_top_left, - target_size, - dtype=prompt_embeds.dtype, - text_encoder_projection_dim=text_encoder_projection_dim, - ) - if negative_original_size is not None and negative_target_size is not None: - negative_add_time_ids = self._get_add_time_ids( - negative_original_size, - negative_crops_coords_top_left, - negative_target_size, - dtype=prompt_embeds.dtype, - text_encoder_projection_dim=text_encoder_projection_dim, - ) - else: - negative_add_time_ids = add_time_ids - - if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], - dim=0) - add_text_embeds = torch.cat( - [negative_pooled_prompt_embeds, add_text_embeds], dim=0) - add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], - dim=0) - - prompt_embeds = prompt_embeds.to(device) - add_text_embeds = add_text_embeds.to(device) - add_time_ids = add_time_ids.to(device).repeat( - batch_size * num_images_per_prompt, 1) - - if ip_adapter_image is not None: - image_embeds, negative_image_embeds = self.encode_image( - ip_adapter_image, device, num_images_per_prompt) - if self.do_classifier_free_guidance: - image_embeds = torch.cat([negative_image_embeds, image_embeds]) - image_embeds = image_embeds.to(device) - - # 8. Denoising loop - num_warmup_steps = max( - len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - # 8.1 Apply denoising_end - if (self.denoising_end is not None - and isinstance(self.denoising_end, float) - and self.denoising_end > 0 and self.denoising_end < 1): - discrete_timestep_cutoff = int( - round(self.scheduler.config.num_train_timesteps - - (self.denoising_end * - self.scheduler.config.num_train_timesteps))) - num_inference_steps = len( - list( - filter(lambda ts: ts >= discrete_timestep_cutoff, - timesteps))) - timesteps = timesteps[:num_inference_steps] - - # 9. Optionally get Guidance Scale Embedding - timestep_cond = None - if self.unet.config.time_cond_proj_dim is not None: - guidance_scale_tensor = torch.tensor(self.guidance_scale - - 1).repeat( - batch_size * - num_images_per_prompt) - timestep_cond = self.get_guidance_scale_embedding( - guidance_scale_tensor, - embedding_dim=self.unet.config.time_cond_proj_dim).to( - device=device, dtype=latents.dtype) - - self._num_timesteps = len(timesteps) - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t in enumerate(timesteps): - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat( - [latents] * - 2) if self.do_classifier_free_guidance else latents - - latent_model_input = self.scheduler.scale_model_input( - latent_model_input, t) - - # predict the noise residual - added_cond_kwargs = { - "text_embeds": add_text_embeds, - "time_ids": add_time_ids - } - if ip_adapter_image is not None: - added_cond_kwargs["image_embeds"] = image_embeds - - t = t.to(latent_model_input.dtype) - feed_dict = { - 'sample': latent_model_input, - 'timesteps': t.unsqueeze(0), - 'encoder_hidden_states': prompt_embeds, - 'text_embeds': add_text_embeds, - 'time_ids': add_time_ids, - } - ok = self.session.run(feed_dict, self.outputs, self.stream) - assert ok, "Runtime execution failed" - noise_pred = self.outputs['pred'] - torch.cuda.synchronize() - - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + self.guidance_scale * \ - (noise_pred_text - noise_pred_uncond) - - if self.do_classifier_free_guidance and self.guidance_rescale > 0.0: - # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf - noise_pred = rescale_noise_cfg( - noise_pred, - noise_pred_text, - guidance_rescale=self.guidance_rescale) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, - t, - latents, - **extra_step_kwargs, - return_dict=False)[0] - - if callback_on_step_end is not None: - callback_kwargs = {} - for k in callback_on_step_end_tensor_inputs: - callback_kwargs[k] = locals()[k] - callback_outputs = callback_on_step_end( - self, i, t, callback_kwargs) - - latents = callback_outputs.pop("latents", latents) - prompt_embeds = callback_outputs.pop( - "prompt_embeds", prompt_embeds) - negative_prompt_embeds = callback_outputs.pop( - "negative_prompt_embeds", negative_prompt_embeds) - add_text_embeds = callback_outputs.pop( - "add_text_embeds", add_text_embeds) - negative_pooled_prompt_embeds = callback_outputs.pop( - "negative_pooled_prompt_embeds", - negative_pooled_prompt_embeds) - add_time_ids = callback_outputs.pop("add_time_ids", - add_time_ids) - negative_add_time_ids = callback_outputs.pop( - "negative_add_time_ids", negative_add_time_ids) - - # call the callback, if provided - if i == len(timesteps) - 1 or ( - (i + 1) > num_warmup_steps and - (i + 1) % self.scheduler.order == 0): - progress_bar.update() - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) - - if XLA_AVAILABLE: - xm.mark_step() - - if not output_type == "latent": - # make sure the VAE is in float32 mode, as it overflows in float16 - needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast - - if needs_upcasting: - self.upcast_vae() - latents = latents.to( - next(iter(self.vae.post_quant_conv.parameters())).dtype) - - image = self.vae.decode(latents / self.vae.config.scaling_factor, - return_dict=False)[0] - - # cast back to fp16 if needed - if needs_upcasting: - self.vae.to(dtype=torch.float16) - else: - image = latents - - if not output_type == "latent": - # apply watermark if available - if self.watermark is not None: - image = self.watermark.apply_watermark(image) - - image = self.image_processor.postprocess(image, - output_type=output_type) - - # Offload all models - self.maybe_free_model_hooks() - - if not return_dict: - return (image, ) - - return StableDiffusionXLPipelineOutput(images=image) diff --git a/examples/models/contrib/sdxl/run_sdxl.py b/examples/models/contrib/sdxl/run_sdxl.py deleted file mode 100755 index 077b93f9ab7e..000000000000 --- a/examples/models/contrib/sdxl/run_sdxl.py +++ /dev/null @@ -1,96 +0,0 @@ -import argparse -import time - -import numpy as np -import torch -from pipeline_stable_diffusion_xl import StableDiffusionXLPipeline - -import tensorrt_llm - -world_size = tensorrt_llm.mpi_world_size() -rank = tensorrt_llm.mpi_rank() - - -def parseArgs(): - parser = argparse.ArgumentParser( - description='run SDXL with the UNet TensorRT engine.') - parser.add_argument('--model_dir', - type=str, - default='stabilityai/stable-diffusion-xl-base-1.0') - parser.add_argument('--size', type=int, default=1024) - parser.add_argument('--seed', type=int, default=233) - parser.add_argument('--num_inference_steps', type=int, default=50) - parser.add_argument( - '--prompt', - type=str, - default= - "masterpiece, gouache painting, 1girl, distant view, lone boat, willow trees" - ) - parser.add_argument('--engine_dir', - type=str, - default=None, - help='engine directory') - parser.add_argument('--num-warmup-runs', type=int, default=3) - parser.add_argument('--avg-runs', type=int, default=10) - parser.add_argument("--ignore_ratio", - type=float, - default=0.2, - help="Ignored ratio of the slowest and fastest steps") - parser.add_argument("--output", - type=str, - default="output.png", - help="Output file name") - return parser.parse_args() - - -if __name__ == "__main__": - args = parseArgs() - model_dir = args.model_dir - size = args.size - seed = args.seed - prompt = args.prompt - num_inference_steps = args.num_inference_steps - engine_dir = f'sdxl_s{size}_w{world_size}' if args.engine_dir is None else args.engine_dir - num_warmup_runs = args.num_warmup_runs - avg_runs = args.avg_runs - output_file = args.output - - pipeline = StableDiffusionXLPipeline.from_pretrained( - model_dir, - torch_dtype=torch.float16, - use_safetensors=True, - ) - pipeline.set_progress_bar_config(disable=rank != 0) - pipeline.prepare(engine_dir, size) - pipeline.to('cuda') - - # warm up - for i in range(num_warmup_runs): - image = pipeline( - num_inference_steps=num_inference_steps, - prompt=prompt, - generator=torch.Generator(device="cuda").manual_seed(seed), - height=size, - width=size).images[0] - - latency_list = [] - for i in range(avg_runs): - st = time.time() - image = pipeline( - num_inference_steps=num_inference_steps, - prompt=prompt, - generator=torch.Generator(device="cuda").manual_seed(seed), - height=size, - width=size, - ).images[0] - ed = time.time() - latency_list.append(ed - st) - - latency_list = sorted(latency_list) - ignored_count = int(args.ignore_ratio * len(latency_list) / 2) - if ignored_count > 0: - latency_list = latency_list[ignored_count:-ignored_count] - - if rank == 0: - print(f"Avg latency: {np.sum(latency_list) / len(latency_list):.5f} s") - image.save(output_file) diff --git a/examples/models/contrib/skywork/README.md b/examples/models/contrib/skywork/README.md deleted file mode 100644 index 59b22a861987..000000000000 --- a/examples/models/contrib/skywork/README.md +++ /dev/null @@ -1,113 +0,0 @@ -# Skywork - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document elaborates how to build the [Skywork](https://huggingface.co/Skywork/) model to runnable engines on single GPU node and perform a summarization task using these engines. - -## Overview -The TensorRT LLM Skywork implementation is based on the LLaMA model. The implementation can -be found in [tensorrt_llm/models/llama/model.py](../../../../tensorrt_llm/models/llama/model.py). -The TensorRT LLM Skywork example code lies in [`examples/models/contrib/skywork`](./): - -* [`convert_checkpoint.py`](../../core/llama/convert_checkpoint.py) converts the Huggingface Model of Skywork into TensorRT LLM checkpoint. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 & BF16 - -## Usage - -This section gives a whole process where we convert HF models, build TensorRT LLM engines and ultimately perform summarization. - -### 1. Clone Code and Weights from Huggingface - -To download checkpoints from HF, you need to have `git-lfs` installed in your machine: - -```bash -pip install -r requirements.txt && sudo apt-get install git-lfs -``` - -Then clone the HF repository with: - -```bash -# Skywork 13B Base Model -git clone https://huggingface.co/Skywork/Skywork-13B-base -``` - -### 2. Convert HF Model to TRT Checkpoint - -```bash -cd examples/models/core/llama - -# fp16 model -python3 convert_checkpoint.py --model_dir ./Skywork-13B-base \ - --dtype float16 \ - --output_dir ./skywork-13b-base/trt_ckpt/fp16 - -# bf16 model -python3 convert_checkpoint.py --model_dir ./Skywork-13B-base \ - --dtype bfloat16 \ - --output_dir ./skywork-13b-base/trt_ckpt/bf16 -``` - -### 3. Build TensorRT Engine(s) - -```bash -# fp16 -trtllm-build --checkpoint_dir ./skywork-13b-base/trt_ckpt/fp16 \ - --gemm_plugin float16 \ - --gpt_attention_plugin float16 \ - --context_fmha enable \ - --max_batch_size 32 \ - --max_input_len 512 \ - --max_seq_len 1024 \ - --output_dir ./skywork-13b-base/trt_engine/fp16 - -# bf16 -trtllm-build --checkpoint_dir ./skywork-13b-base/trt_ckpt/bf16 \ - --gemm_plugin bfloat16 \ - --gpt_attention_plugin bfloat16 \ - --context_fmha enable \ - --max_batch_size 32 \ - --max_input_len 512 \ - --max_seq_len 1024 \ - --output_dir ./skywork-13b-base/trt_engine/bf16 -``` - -### 4. Summarization using the Engines - -After building TRT engines, we can use them to perform various tasks. TensorRT LLM provides handy code to run summarization on [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset and get [ROUGE](https://en.wikipedia.org/wiki/ROUGE_(metric)) scores. The `ROUGE-1` score can be used to validate model implementations. - -```bash -# fp16 -python ../../../summarize.py --hf_model_dir ./Skywork-13B-base \ - --test_hf \ - --batch_size 32 \ - --max_input_length 512 \ - --output_len 512 \ - --test_trt_llm \ - --engine_dir ./skywork-13b-base/trt_engine/fp16 \ - --data_type fp16 \ - --check_accuracy \ - --tensorrt_llm_rouge1_threshold=14 - -# bf16 -python ../../../summarize.py --hf_model_dir ./Skywork-13B-base \ - --test_hf \ - --batch_size 32 \ - --max_input_length 512 \ - --output_len 512 \ - --test_trt_llm \ - --engine_dir ./skywork-13b-base/trt_engine/bf16 \ - --data_type bf16 \ - --check_accuracy \ - --tensorrt_llm_rouge1_threshold=14 -``` diff --git a/examples/models/contrib/skywork/requirements.txt b/examples/models/contrib/skywork/requirements.txt deleted file mode 100644 index 88232baef811..000000000000 --- a/examples/models/contrib/skywork/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score -sentencepiece>=0.1.99 diff --git a/examples/models/contrib/smaug/README.md b/examples/models/contrib/smaug/README.md deleted file mode 100644 index 2d96da5854e1..000000000000 --- a/examples/models/contrib/smaug/README.md +++ /dev/null @@ -1,62 +0,0 @@ -# Smaug - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document elaborates how to build the [Smaug-72B-v0.1](https://huggingface.co/abacusai/Smaug-72B-v0.1) model to runnable engines on multi-GPU node and perform a summarization task using these engines. - -## Overview - -The TensorRT LLM support for Smaug-72B-v0.1 is based on the LLaMA model, the implementation can be found in [tensorrt_llm/models/llama/model.py](../../../../tensorrt_llm/models/llama/model.py). Smaug model resembles LLaMA very much except it uses bias term in its attention module, we therefore reuse the [LLaMA example code](../../core/llama) for Smaug, - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert the LLaMA model into TensorRT LLM checkpoint format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - -* FP16 - -## Usage - -This section gives a whole process where we convert HF models, build TensorRT LLM engines and ultimately perform summarization. - -### Build TensorRT engine(s) - -Run the following commands and TRT-LLM will first transforms a HF model into its own checkpoint format, then builds a TRT engine based on the checkpoint - -```bash -python ../../../llama/convert_checkpoint.py \ - --model_dir ./Smaug-72B-v0.1 \ - --output_dir ./tllm_checkpoint_8gpu_tp8 \ - --dtype float16 \ - --tp_size 8 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_8gpu_tp8 \ - --output_dir ./Smaug_72B_tp8 \ - --gemm_plugin float16 \ - --gpt_attention_plugin float16 \ - --context_fmha=enable \ - --max_batch_size 64 \ - --remove_input_padding=enable -``` - -### Run Summarization - -After building TRT engine, we can use it to perform various tasks. TensorRT LLM provides handy code to run summarization on [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset and get [ROUGE](https://en.wikipedia.org/wiki/ROUGE_(metric)) scores. The `ROUGE-1` score can be used to validate model implementations. - -```bash -mpirun -n 8 -allow-run-as-root python ../../../summarize.py \ - --hf_model_dir ../Smaug-72B-v0.1 \ - --engine_dir ./Smaug_72B_tp8 \ - --data_type fp16 \ - --test_hf \ - --hf_device_map_auto \ - --test_trt_llm -``` diff --git a/examples/models/contrib/smaug/requirements.txt b/examples/models/contrib/smaug/requirements.txt deleted file mode 100644 index 88232baef811..000000000000 --- a/examples/models/contrib/smaug/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score -sentencepiece>=0.1.99 diff --git a/examples/models/core/granite/README.md b/examples/models/core/granite/README.md deleted file mode 100644 index b6cf34c50c5f..000000000000 --- a/examples/models/core/granite/README.md +++ /dev/null @@ -1,89 +0,0 @@ -# Granite - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a [Granite 3.0](https://huggingface.co/collections/ibm-granite/granite-30-language-models-66fdb59bbb54785c3512114f) model in TensorRT-LLM. - -The TensorRT LLM Granite implementation is based on the LLaMA model, with Mixture of Experts (MoE) enabled. The implementation can be found in [`llama/model.py`](../../../../tensorrt_llm/models/llama/model.py). See the LLaMA example [`examples/models/core/llama`](../llama) for details. - -- [Granite 3.0](#Granite) - - [Download model checkpoints](#download-model-checkpoints) - - [Convert weights from HF Transformers to TensorRT LLM format](#Convert-weights-from-HF-Transformers-to-TensorRT-LLM-format) - - [Build TensorRT engine](#build-tensorrt-engine) - - [Run Engine](#run-engine) - -## Download model checkpoints - -First, download the HuggingFace BF16 checkpoints of Granite 3.0 model. - -```bash -HF_MODEL="granite-3.0-8b-instruct" # or granite-3.0-3b-a800m-instruct -# clone the model we want to build -git clone https://huggingface.co/ibm-granite/${HF_MODEL} tmp/hf_checkpoints/${HF_MODEL} -``` - -## Convert weights from HF Transformers to TensorRT LLM format -Set environment variables and necessary directory: - -```bash -PREC_RAW="bfloat16" -TP=1 -mkdir -p tmp/trt_engines -``` - -### BF16 -Convert the weights using the `convert_checkpoint.py` script: - -```bash -ENGINE="${HF_MODEL}_${PREC_RAW}_tp${TP}" -export TRTLLM_DISABLE_UNIFIED_CONVERTER=1 # The current checkpoint conversion code requires legacy path -python3 ../llama/convert_checkpoint.py --model_dir tmp/hf_checkpoints/${HF_MODEL} \ - --output_dir tmp/tllm_checkpoints/${ENGINE} \ - --dtype ${PREC_RAW} \ - --tp_size ${TP} \ - --use_embedding_sharing - - -``` -### FP8 PTQ -Notes: -- Currently quantize.py does not support Expert Parallelism (EP) mode yet. User should use `../llama/convert_checkpoint.py` and specify `--moe_ep_size 1` instead, if needed. -- TensorRT LLM uses static quantization methods, which is expected to be faster at runtime as compared to dynamic quantization methods. This comes at a cost of an offline calibration step during quantization. `batch_size` and `calib_size` can be adjusted to shorten the calibration time. Please refer to `../../../quantization/README.md` for explanation. - -```bash -PREC_QUANT="fp8" -ENGINE="${HF_MODEL}_${PREC_QUANT}_tp${TP}" -python ../../../quantization/quantize.py --model_dir tmp/hf_checkpoints/${HF_MODEL} \ - --dtype ${PREC_RAW} \ - --qformat ${PREC_QUANT} \ - --kv_cache_dtype ${PREC_QUANT} \ - --output_dir tmp/tllm_checkpoints/${ENGINE} \ - --batch_size 1 \ - --calib_size 128 \ - --tp_size ${TP} - -``` - -## Build TensorRT engine -```bash -# Enable fp8 context fmha to get further acceleration by setting `--use_fp8_context_fmha enable` -# Use --workers to enable parallel build -trtllm-build --checkpoint_dir ./tmp/tllm_checkpoints/${ENGINE} \ - --output_dir ./tmp/trt_engines/${ENGINE} \ - --gpt_attention_plugin ${PREC_RAW} \ - --gemm_plugin ${PREC_RAW} \ - --workers ${TP} -``` - -## Run Engine -Test your engine with the [run.py](../../../run.py) script: - -```bash -mpirun -n ${TP} --allow-run-as-root python ../../../run.py --engine_dir ./tmp/trt_engines/${ENGINE} --tokenizer_dir tmp/hf_checkpoints/${HF_MODEL} --max_output_len 20 --input_text "The future of AI is" -``` - -For more usage examples see [`examples/models/core/llama/README.md`](../llama/README.md) diff --git a/examples/models/core/mixtral/README.md b/examples/models/core/mixtral/README.md deleted file mode 100644 index 079753f88c92..000000000000 --- a/examples/models/core/mixtral/README.md +++ /dev/null @@ -1,219 +0,0 @@ -# Mixtral - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a Mixtral model in TensorRT LLM on both single GPU, single node multi-GPU and -multi-node multi-GPU. Mixtral 8x22B is also supported and can be replace Mixtral 8x7B below as long as GPU memory is -sufficient. - -## Overview - -The TensorRT LLM Mixtral implementation is based on the LLaMA model, with Mixture of Experts enabled. The implementation can -be found in [tensorrt_llm/models/llama/model.py](../../../../tensorrt_llm/models/llama/model.py). -See the LLaMA example [`examples/models/core/llama`](../llama) for details. - -### Build TensorRT engine(s) - -#### Download Mixtral 8x7b weights -Get the weights by downloading from HF https://huggingface.co/mistralai/Mixtral-8x7B-v0.1. -See also https://huggingface.co/docs/transformers/main/en/model_doc/mixtral - -```bash -git lfs install -git clone https://huggingface.co/mistralai/Mixtral-8x7B-v0.1 -``` - -#### Download Mixtral 8x22b weights -Get the weights by downloading from HF https://huggingface.co/mistralai/Mixtral-8x22B-v0.1. -See also https://huggingface.co/docs/transformers/main/en/model_doc/mixtral - -```bash -git lfs install -git clone https://huggingface.co/mistralai/Mixtral-8x22B-v0.1 -``` - -We use the LLaMA `convert_checkpoint.py` script to convert and build the model. TensorRT LLM LLaMA builds TensorRT engine(s) from HF checkpoint provided by `--model_dir`. -If no checkpoint directory is specified, TensorRT LLM will build engine(s) with dummy weights. - -`trtllm-build` uses one GPU by default, but if you have already more GPUs available at build time, -you may enable parallel builds to make the engine building process faster by adding the `--workers` argument. - -Here are some examples: - -```bash -# Build Mixtral8x7B with pipeline parallelism -python ../llama/convert_checkpoint.py --model_dir ./Mixtral-8x7B-v0.1 \ - --output_dir ./tllm_checkpoint_mixtral_2gpu \ - --dtype float16 \ - --pp_size 2 -trtllm-build --checkpoint_dir ./tllm_checkpoint_mixtral_2gpu \ - --output_dir ./trt_engines/mixtral/pp2 \ - --gemm_plugin float16 - -``` - -```bash -# Build Mixtral8x7B with tensor parallelism -python ../llama/convert_checkpoint.py --model_dir ./Mixtral-8x7B-v0.1 \ - --output_dir ./tllm_checkpoint_mixtral_2gpu \ - --dtype float16 \ - --tp_size 2 \ - --moe_tp_size 2 -trtllm-build --checkpoint_dir ./tllm_checkpoint_mixtral_2gpu \ - --output_dir ./trt_engines/mixtral/tp2 \ - --gemm_plugin float16 - - -# Build Mixtral8x22B with tensor parallelism and expert parallelism -python ../llama/convert_checkpoint.py --model_dir ./Mixtral-8x22B-v0.1 \ - --output_dir ./tllm_checkpoint_mixtral_8gpu \ - --dtype float16 \ - --tp_size 8 \ - --moe_tp_size 2 \ - --moe_ep_size 4 -trtllm-build --checkpoint_dir ./tllm_checkpoint_mixtral_8gpu \ - --output_dir ./trt_engines/mixtral/tp2ep4 \ - --gemm_plugin float16 -``` - -Then, you can test your engine with the [run.py](../../../run.py) script: - -```bash -mpirun -n 2 python3 ../../../run.py --engine_dir ./trt_engines/mixtral/tp2 --tokenizer_dir ./Mixtral-8x7B-v0.1 --max_output_len 8 --input_text "I love french quiche" -``` - -For more examples see [`examples/models/core/llama/README.md`](../llama/README.md) - -### Parallelism Modes - -Mixture of Experts supports 3 parallelism modes, these are Expert Parallelism (EP), Tensor Parallelism (TP), and the hybrid of the two (TP+EP). - -In TP mode (default) expert weight matrices are sliced evenly between all GPUs, so that all GPUs work together to calculate the result for each expert. - -In EP mode each GPU is assigned a subset of the expert weights matrices, so each GPU works independently to calculate the result for its assigned experts. This may cause load balancing issues where some GPUs have more work than others, thus increasing latency. - -In TP+EP mode, both strategies are used simultaneously. This means each GPU handles a portion of the expert weights matrices (as in EP mode) and these weights are further sliced across multiple GPUs (as in TP mode). This hybrid approach aims to balance the workload more evenly across GPUs, enhancing efficiency and reducing the likelihood of bottlenecks associated with EP mode alone. - -You can enable Expert Parallel or hybrid parallel by setting `--moe_tp_size` and `--moe_ep_size` when calling `convert_coneckpoint.py`. If only `--moe_tp_size` is provided, TRT-LLM will use Tensor Parallel for the MoE model; if only `--moe_ep_size` is provided, TRT-LLM will use Expert Parallel; if both are provided, the hybrid parallel will be used. - -Be sure that the product of `moe_tp_size` and `moe_ep_size` should equal to `tp_size`, since the total number of MoE parallelism across all GPUs must match the total number of parallelism in other parts of the model. - -```bash -# Build Mixtral8x7B with Expert Parallelism -python ../llama/convert_checkpoint.py --model_dir ./Mixtral-8x7B-v0.1 \ - --output_dir ./tllm_checkpoint_mixtral_2gpu \ - --dtype float16 \ - --tp_size 2 \ - --moe_ep_size 2 -trtllm-build --checkpoint_dir ./tllm_checkpoint_mixtral_2gpu \ - --output_dir ./trt_engines/mixtral/ep2 \ - --gemm_plugin float16 - -# Build Mixtral8x7B with Expert Parallelism and Tensor Parallelism -python ../llama/convert_checkpoint.py --model_dir ./Mixtral-8x7B-v0.1 \ - --output_dir ./tllm_checkpoint_mixtral_4gpu \ - --dtype float16 \ - --tp_size 4 \ - --moe_tp_size 2 \ - --moe_ep_size 2 -trtllm-build --checkpoint_dir ./tllm_checkpoint_mixtral_4gpu \ - --output_dir ./trt_engines/mixtral/tp2ep2 \ - --gemm_plugin float16 -``` - -### Normalization Modes - -MOE Supports different normalization modes which influence how the scales are calculated for the final weighted sum in -of the different top-k values. - -- 0 (NONE) corresponds to: `scales = topk(softmax(routing values))` -- 1 (RENORM) corresponds to: `scales = softmax(topk(routing values))` -- 2 (SPARSE_MIXER) corresponds to: `scales = sparsemixer(routing values)` - -Mixtral uses `RENORM` mode, this is set as the default. To use a different mode use the `--moe_normalization_mode` flag. -See [tensorrt_llm/layers/moe.py](../../../../tensorrt_llm/layers/moe.py#L56) for available values - - -## Quantization - -### Weight-only Quantization - -Mixtral supports weight only quantization - -```bash -# Build Mixtral8x7B with weight only -python ../llama/convert_checkpoint.py --model_dir ./Mixtral-8x7B-v0.1 \ - --output_dir ./tllm_checkpoint_mixtral_2gpu \ - --dtype float16 \ - --tp_size 2 \ - --use_weight_only \ - --weight_only_precision int8 -trtllm-build --checkpoint_dir ./tllm_checkpoint_mixtral_2gpu \ - --output_dir ./trt_engines/mixtral/tp2 \ - --gemm_plugin float16 -``` - -### FP8 Post-Training Quantization - -Mixtral supports FP8 quantization, using Modelopt. See [`examples/models/core/llama/README.md`](../llama/README.md#fp8-post-training-quantization) for full details on installing Modelopt - -```bash -# Quantize HF Mixtral into FP8 and export trtllm checkpoint -python ../../../quantization/quantize.py --model_dir ./Mixtral-8x7B-v0.1 \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir ./tllm_checkpoint_mixtral_2gpu \ - --calib_size 512 \ - --tp_size 2 - -# Build trtllm engines from the trtllm checkpoint -# Enable fp8 context fmha to get further acceleration by setting `--use_fp8_context_fmha enable` -trtllm-build --checkpoint_dir ./tllm_checkpoint_mixtral_2gpu \ - --output_dir ./engine_outputs \ - --workers 2 -``` - -### AWQ Quantization - -Mixtral supports AWQ quantization using [AutoAWQ](https://github.com/casper-hansen/AutoAWQ). - -```bash -# Convert AutoAWQ HF checkpoints into TRT-LLM checkpoint -python ../llama/convert_checkpoint.py --model_dir ./tmp/mixtral-8x7b-v0.1-AWQ/ \ - --output_dir ./tllm_checkpoint_mixtral_awq_1gpu - -# Build trtllm engines from the trtllm checkpoint -trtllm-build --checkpoint_dir ./tllm_checkpoint_mixtral_awq_1gpu \ - --output_dir ./engine_outputs -``` - -You may found `quant_algo = W4A16_GPTQ` in the configuration file of the converted checkpoints, and that's because AutoAWQ is using exactly the same components as GPTQ. - -### NVFP4 Post-Training Quantization - -Mixtral supports NVFP4 quantization. - -```bash -# Quantize HF Mixtral into FP8 and export trtllm checkpoint -python ../../../quantization/quantize.py --model_dir ./Mixtral-8x7B-v0.1 \ - --dtype float16 \ - --qformat nvfp4 \ - --kv_cache_dtype fp8 \ - --output_dir ./tllm_checkpoint_mixtral_nvfp4_1gpu \ - --calib_size 512 \ - --tp_size 1 - -# Build trtllm engines from the trtllm checkpoint -# Enable fp8 context fmha to get further acceleration by setting `--use_fp8_context_fmha enable` -trtllm-build --checkpoint_dir ./tllm_checkpoint_mixtral_nvfp4_1gpu \ - --output_dir ./engine_outputs -``` - -## OOTB - -Mixtral supports OOTB operation without the plugin, however this comes at a significant performance cost. Users should prefer using the plugin path whenever possible diff --git a/examples/models/core/mixtral/requirements.txt b/examples/models/core/mixtral/requirements.txt deleted file mode 100644 index c1b98e5d3f7a..000000000000 --- a/examples/models/core/mixtral/requirements.txt +++ /dev/null @@ -1,4 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -transformers==4.56.0 -accelerate==0.25.0 diff --git a/examples/models/core/multimodal/README.md b/examples/models/core/multimodal/README.md index 159e26d4c7d8..924561126cfb 100644 --- a/examples/models/core/multimodal/README.md +++ b/examples/models/core/multimodal/README.md @@ -1,1169 +1,17 @@ # Multi-Modal -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. +The engine-build multimodal workflow that used to live here +(`build_multimodal_engine.py` / `run.py` / `eval.py` on top of +`trtllm-build`) was removed together with the legacy TensorRT backend. -This document shows how to run multimodal pipelines with TensorRT-LLM, e.g. from image+text input modalities to text output. +Multimodal models are supported on the PyTorch backend. See: -Multimodal models' LLM part has an additional parameter `--max_multimodal_len` compared to LLM-only build commands. Under the hood, `max_multimodal_len` and `max_prompt_embedding_table_size` are effectively the same concept, i.e., prepended/concatenated embeddings (either multimodal feature embeddings or prompt tuning embeddings) to the LLM input embeddings. The multimodal features from the visual encoder of shape `[batch_size, num_visual_features, visual_hidden_dim]` is flattened as `[batch_size * num_visual_features, visual_hidden_dim]` and passed like a prompt embedding table. - -We first describe three runtime modes for running multimodal models and how to run each model on a single GPU. We then provide general guidelines on using tensor parallelism for the LLM part of the pipeline. - -- [Runtime Mode](#runtime-modes) -- [BLIP2](#blip2) -- [CogVLM](#cogvlm) -- [Deplot](#deplot) -- [Fuyu](#fuyu) -- [Gemma3](#gemma3) -- [InternLM-XComposer2](#internlm-xcomposer2) -- [InternVL2](#internvl2) -- [Kosmos-2](#kosmos-2) -- [LLaVA, LLaVa-NeXT, LLaVA-OneVision and VILA](#llava-llava-next-llava-onevision-and-vila) -- [MLLaMA](#mllama) -- [NeVA](#neva) -- [Nougat](#nougat) -- [Phi-3-vision](#phi-3-vision) -- [Phi-4-multimodal](#phi-4-multimodal) -- [Qwen2-VL](#qwen2-vl) -- [Qwen-Image-Bench Evaluator](#qwen-image-bench-evaluator) -- [Video NeVA](#video-neva) -- [Dataset Evaluation](#dataset-evaluation) -- [Enabling Tensor Parallelism for multi-GPU](#enabling-tensor-parallelism-for-multi-gpu) -- [Enabling Embedding Table Offloading](#enabling-embedding-table-offloading) - -## Runtime Modes -TensorRT LLM supports three runtime modes for running multimodal models. -- `cpp_llm_only` (default): vision engine runs in python runtime, LLM in pybind C++ runtime -- `python`: everything runs in python runtime -- `cpp`: everything runs in C++ runtime - -This can be specified by the `--session RUNTIME_MODE` argument in `run.py` (see instructions of each model below). -Not all models supports end-to-end `cpp` mode, the checked ones below are supported. See footnotes for reasons models are unsupported -- [ ] BLIP-2-T5 [^1] -- [x] BLIP-2-OPT -- [x] CogVLM -- [ ] Deplot [^1] -- [ ] Pix2Struct [^1] -- [x] Fuyu -- [x] InternVL2-2b -- [x] Kosmos-2 -- [x] LLaVA -- [ ] LLaVA-NeXT / OneVision [^2] -- [x] VILA [^3] -- [ ] Mllama [^1] -- [x] NeVA -- [ ] Nougat [^1] -- [ ] Phi-3-Vision [^2] -- [ ] Phi-4-multimodal -- [ ] Qwen2-VL [^4] -- [x] Video-NeVA - - -[^1]: Model uses cross attention to feed visiual features to LLM decoder, which is not supported -[^2]: Model requires post processing its encoder output features, which is not supported -[^3]: Currently C++ runtime only supports single image per request (VILA mode 2) -[^4]: Vision encoder requires additional inputs not supported by the C++ runtime - -## BLIP2 - -This BLIP section covers both BLIP2-OPT and BLIP2-T5, with minor changes needed when switching the LLM backbone. - -1. Download Huggingface weights and convert original checkpoint to TRT-LLM checkpoint format - following example in `examples/models/contrib/opt/README.md` and `examples/models/core/enc_dec/README.md`. - - ```bash - export MODEL_NAME="blip2-opt-2.7b" # options: blip2-opt-6.7b, blip2-flan-t5-xl, blip2-flan-t5-xxl - git clone https://huggingface.co/Salesforce/${MODEL_NAME} tmp/hf_models/${MODEL_NAME} - ``` - - For BLIP2-OPT family, - ```bash - python ../../contrib/opt/convert_checkpoint.py --model_type blip2 \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --output_dir tmp/trt_models/${MODEL_NAME}/fp16/1-gpu \ - --dtype float16 - ``` - - For BLIP2-T5 family, - ```bash - python ../enc_dec/convert_checkpoint.py --model_type blip2 \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --output_dir tmp/trt_models/${MODEL_NAME}/bfloat16 \ - --tp_size 1 \ - --pp_size 1 \ - --dtype bfloat16 - ``` - -2. Build TRT-LLM engine from TRT-LLM checkpoint - - For BLIP2-OPT family, - ```bash - trtllm-build \ - --checkpoint_dir tmp/trt_models/${MODEL_NAME}/fp16/1-gpu \ - --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/llm \ - --gemm_plugin float16 \ - --max_beam_width 1 \ - --max_batch_size 8 \ - --max_seq_len 1024 \ - --max_input_len 924 \ - --max_multimodal_len 256 # 8 (max_batch_size) * 32 (num_visual_features) - ``` - - For BLIP2-T5 family, - ```bash - trtllm-build --checkpoint_dir tmp/trt_models/${MODEL_NAME}/bfloat16/encoder \ - --output_dir tmp/trt_engines/${MODEL_NAME}/bfloat16/llm/encoder \ - --paged_kv_cache disable \ - --moe_plugin disable \ - --gemm_plugin bfloat16 \ - --bert_attention_plugin bfloat16 \ - --gpt_attention_plugin bfloat16 \ - --remove_input_padding enable \ - --context_fmha disable \ - --max_beam_width 1 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_multimodal_len 256 # 8 (max_batch_size) * 32 (num_visual_features) - - trtllm-build --checkpoint_dir tmp/trt_models/${MODEL_NAME}/bfloat16/decoder \ - --output_dir tmp/trt_engines/${MODEL_NAME}/bfloat16/llm/decoder \ - --paged_kv_cache disable \ - --moe_plugin disable \ - --gemm_plugin bfloat16 \ - --bert_attention_plugin bfloat16 \ - --gpt_attention_plugin bfloat16 \ - --remove_input_padding enable \ - --context_fmha disable \ - --max_beam_width 1 \ - --max_batch_size 8 \ - --max_seq_len 1024 \ - --max_encoder_input_len 924 \ - --max_input_len 1 # Same command for decoder but don't set --max_multimodal_len - ``` - - **NOTE**: `max_multimodal_len = max_batch_size * num_visual_features`, so if you change max_batch_size, max multimodal length **MUST** be changed accordingly. - -3. Build TensorRT engines for vision encoders - - ```bash - python build_multimodal_engine.py --model_type blip2 --model_path tmp/hf_models/${MODEL_NAME} --output_dir tmp/trt_engines/${MODEL_NAME}/bfloat16/vision --max_batch_size 8 - ``` - - The built engines are located in `tmp/trt_engines/${MODEL_NAME}/bfloat16/vision` for BLIP2-T5, similarly for BLIP-OPT. - - To run the BLIP2 pipeline with batch size > 1, change `--max_batch_size` argument to `build_multimodal_engine.py` accordingly. - -4. Assemble everything into BLIP2 pipeline - - For BLIP2-OPT family, - ```bash - python run.py \ - --max_new_tokens 30 \ - --input_text "Question: which city is this? Answer:" \ - --hf_model_dir tmp/hf_models/${MODEL_NAME} \ - --engine_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu - ``` - - For BLIP2-T5 family, - ```bash - python run.py \ - --max_new_tokens 30 \ - --input_text "Question: which city is this? Answer:" \ - --hf_model_dir tmp/hf_models/${MODEL_NAME} \ - --engine_dir tmp/trt_engines/${MODEL_NAME}/bfloat16 - ``` - -5. (Optional) INT8/INT4 weight-only quantization for OPT can be enabled using commands as follows (take `INT4` as an example, while `INT8` is the default precision for weight-only quantization): - ```bash - python ../../contrib/opt/convert_checkpoint.py \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --dtype float16 \ - --output_dir tmp/trt_models/${MODEL_NAME}/int4_weightonly/1-gpu \ - --use_weight_only \ - --weight_only_precision int4 - - trtllm-build \ - --checkpoint_dir tmp/trt_models/${MODEL_NAME}/int4_weightonly/1-gpu \ - --output_dir tmp/trt_engines/${MODEL_NAME}/int4_weightonly/1-gpu/llm \ - --gemm_plugin float16 \ - --max_beam_width 1 \ - --max_batch_size 8 \ - --max_multimodal_len 256 \ - --max_input_len 924 \ - --max_seq_len 1024 - ``` - - The built OPT engines lie in `tmp/trt_engines/${MODEL_NAME}/int4_weightonly/1-gpu/llm`. - You should use this directory without the `llm` part as `--engine_dir` argument to `run.py` - - **NOTE:** INT8/INT4 option is not supported for BLIP2-T5, because quantization support has not been - added for encoder-decoder models yet. - -## CogVLM - -Currently, CogVLM only support bfloat16 precision. - -1. Download Huggingface weights - - ```bash - export MODEL_NAME="cogvlm-chat-hf" - git clone https://huggingface.co/THUDM/${MODEL_NAME} tmp/hf_models/${MODEL_NAME} - export TOKENIZER_NAME="vicuna-7b-v1.5" - git clone https://huggingface.co/lmsys/${TOKENIZER_NAME} tmp/hf_models/${TOKENIZER_NAME} - ``` - - Because currently onnx doesn't support `xops.memory_efficient_attention`, we need to modify some source code of the huggingface CogVLM. - ``` - cd tmp/hf_models/${MODEL_NAME} - sed -i '4s/.*//;40s/.*/ out = self.attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).transpose(1, 2).contiguous()/;41s/.*//;42s/.*//' visual.py # It will replace memory_efficient_attention with some basic ops - ``` - -2. Convert Huggingface weights into TRT-LLM checkpoints and build TRT engines using scripts in `examples/cogvlm` - - CogVLM uses a Vit encoder as LLM encoder and a modified Llama as decoder. - - ```bash - python ../../contrib/cogvlm/convert_checkpoint.py --model_dir tmp/hf_models/${MODEL_NAME} --output_dir tmp/trt_models/${MODEL_NAME} --dtype bfloat16 --use_prompt_tuning - - trtllm-build --checkpoint_dir tmp/trt_models/${MODEL_NAME} \ - --output_dir tmp/trt_engines/${MODEL_NAME}/bf16/1-gpu/llm \ - --gemm_plugin bfloat16 \ - --gpt_attention_plugin bfloat16 \ - --remove_input_padding enable \ - --max_batch_size 48 \ - --max_input_len 2048 \ - --max_seq_len 3076 \ - --paged_kv_cache enable \ - --bert_attention_plugin disable \ - --moe_plugin disable \ - --max_multimodal_len 61440 # 48 (max_batch_size) * 1280 (max_num_visual_features) - ``` - -3. Generate TensorRT engines for visual components and combine everything into final pipeline. - - ```bash - python build_multimodal_engine.py --model_type cogvlm --model_path tmp/hf_models/${MODEL_NAME} --max_batch_size 48 --output_dir tmp/trt_engines/${MODEL_NAME}/bf16/1-gpu/vision - - python run.py \ - --max_new_tokens 1000 \ - --input_text " [INST] please describe this image in detail [/INST] " \ - --hf_model_dir tmp/hf_models/${TOKENIZER_NAME} \ - --engine_dir tmp/trt_engines/${MODEL_NAME}/bf16/1-gpu \ - --batch_size 1 \ - --top_p 0.4 \ - --top_k 1 \ - --temperature 0.2 \ - --repetition_penalty 1.2 \ - --enable_context_fmha_fp32_acc - - CogVLM uses model_runner_cpp for its LLM decoder by default. To switch to model_runner, set `--session python` in the command mentioned above. - ``` - -## Deplot - -1. Download Huggingface weights and convert original checkpoint to TRT-LLM checkpoint format - following example in `examples/models/core/enc_dec/README.md`. - - ```bash - export MODEL_NAME="deplot" - git clone https://huggingface.co/google/${MODEL_NAME} tmp/hf_models/${MODEL_NAME} - - python ../enc_dec/convert_checkpoint.py --model_type pix2struct \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --output_dir tmp/trt_models/${MODEL_NAME}/float16 \ - --tp_size 1 \ - --pp_size 1 \ - --dtype float16 - ``` - -2. Build TRT-LLM engine from TRT-LLM checkpoint - - ```bash - trtllm-build --checkpoint_dir tmp/trt_models/${MODEL_NAME}/float16/decoder \ - --output_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/float16/llm/decoder \ - --paged_kv_cache disable \ - --moe_plugin disable \ - --gemm_plugin float16 \ - --bert_attention_plugin float16 \ - --gpt_attention_plugin float16 \ - --remove_input_padding enable \ - --context_fmha disable \ - --max_beam_width 1 \ - --max_batch_size 8 \ - --max_seq_len 2558 \ - --max_encoder_input_len 2048 \ - --max_input_len 1 - ``` - - The built deplot engines are located in `tmp/trt_engines/${MODEL_NAME}/1-gpu/float16`. - -3. Build TensorRT engines for visual components - - ```bash - python build_multimodal_engine.py --model_type pix2struct --model_path tmp/hf_models/${MODEL_NAME} --max_batch_size 8 --output_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/float16/vision - ``` - - The built visual engines are located in `tmp/trt_engines/${MODEL_NAME}/1-gpu/float16/vision`. - - To run the deplot pipeline with batch size > 1, change `--max_batch_size` argument to `build_multimodal_engine.py` accordingly. - -4. Assemble everything into deplot pipeline - - ```bash - python run.py \ - --max_new_tokens 100 \ - --input_text "" \ - --hf_model_dir tmp/hf_models/${MODEL_NAME} \ - --engine_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/float16 - ``` - -## Fuyu - -1. Download Huggingface weights - - ```bash - export MODEL_NAME="fuyu-8b" - git clone https://huggingface.co/adept/${MODEL_NAME} tmp/hf_models/${MODEL_NAME} - ``` - -2. Convert Huggingface weights into TRT-LLM checkpoints and build TRT engines using scripts in `examples/models/core/gpt`. - The LLM portion of Fuyu uses a Persimmon model - ```bash - python ../gpt/convert_checkpoint.py \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --output_dir tmp/trt_models/${MODEL_NAME}/fp16/1-gpu \ - --dtype float16 \ - --gpt_variant persimmon - - trtllm-build \ - --checkpoint_dir tmp/trt_models/${MODEL_NAME}/fp16/1-gpu \ - --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/llm \ - --gemm_plugin float16 \ - --use_fused_mlp=enable \ - --max_batch_size 1 \ - --max_input_len 2048 \ - --max_seq_len 2560 \ - --max_multimodal_len 2048 - ``` - -3. Generate TensorRT engines for visual components and combine everything into final pipeline. - - ```bash - python build_multimodal_engine.py --model_type fuyu --model_path tmp/hf_models/${MODEL_NAME} --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/vision - - python run.py \ - --hf_model_dir tmp/hf_models/${MODEL_NAME} \ - --engine_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu - ``` - -## Gemma3 - -**NOTE: We only support Gemma3 VLMs in Pytorch workflow.** - -Gemma3VL decoder requires a custom attention mask while processing images. During the context phase: -- Text tokens attend to other tokens in a causal fashion (standard autoregressive behavior) -- Image tokens attend to other tokens in a causal fashion AND attend to other tokens from the same image in a bidirectional manner - -**Reference:** [Gemma3 Model Documentation](https://huggingface.co/docs/transformers/en/model_doc/gemma3) - -We support this custom mask with FlashInfer attention backend. - -### Requirements - -To ensure expected behavior with Gemma3VL, the following configurations are **required**: -- **Attention Backend**: Use the FlashInfer attention backend -- **Chunked Prefill**: Must be disabled -- **KV Cache Reuse**: Must be disabled - -### Quick Start - -#### 1. Download Model Weights - -```bash -export MODEL_NAME="gemma-3-27b-it" -git clone https://huggingface.co/google/${MODEL_NAME} -``` - -#### 2. Interactive Testing - -Use the `quickstart_multimodal.py` script for quick testing: - -```bash -python3 examples/llm-api/quickstart_multimodal.py \ - --model_dir ${MODEL_NAME}/ \ - --modality image \ - --image_format pil \ - --attention_backend FLASHINFER \ - --disable_kv_cache_reuse -``` - -#### 3. Model Serving - -Serve the model using `trtllm-serve` with the required llmapi arguments mentioned in a yaml file: - -```bash -# Create the configuration file -cat > extra-llm-api-options.yaml << 'EOF' -cuda_graph_config: null -attn_backend: "FLASHINFER" -enable_chunked_prefill: false -kv_cache_config: - enable_block_reuse: false -EOF - -# Serve the model -trtllm-serve ${MODEL_NAME}/ \ - --backend pytorch \ - --tp_size 1 \ - --port 8000 \ - --max_batch_size 4 \ - --config extra-llm-api-options.yaml -``` - -### Supported Model Variants - -Currently supported Gemma3 variants: 4B, 12B, 27B - - -## InternLM-XComposer2 - -**NOTE: We only support InternLM-XComposer-VL-7b for now** - -Firstly, please install transformers with 4.45.2 -```bash - pip install -r requirements-internlm-xcomposer2.txt -``` - -1. Convert Huggingface weights to TRT-LLM checkpoint format using `examples/models/contrib/internlm/README.md`. - -2. Use `trtllm-build` command to build TRT-LLM engine for OPT. - -3. The full list of commands is as follows: - - ```bash - export MODEL_NAME=internlm-xcomposer2-vl-7b - git lfs clone https://huggingface.co/internlm/${MODEL_NAME} tmp/hf_models/${MODEL_NAME} - - python ../internlm2/convert_checkpoint.py \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --dtype float16 \ - --output_dir tmp/trt_models/${MODEL_NAME}/fp16/1-gpu - - trtllm-build \ - --checkpoint_dir tmp/trt_models/${MODEL_NAME}/fp16/1-gpu \ - --output_dir trt_engines/${MODEL_NAME}/fp16/1-gpu/llm \ - --gemm_plugin float16 \ - --lora_plugin float16 \ - --lora_dir . \ - --max_lora_rank 256 \ - --max_input_len 1536 \ - --max_batch_size 48 \ - --max_multimodal_len 58800 # 58800 = 1225(visual token/img) * 48 (max_batch_size), as each image corresponds to 1225 visual tokens in the ViT here - - python build_multimodal_engine.py \ - --model_type internlm-xcomposer2 \ - --model_path tmp/hf_models/${MODEL_NAME} \ - --output_dir trt_engines/${MODEL_NAME}/fp16/1-gpu/vision \ - --max_batch_size 48 - - python run.py \ - --max_new_tokens 200 \ - --hf_model_dir tmp/hf_models/${MODEL_NAME} \ - --engine_dir trt_engines/${MODEL_NAME}/fp16/1-gpu \ - --batch_size 1 - ``` - -## InternVL2 - -[InternVL Family](https://github.com/OpenGVLab/InternVL): Closing the Gap to Commercial Multimodal Models with Open-Source Suites —— A Pioneering Open-Source Alternative to GPT-4o. Here we show how to deploy InternVL2‑1B/InternVL2‑2B/InternVL2‑4B/InternVL2‑8B/InternVL2‑26B in TensorRT-LLM. - -Firstly, please install transformers with 4.37.2 -```bash - pip install transformers==4.37.2 -``` - -1. Download Huggingface weights - - For InternVL2-1B - ```bash - export MODEL_NAME="InternVL2-1B" - git clone https://huggingface.co/OpenGVLab/${MODEL_NAME} tmp/hf_models/${MODEL_NAME} - export LLM_MODEL_NAME="qwen" - ``` - - - For InternVL2-2B/InternVL2‑8B/InternVL2‑26B - ```bash - export MODEL_NAME="InternVL2-2B" # or InternVL2‑8B, InternVL2‑26B - git clone https://huggingface.co/OpenGVLab/${MODEL_NAME} tmp/hf_models/${MODEL_NAME} - export LLM_MODEL_NAME="internlm2" - ``` - - - For InternVL2-4B - ```bash - export MODEL_NAME="InternVL2-4B" - git clone https://huggingface.co/OpenGVLab/${MODEL_NAME} tmp/hf_models/${MODEL_NAME} - export LLM_MODEL_NAME="phi" - ``` - -2. Convert Huggingface weights into TRT-LLM checkpoints - ```bash - python ../${LLM_MODEL_NAME}/convert_checkpoint.py \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --output_dir tmp/trt_models/${MODEL_NAME}/fp16/1-gpu \ - --dtype float16 - ``` - -3. Build TRT engines - ```bash - trtllm-build \ - --checkpoint_dir tmp/trt_models/${MODEL_NAME}/fp16/1-gpu \ - --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/llm \ - --gemm_plugin auto \ - --max_batch_size 1 \ - --max_input_len 4096 \ - --max_seq_len 4608 \ - --max_multimodal_len 3328 - ``` - -4. Generate TensorRT engines for visual components and combine everything into final pipeline. - ```bash - python build_multimodal_engine.py --model_type internvl --model_path tmp/hf_models/${MODEL_NAME} --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/vision - python run.py \ - --hf_model_dir tmp/hf_models/${MODEL_NAME} \ - --engine_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/ \ - --image_path tmp/hf_models/${MODEL_NAME}/examples/image1.jpg - ``` - -5. (Optional) FP8 and INT8 SmoothQuant quantization is supported for the InternVL2-4B variant (LLM model only). - - ```bash - # FP8 quantization - python ../../../quantization/quantize.py \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --output_dir tmp/trt_models/${MODEL_NAME}/fp8/1-gpu \ - --dtype bfloat16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 - - # INT8 SmoothQuant quantization - python ../../../quantization/quantize.py \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --output_dir tmp/trt_models/${MODEL_NAME}/int8/1-gpu \ - --dtype bfloat16 \ - --qformat int8_sq - ``` - - Then follow the same `trtllm-build`, `build_multimodal_engine.py` and `run.py` steps as before. - - -## Kosmos-2 - -1. Download Huggingface weights - - ```bash - export MODEL_NAME="kosmos-2" - git clone https://huggingface.co/microsoft/kosmos-2-patch14-224 tmp/hf_models/${MODEL_NAME} - ``` - -2. Convert Huggingface weights into TRT-LLM checkpoints and build TRT engines using scripts in `examples/models/core/gpt`. - ```bash - python ../gpt/convert_checkpoint.py \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --output_dir tmp/trt_models/${MODEL_NAME}/fp16/1-gpu \ - --dtype float16 \ - --gpt_variant ${MODEL_NAME} - - trtllm-build \ - --checkpoint_dir tmp/trt_models/${MODEL_NAME}/fp16/1-gpu \ - --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/llm \ - --gpt_attention_plugin float16 \ - --gemm_plugin float16 \ - --max_batch_size 1 \ - --max_input_len 512 \ - --max_seq_len 1024 \ - --max_multimodal_len 64 # 1 (max_batch_size) * 64 (num_visual_features) - ``` - -3. Generate TensorRT engines for visual components and combine everything into final pipeline. - - ```bash - python build_multimodal_engine.py --model_type kosmos-2 --model_path tmp/hf_models/${MODEL_NAME} --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/vision - - python run.py \ - --hf_model_dir tmp/hf_models/${MODEL_NAME} \ - --engine_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu - ``` - -## LLaVA, LLaVa-NeXT, LLaVA-OneVision and VILA - -[LLaVA](https://github.com/haotian-liu/LLaVA) and [VILA](https://github.com/Efficient-Large-Model/VILA) are both visual language models (VLM) that can be deployed in TensorRT LLM with many quantization options. [LLaVA-NeXT](https://huggingface.co/collections/llava-hf/llava-next-65f75c4afac77fd37dbbe6cf) is an extension of LLaVA. TRT-LLM currently supports [Mistral-7b](https://huggingface.co/llava-hf/llava-v1.6-mistral-7b-hf) and [ Nous-Hermes-2-Yi-34B](https://huggingface.co/llava-hf/llava-v1.6-34b-hf) variant of LLaVA-NeXT. [LLaVA-OneVision](https://huggingface.co/collections/llava-hf/llava-onevision-66bb1e9ce8856e210a7ed1fe) is another extension of LLaVA. - -1. Download Huggingface model weights. These models have both visual and LLM components - unlike BLIP2 example which downloads only LLM components from Huggingface. - - For LLaVA, - - ```bash - export MODEL_NAME="llava-1.5-7b-hf" # also llava-1.5-13b-hf - git clone https://huggingface.co/llava-hf/${MODEL_NAME} tmp/hf_models/${MODEL_NAME} - ``` - For LLaVA-NeXT, - - ```bash - export MODEL_NAME="llava-v1.6-mistral-7b-hf" #for 34b variant "llava-v1.6-34b-hf" - git clone https://huggingface.co/llava-hf/${MODEL_NAME} tmp/hf_models/${MODEL_NAME} - ``` - - For LLaVA-OneVision, - - ```bash - export MODEL_NAME="llava-onevision-qwen2-7b-ov-hf" # also llava-onevision-qwen2-0.5b-ov-hf, llava-onevision-qwen2-72b-ov-hf, etc - git clone https://huggingface.co/llava-hf/${MODEL_NAME} tmp/hf_models/${MODEL_NAME} - ``` - - For VILA, we need a few more steps until it is added to HF model zoo - - ```bash - # install the following dependency - pip install -r requirements-vila.txt - - # clone original VILA repo - export VILA_PATH="tmp/hf_models/VILA" - git clone https://github.com/Efficient-Large-Model/VILA.git ${VILA_PATH} - - # download VILA checkpoints - export MODEL_NAME="vila1.5-3b" # NOTE: name must contain vila or VILA! it's used to identify whether we need to register the non-HF VILA codebase in HF Auto class - git clone https://huggingface.co/Efficient-Large-Model/${MODEL_NAME} tmp/hf_models/${MODEL_NAME} - ``` - -2. Generate TRT-LLM engine for LLaMA following example in `examples/models/core/llama/README.md` and `examples/models/core/qwen/README.md` - - ```bash - python ../llama/convert_checkpoint.py \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --output_dir tmp/trt_models/${MODEL_NAME}/fp16/1-gpu \ - --dtype float16 - - # for LLaVA - trtllm-build \ - --checkpoint_dir tmp/trt_models/${MODEL_NAME}/fp16/1-gpu \ - --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/llm \ - --gemm_plugin float16 \ - --use_fused_mlp=enable \ - --max_batch_size 1 \ - --max_input_len 2048 \ - --max_seq_len 2560 \ - --max_multimodal_len 576 # 1 (max_batch_size) * 576 (num_visual_features) - - # for LLaVA-NeXT - trtllm-build \ - --checkpoint_dir tmp/trt_models/${MODEL_NAME}/fp16/1-gpu \ - --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/llm \ - --gpt_attention_plugin float16 \ - --gemm_plugin float16 \ - --use_fused_mlp=enable \ - --max_batch_size 1 \ - --max_input_len 4096 \ - --max_seq_len 5120 \ - --max_num_tokens 4096 \ - --max_multimodal_len 4096 # 1 (max_batch_size) * 4096 (max_input_len) - - # for VILA - trtllm-build \ - --checkpoint_dir tmp/trt_models/${MODEL_NAME}/fp16/1-gpu \ - --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/llm \ - --gemm_plugin float16 \ - --use_fused_mlp=enable \ - --max_batch_size 1 \ - --max_input_len 2048 \ - --max_seq_len 2560 \ - --max_multimodal_len 196 # 1 (max_batch_size) * 196 (num_visual_features) - - # for LLaVA-OneVision - python ../qwen/convert_checkpoint.py \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --output_dir tmp/trt_models/${MODEL_NAME}/fp16/1-gpu \ - --dtype float16 - - trtllm-build \ - --checkpoint_dir tmp/trt_models/${MODEL_NAME}/fp16/1-gpu \ - --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/llm \ - --gemm_plugin float16 \ - --use_fused_mlp=enable \ - --max_batch_size 1 \ - --max_input_len 7228 \ - --max_seq_len 7328 \ - --max_multimodal_len 7128 # max_batch_size * num_visual_features(depends on the image size or the specified video num frame) - ``` - -3. Build TensorRT engines for visual components - - ```bash - python build_multimodal_engine.py --model_path tmp/hf_models/${MODEL_NAME} --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/vision --model_type llava # for LLaVA - - python build_multimodal_engine.py --model_path tmp/hf_models/${MODEL_NAME} --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/vision --model_type llava_next --max_batch_size 5 # 1 (max_batch_size) * 5 (because LLAVA-NeXT visual encoder can have at most 5 patches) # for LLaVA-NeXT - - python build_multimodal_engine.py --model_path tmp/hf_models/${MODEL_NAME} --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/vision --model_type llava_onevision --max_batch_size 32 # max_batch_size * patch for image or frame for video # for LLaVA-OneVision - - python build_multimodal_engine.py --model_path tmp/hf_models/${MODEL_NAME} --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/vision --model_type vila --vila_path ${VILA_PATH} # for VILA - ``` - - ```bash - python run.py \ - --max_new_tokens 30 \ - --hf_model_dir tmp/hf_models/${MODEL_NAME} \ - --engine_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu \ - --input_text "\n Which city is this?" # for LLaVA and for LLaVA-NeXT - - python run.py \ - --max_new_tokens 30 \ - --hf_model_dir tmp/hf_models/${MODEL_NAME} \ - --engine_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu \ - --image_path=https://github.com/Efficient-Large-Model/VILA/raw/main/demo_images/av.png,https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png \ - --input_text "\n Please elaborate what you see in the images?","\n Which city is this?" \ - --batch_size=2 # for LLaVA - ``` - - Note that Llava can support N pairs inference batching, `--batch_size=N` should be used. There should be N images listed under `--image_path` and N text prompts listed under `--input_text`. Don't forget to set the `--max_batch_size` and `--max_multimodal_len` during engine building. - - For LLaVA-OneVision, you can use either image or video as inputs. - ```bash - python run.py \ - --max_new_tokens 30 \ - --hf_model_dir tmp/hf_models/${MODEL_NAME} \ - --engine_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu \ - --input_text "What is shown in this image?" \ - --image_path image.png - - python run.py \ - --max_new_tokens 30 \ - --hf_model_dir tmp/hf_models/${MODEL_NAME} \ - --engine_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu \ - --input_text "Why is this video funny?" \ - --video_path video.mp4 - --video_num_frames 8 # sample uniformly 8 frames from the video, up to 32 frames - ``` - - For VILA, you can use either local file or web url as input images. - Suppose you have a local image `av.png` downloaded from `https://github.com/Efficient-Large-Model/VILA/blob/main/demo_trt_llm/av.png` and the url of `merlion.png` - ```bash - wget -O av.png https://raw.githubusercontent.com/Efficient-Large-Model/VILA/main/demo_images/av.png - - python run.py \ - --max_new_tokens 30 \ - --hf_model_dir tmp/hf_models/${MODEL_NAME} \ - --engine_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu \ - --image_path=av.png,https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png \ - --input_text="\n\n Please elaborate what you see in the images?" \ - --batch_size=1 # for VILA mode 1 - - python run.py \ - --max_new_tokens 30 \ - --hf_model_dir tmp/hf_models/${MODEL_NAME} \ - --engine_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu \ - --image_path=av.png,https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png \ - --input_text="\n Please elaborate what you see in the images?","\n Which city is this?" \ - --batch_size=2 \ - --check_accuracy # for VILA mode 2 - ``` - - Note that VILA can support different modes in terms of batching: - - Mode 1: if you want to query N images as a whole using a prompt, `--batch_size=1` should be used (which is the default value). Example is given above. - - Mode 2: if you want to query N pairs, `--batch_size=N` should be used. There should be N images listed under `--image_path` and N text prompts listed under `--input_text`. Don't forget to set the `--max_batch_size` and `--max_multimodal_len` during engine building. - - Note: use `--run_profiling` for performance measurement, use `--check_accuracy` for accuracy check. - -4. (Optional) Different quantization methods supported in LLaMA and Qwen can be applied to LLaVA/VILA/LLaVA-OneVision as well, such as INT4/INT8 weight-only, SmoothQuant, and INT4 Activation-Aware Quantization (AWQ). Detailed instructions can be found in LLaMA [README](../llama/README.md) and Qwen [README](../qwen/README.md). - - For example, - - ```bash - # INT4 weight only - python ../llama/convert_checkpoint.py \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --dtype float16 \ - --output_dir tmp/trt_models/${MODEL_NAME}/int4_weightonly/1-gpu \ - --use_weight_only \ - --weight_only_precision int4 - - # INT4 AWQ - python ../../../quantization/quantize.py \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --output_dir tmp/trt_models/${MODEL_NAME}/int4_awq/1-gpu \ - --dtype float16 \ - --qformat int4_awq \ - --calib_size 32 - ``` - - Then follow the same `trtllm-build` and `run.py` steps as before. NOTE: for `trtllm-build` command, do not use `--use_fused_mlp=enable` in these quantization modes. - -## MLLaMA - -This section shows how to build and run a LLaMA-3.2 Vision model in TensorRT-LLM. We use [Llama-3.2-11B-Vision/](https://huggingface.co/meta-llama/Llama-3.2-11B-Vision) as an example. - -For LLaMA-3.2 text model, please refer to the [examples/models/core/llama/README.md](../llama/README.md) because it shares the model architecture of llama. - -### Support data types - * BF16 - * Tensor Parallel - * INT8 & INT4 Weight-Only - * FP8 - -### Build and run vision model - -* build engine of vision encoder model - -```bash -python examples/models/core/multimodal/build_multimodal_engine.py --model_type mllama \ - --model_path Llama-3.2-11B-Vision/ \ - --output_dir /tmp/mllama/trt_engines/vision/ -``` - -* build engine of decoder model - -```bash -python examples/models/core/mllama/convert_checkpoint.py --model_dir Llama-3.2-11B-Vision/ \ - --output_dir /tmp/mllama/trt_ckpts \ - --dtype bfloat16 - -trtllm-build --checkpoint_dir /tmp/mllama/trt_ckpts \ - --output_dir /tmp/mllama/trt_engines/llm/ \ - --max_num_tokens 4096 \ - --max_seq_len 2048 \ - --workers 1 \ - --gemm_plugin auto \ - --max_batch_size 4 \ - --max_encoder_input_len 4100 \ - --input_timing_cache model.cache -``` - -Note that for instruct Vision model, please set the `max_encoder_input_len` as `6404`. - -* Run test on multimodal/run.py with C++ runtime (LLM part only) - -```bash -python3 examples/models/core/multimodal/run.py --engine_dir /tmp/mllama/trt_engines/ \ - --hf_model_dir Llama-3.2-11B-Vision/ \ - --image_path https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg \ - --input_text "<|image|><|begin_of_text|>If I had to write a haiku for this one" \ - --max_new_tokens 50 \ - --batch_size 2 - -Use model_runner_cpp by default. To switch to model_runner, set `--session python` in the command mentioned above. - -python3 examples/models/core/multimodal/eval.py \ - --engine_dir /tmp/mllama/trt_engines/ \ - --hf_model_dir Llama-3.2-11B-Vision/ \ - --test_trtllm \ - --accuracy_threshold 65 \ - --eval_task lmms-lab/ai2d -``` - -### Run MLLaMA decoder part by FP8 - -```bash -# install modelopt 0.21.0 -pip install nvidia-modelopt[torch]~=0.21.0 - -python ./examples/quantization/quantize.py --model_dir Llama-3.2-11B-Vision/ \ - --dtype bfloat16 \ - --qformat fp8 \ - --output_dir /tmp/llama-3.2-11B-Vision/fp8/ \ - --kv_cache_dtype fp8 \ - --calib_size 512 \ - --calib_dataset scienceqa - -trtllm-build --checkpoint_dir /tmp/llama-3.2-11B-Vision/fp8/ \ - --output_dir /tmp/trt_engines/llama-3.2-11B-Vision/fp8/llm \ - --max_num_tokens 4096 \ - --max_seq_len 2048 \ - --workers 1 \ - --gemm_plugin auto \ - --max_batch_size 4 \ - --max_encoder_input_len 4100 \ - --input_timing_cache model.cache \ - --use_paged_context_fmha enable \ - --use_fp8_context_fmha enable - -# copy visiual engine directory `/tmp/mllama/trt_engines/vision/` to fp8 engine directory `/tmp/trt_engines/llama-3.2-11B-Vision/fp8/vision` - -python3 examples/models/core/multimodal/run.py --engine_dir /tmp/trt_engines/llama-3.2-11B-Vision/fp8/ \ - --hf_model_dir Llama-3.2-11B-Vision/ \ - --image_path https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg \ - --input_text "<|image|><|begin_of_text|>If I had to write a haiku for this one" \ - --max_new_tokens 50 \ - --batch_size 2 - -python3 examples/models/core/multimodal/eval.py --engine_dir /tmp/trt_engines/llama-3.2-11B-Vision/fp8/ \ - --hf_model_dir Llama-3.2-11B-Vision/ \ - --test_trtllm \ - --accuracy_threshold 65 \ - --eval_task lmms-lab/ai2d -``` - -Note that for instruct Vision model, please set the `max_encoder_input_len` as `6404`. - -## NeVA - -[NeVA](https://docs.nvidia.com/nemo-framework/user-guide/24.12/nemotoolkit/multimodal/mllm/neva.html) is a groundbreaking addition to the NeMo Multimodal ecosystem. This model seamlessly integrates large language-centric models with a vision encoder, that can be deployed in TensorRT-LLM. - -1. Generate TRT-LLM engine for NVGPT following example in `examples/models/core/gpt/README.md`. To adhere to the NVGPT conventions of the conversion script, some layer keys have to be remapped using `--nemo_rename_key`. - - ```bash - export MODEL_NAME="neva" - python ../gpt/convert_checkpoint.py \ - --nemo_ckpt_path ./${MODEL_NAME}.nemo \ - --dtype bfloat16 \ - --output_dir tmp/trt_models/${MODEL_NAME} \ - --nemo_rename_key model:model.language_model \ - attention.linear_qkv.layer_norm_bias:input_layernorm.bias \ - attention.linear_qkv.layer_norm_weight:input_layernorm.weight \ - mlp.linear_fc1.layer_norm_bias:post_attention_layernorm.bias \ - mlp.linear_fc1.layer_norm_weight:post_attention_layernorm.weight \ - linear_qkv:query_key_value \ - linear_fc1:dense_h_to_4h \ - linear_fc2:dense_4h_to_h \ - linear_proj:dense \ - decoder:encoder - - trtllm-build \ - --checkpoint_dir tmp/trt_models/${MODEL_NAME} \ - --output_dir tmp/trt_engines/${MODEL_NAME}/bf16/1-gpu/llm \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 \ - --max_batch_size 1 \ - --max_input_len 2048 \ - --max_seq_len 2560 \ - --max_multimodal_len 729 # 1 (max_batch_size) * 729 (num_visual_features) - ``` - -2. Build TensorRT engines for visual components - - ```bash - python build_multimodal_engine.py --model_path ./${MODEL_NAME}.nemo --model_type neva --output_dir tmp/trt_engines/${MODEL_NAME}/bf16/1-gpu/vision - ``` - - ```bash - python run.py \ - --max_new_tokens 30 \ - --hf_model_dir tmp/trt_models/${MODEL_NAME} \ - --engine_dir tmp/trt_engines/${MODEL_NAME}/bf16/1-gpu \ - --input_text "Question: which city is this? Answer:" - ``` - - Note: use `--run_profiling` for performance measurement, use `--check_accuracy` for accuracy check. - -## Nougat - -1. Download Huggingface weights - - ```bash - export MODEL_NAME="nougat-base" # also nougat-small - git clone https://huggingface.co/facebook/${MODEL_NAME} tmp/hf_models/${MODEL_NAME} - ``` - -2. Convert Huggingface weights into TRT-LLM checkpoints and build TRT engines using scripts in `examples/models/core/enc_dec` - - Nougat uses mBART architecture but replaces the LLM encoder with a Swin Transformer encoder. - To achieve this, we add an extra `--nougat` flag (over mBART example) to - `convert_checkpoint.py` in `examples/models/core/enc_dec` and `trtllm-build`. - - ```bash - python ../enc_dec/convert_checkpoint.py --model_type bart \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --output_dir tmp/trt_models/${MODEL_NAME}/bfloat16 \ - --tp_size 1 \ - --pp_size 1 \ - --dtype bfloat16 \ - --nougat - - trtllm-build --checkpoint_dir tmp/trt_models/${MODEL_NAME}/bfloat16/decoder \ - --output_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/bfloat16/llm/decoder \ - --paged_kv_cache disable \ - --moe_plugin disable \ - --gemm_plugin bfloat16 \ - --bert_attention_plugin bfloat16 \ - --gpt_attention_plugin bfloat16 \ - --remove_input_padding enable \ - --max_beam_width 1 \ - --max_batch_size 1 \ - --max_seq_len 101 \ - --max_input_len 1 \ - --max_encoder_input_len 588 # 1 (max_batch_size) * 588 (num_visual_features) - ``` - -3. Generate TensorRT engines for visual components and combine everything into final pipeline. - - ```bash - python build_multimodal_engine.py --model_type nougat --model_path tmp/hf_models/${MODEL_NAME} --output_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/bfloat16/vision - - python run.py \ - --hf_model_dir tmp/hf_models/${MODEL_NAME} \ - --engine_dir tmp/trt_engines/${MODEL_NAME}/1-gpu/bfloat16 - ``` - - Note: Nougat models usually do not need a text prompt. - - -## Phi-3-vision - -1. Download Huggingface weights - - ```bash - export MODEL_NAME="Phi-3-vision-128k-instruct" # or Phi-3.5-vision-instruct - git clone https://huggingface.co/microsoft/${MODEL_NAME} tmp/hf_models/${MODEL_NAME} - ``` - -2. Convert Huggingface weights into TRT-LLM checkpoints and build TRT engines using scripts in `examples/models/core/phi`. - ```bash - python ../phi/convert_checkpoint.py \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --output_dir tmp/trt_models/${MODEL_NAME}/fp16/1-gpu \ - --dtype float16 - - trtllm-build \ - --checkpoint_dir tmp/trt_models/${MODEL_NAME}/fp16/1-gpu \ - --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/llm \ - --gpt_attention_plugin float16 \ - --gemm_plugin float16 \ - --max_batch_size 1 \ - --max_input_len 4096 \ - --max_seq_len 4608 \ - --max_multimodal_len 4096 - ``` - -3. Generate TensorRT engines for visual components and combine everything into final pipeline. - - ```bash - python build_multimodal_engine.py --model_type phi-3-vision --model_path tmp/hf_models/${MODEL_NAME} --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/vision - - python run.py \ - --hf_model_dir tmp/hf_models/${MODEL_NAME} \ - --kv_cache_free_gpu_memory_fraction 0.7 \ - --engine_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/ \ - --image_path=https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png - ``` -## Phi-4-multimodal -Navigate to the folder `TensorRT-LLM/examples/models/core/multimodal` - -1. Download Huggingface weights - - ```bash - export MODEL_NAME="Phi-4-multimodal-instruct" - export HF_DIR="tmp/hf_models/${MODEL_NAME}" - export CKPT_DIR="tmp/trt_models/${MODEL_NAME}/fp16/1-gpu" - export ENGINE_DIR="tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu" - git clone https://huggingface.co/microsoft/${MODEL_NAME} ${HF_DIR} - - ``` - -2. Convert Huggingface weights into TRT-LLM checkpoints and build TRT engines using scripts in `examples/models/core/phi`. - ```bash - python ../phi/convert_checkpoint.py \ - --model_dir ${HF_DIR} \ - --output_dir ${CKPT_DIR} \ - --dtype float16 - - trtllm-build \ - --checkpoint_dir ${CKPT_DIR} \ - --output_dir ${ENGINE_DIR} \ - --gpt_attention_plugin float16 \ - --gemm_plugin float16 \ - --max_batch_size 1 \ - --max_input_len 4096 \ - --max_seq_len 4608 \ - --max_multimodal_len 4096 - ``` - -3. Generate TensorRT engines for visual components and combine everything into final pipeline. -*Note: the encoders are not the TRT engines but are pure Pytorch ones* - - ```bash - python build_multimodal_engine.py --model_type phi-4-multimodal --model_path ${HF_DIR} --output_dir ${ENGINE_DIR} - - python run.py \ - --hf_model_dir ${HF_DIR} \ - --kv_cache_free_gpu_memory_fraction 0.7 \ - --engine_dir ${ENGINE_DIR} \ - --image_path=https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png - --audio_path=${HF_DIR}/examples/what_is_shown_in_this_image.wav - ``` -## Qwen2-VL -[Qwen2-VL Family](https://github.com/QwenLM/Qwen2-VL): is the latest version of the vision language models in the Qwen model families. Here we show how to deploy Qwen2-VL 2B and 7B in TensorRT-LLM. - -Firstly, please install transformers and qwen-vl-utils -```bash -pip install -r requirements-qwen2vl.txt -``` -### Support data types - * FP16 - * FP8 - -### Build and run vision model -* Download Huggingface weights - ```bash - export MODEL_NAME="Qwen2-VL-7B-Instruct" # or Qwen2-VL-2B-Instruct - git clone https://huggingface.co/Qwen/${MODEL_NAME} tmp/hf_models/${MODEL_NAME} - ``` -* Build engine of decoder model - - ```bash - python3 ../qwen/convert_checkpoint.py \ - --model_dir=tmp/hf_models/${MODEL_NAME} \ - --output_dir=tmp/trt_models/${MODEL_NAME}/fp16/1-gpu \ - --dtype float16 - - trtllm-build --checkpoint_dir tmp/trt_models/${MODEL_NAME}/fp16/1-gpu \ - --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/llm \ - --gemm_plugin=float16 \ - --gpt_attention_plugin=float16 \ - --max_batch_size=4 \ - --max_input_len=2048 \ - --max_seq_len=3072 \ - --max_multimodal_len=1296 #(max_batch_size) * 324 (num_visual_features), this's for image_shape=[504,504] - ``` - -* Build engine of vision encoder model - ```bash - python build_multimodal_engine.py --model_type qwen2_vl --model_path tmp/hf_models/${MODEL_NAME} --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/vision - ``` - -* Run test on multimodal/run.py with C++ runtime (LLM part only) - ```bash - python3 run.py \ - --hf_model_dir tmp/hf_models/${MODEL_NAME} \ - --engine_dir tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/ - ``` -### Run Qwen2-VL decoder part by FP8 -* Build engine - ```bash - python ./examples/quantization/quantize.py \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir tmp/trt_models/${MODEL_NAME}/fp8/1-gpu \ - --calib_size 512 - - trtllm-build --checkpoint_dir tmp/trt_models/${MODEL_NAME}/fp8/1-gpu \ - --output_dir tmp/trt_engines/${MODEL_NAME}/fp8/1-gpu/llm \ - --max_input_len=2048 \ - --max_seq_len 3072 \ - --gemm_plugin auto \ - --max_batch_size 4 \ - --max_multimodal_len=1296 - - # copy visiual engine directory `tmp/trt_engines/${MODEL_NAME}/fp16/1-gpu/vision/` to fp8 engine directory `tmp/trt_engines/${MODEL_NAME}/fp8/1-gpu/vision` - ``` -* Run test on multimodal/run.py with C++ runtime (LLM part only) - ```bash - python3 run.py \ - --hf_model_dir tmp/hf_models/${MODEL_NAME} \ - --engine_dir tmp/trt_engines/${MODEL_NAME}/fp8/1-gpu/ - ``` +- [Supported models](https://nvidia.github.io/TensorRT-LLM/models/supported-models.html) +- [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) and the + [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) +- Multimodal serving examples under `examples/llm-api/` and + `examples/serve/` ## Qwen-Image-Bench Evaluator @@ -1183,153 +31,3 @@ python examples/models/core/multimodal/qwen_image_bench_eval.py \ The script evaluates all five Qwen-Image-Bench level-1 dimensions by default: Quality, Aesthetics, Alignment, Real-world Fidelity, and Creative Generation. - -## Video NeVA - -[Video NeVA](https://github.com/NVIDIA/NeMo/blob/main/docs/source/multimodal/mllm/video_neva.rst) is a groundbreaking addition to the NeMo Multimodal ecosystem that could work with video modality. This model seamlessly integrates large language-centric models with a vision encoder, that can be deployed in TensorRT-LLM. - -1. Generate TRT-LLM engine for Nemotron model following example in `examples/models/core/nemotron/README.md`. To adhere to the NVGPT conventions of the conversion script. This will be used as our base LM for inference. - - ```bash - pip install decord # used for loading video - - python3 ../../../quantization/quantize.py \ - --nemo_ckpt_path /path/to/nemotron/model.nemo \ - --dtype bfloat16 \ - --batch_size 64 \ - --qformat full_prec \ - --output_dir nemotron-3/trt_ckpt/bf16/1-gpu - - - trtllm-build \ - --checkpoint_dir nemotron-3/trt_ckpt/bf16/1-gpu \ - --output_dir tmp/trt_engines/nemotron-3/bf16/1-gpu/llm \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 \ - --max_batch_size 1 \ - --max_input_len 4096 \ - --max_seq_len 4352 \ - --max_multimodal_len 3072 # 1 (max_batch_size) * (12 num_frames) * (256 image_token_len) - ``` - -2. Build TensorRT engines for visual components - - ```bash - python build_multimodal_engine.py --model_path /path/to/video/neva/projector.nemo --model_type video-neva --output_dir tmp/trt_engines/nemotron-3/visual_encoder --output_dir tmp/trt_engines/nemotron-3/bf16/1-gpu/vision - ``` - - ```bash - python run.py \ - --max_new_tokens 30 \ - --hf_model_dir nemotron-3/trt_ckpt/bf16/1-gpu \ - --engine_dir tmp/trt_engines/nemotron-3/bf16/1-gpu \ - --input_text "Question: what is in the video? Answer:" \ - --video_path /path/to/your/local/video/file - ``` - - Note: use `--run_profiling` for performance measurement, use `--check_accuracy` for accuracy check. - -## Dataset Evaluation - -This section explains how to evaluate datasets using our provided script, including supported models and configurations. - -### Evaluation Command -To run an evaluation, use the following command: - -```bash -python ./examples/models/core/multimodal/eval.py \ - --model_type \ - --engine_dir \ - --hf_model_dir \ - --dataset_dir \ - --test_trtllm (or --test_hf, or both) \ - --accuracy_threshold \ - --eval_task \ - --max_ite 20 \ - --visual_engine_name -``` - -### Parameters -- `--model_type`: Specify the model type to evaluate. -- `--engine_dir`: Path to the model engines directory. -- `--hf_model_dir`: Path to the Hugging Face model directory. -- `--dataset_dir`: Path to the dataset directory. If not specified, will load the dataset from HF with the `--eval_task` as dataset tag. -- `--test_trtllm` or `--test_hf`: Specify which evaluation framework to use. Both can be used simultaneously. -- `--accuracy_threshold`: Set the accuracy threshold for evaluation. -- `--eval_task`: Specify the evaluation task. Supported tasks: `['lmms-lab/ai2d', 'lmms-lab/VQAv2', 'lmms-lab/MME']`. Default to `'lmms-lab/VQAv2'`. -- `--max_ite`: Maximum number of iterations, default to 20. -- `--visual_engine_name`: Name of the visual engine. - -### Supported Evaluation Tasks -The following evaluation tasks are supported: -- `lmms-lab/ai2d` -- `lmms-lab/VQAv2` -- `lmms-lab/MME` - -### Supported Model Types -The script supports the following model types: -- `blip2` -- `fuyu` -- `kosmos-2` -- `llava` -- `llava_next` -- `llava_onevision` -- `phi-3-vision` -- `qwen2_vl` -- `mllama` -- `vila` -- `cogvlm` -- `neva` -- `internvl` - -**Note:** The models `vila`, `cogvlm`, `neva`, and `internvl` do not support the `--test_hf` evaluation framework. - -## Enabling Tensor Parallelism for multi-GPU - -The LLM part of the pipeline can be run on multiple GPUs using tensor parallelism. -The visual encoder will be replicated on each GPU and operate in a data parallel fashion. - -To enable tensor parallelism, both weight conversion step (from Huggingface to FT format) -and engine building step should use additional arguments. Finally `run.py` should be prefixed -with `mpirun -n NUM_GPUS --allow-run-as-root`. - -The full set of commands to enable 2-way tensor parallelism for LLaVA is: - - ```bash - export MODEL_NAME="llava-1.5-7b-hf" - - python ../llama/convert_checkpoint.py \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --output_dir tmp/trt_models/${MODEL_NAME}/fp16/2-gpu \ - --dtype float16 --tp_size 2 - - trtllm-build \ - --checkpoint_dir tmp/trt_models/${MODEL_NAME}/fp16/2-gpu \ - --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/2-gpu/llm \ - --gemm_plugin float16 \ - --max_batch_size 1 \ - --max_input_len 2048 \ - --max_seq_len 2560 \ - --max_multimodal_len 576 - - python build_multimodal_engine.py --model_type llava --model_path tmp/hf_models/${MODEL_NAME} --output_dir tmp/trt_engines/${MODEL_NAME}/fp16/2-gpu/vision - - mpirun -n 2 --allow-run-as-root \ - python run.py \ - --max_new_tokens 30 \ - --hf_model_dir tmp/hf_models/${MODEL_NAME} \ - --engine_dir tmp/trt_engines/${MODEL_NAME}/fp16/2-gpu \ - ``` -## Enabling Embedding Table Offloading - -Embedding Table Offloading is a memory optimization technique that helps manage large embedding tables more efficiently. It offloads the embedding table to CPU memory and uses a chunked prefetching mechanism during processing. This approach is only available when operating in context chunk mode. - -To enable this feature, use the `--mm_embedding_offloading` argument: -```bash -python run.py \ - --enable_chunked_context \ - --mm_embedding_offloading true \ - --hf_model_dir ${HF_MODEL_PATH} \ - --engine_dir ${ENGINE_PATH} -``` -When not explicitly specified, this feature automatically enables if you're using a multimodal model along with context chunking enabled. diff --git a/examples/models/core/multimodal/__init__.py b/examples/models/core/multimodal/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/examples/models/core/multimodal/build_multimodal_engine.py b/examples/models/core/multimodal/build_multimodal_engine.py deleted file mode 100644 index 59e8bb4ffa25..000000000000 --- a/examples/models/core/multimodal/build_multimodal_engine.py +++ /dev/null @@ -1,12 +0,0 @@ -import argparse - -from tensorrt_llm.tools.multimodal_builder import (MultimodalEngineBuilder, - add_multimodal_arguments) - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser = add_multimodal_arguments(parser) - args = parser.parse_args() - - builder = MultimodalEngineBuilder(args) - builder.build() diff --git a/examples/models/core/multimodal/eval.py b/examples/models/core/multimodal/eval.py deleted file mode 100644 index 01804d290dee..000000000000 --- a/examples/models/core/multimodal/eval.py +++ /dev/null @@ -1,317 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. - -import argparse -import os - -import aiohttp -import datasets -import torch -from transformers import AutoProcessor -from utils import add_common_args - -import tensorrt_llm -import tensorrt_llm.profiler as profiler -from tensorrt_llm import logger -from tensorrt_llm.runtime import MultimodalModelRunner - -SUPPORTED_MODEL_TYPES = { - 'blip2': 'Blip2ForConditionalGeneration', - 'fuyu': 'FuyuForCausalLM', - 'kosmos-2': 'Kosmos2ForConditionalGeneration', - 'llava': 'LlavaForConditionalGeneration', - 'llava_next': 'LlavaNextForConditionalGeneration', - 'llava_onevision': 'LlavaOnevisionForConditionalGeneration', - 'phi-3-vision': 'AutoModelForCausalLM', - 'qwen2_vl': 'Qwen2VLForConditionalGeneration', # not tested for TRT-LLM yet - 'mllama': 'MllamaForConditionalGeneration', - 'vila': None, - 'cogvlm': None, # not tested for TRT-LLM yet - 'neva': None, # not tested for TRT-LLM yet - 'internvl': None, -} -EVAL_TASKS = ['lmms-lab/ai2d', 'lmms-lab/VQAv2', 'lmms-lab/MME'] - - -def parse_arguments(args=None): - parser = argparse.ArgumentParser() - parser = add_common_args(parser) - parser.add_argument('--test_trtllm', - action='store_true', - default=None, - help="Evaluate the TensorRT-LLM.") - parser.add_argument('--test_hf', - action='store_true', - default=None, - help="Evaluate the Huggingface.") - parser.add_argument('--max_ite', type=int, default=20) - parser.add_argument('--eval_task', - type=str, - choices=EVAL_TASKS, - default='lmms-lab/VQAv2') - parser.add_argument('--model_type', - type=str, - default=None, - choices=SUPPORTED_MODEL_TYPES.keys()) - parser.add_argument( - '--accuracy_threshold', - type=float, - default=None, - help= - 'used to check the accuracy of test_trtllm. Should be between 0 and 100.' - ) - parser.add_argument( - '--dataset_dir', - type=str, - default=None, - help="The local directory of the dataset for evaluation; " - "will download the dataset from huggingface hub if not specified.") - parser.add_argument( - '--dataset_cache_dir', - type=str, - default=None, - help="The local cache directory for dataset; " - "will use `~/.cache/huggingface/datasets` if not specified.") - return parser.parse_args(args=args) - - -def load_dataset(args) -> datasets.Dataset: - split_name = 'validation' if 'VQAv2' in args.eval_task else 'test' - - if args.dataset_dir is not None and os.path.exists( - os.path.join(args.dataset_dir, "dataset_info.json")): - logger.info(f"load dataset by load_from_disk from {args.dataset_dir}") - dataset = datasets.load_from_disk(args.dataset_dir) - - else: - logger.info( - f"load dataset by load_dataset from {args.dataset_dir or args.eval_task}" - ) - dataset = datasets.load_dataset( - args.dataset_dir or args.eval_task, - cache_dir=args.dataset_cache_dir, - split=split_name, - storage_options={ - 'client_kwargs': { - 'timeout': aiohttp.ClientTimeout(total=3600) - } - }, - trust_remote_code=True, - ) - return dataset - - -def load_hf_model(args): - if SUPPORTED_MODEL_TYPES[args.model_type] is None: - raise ValueError(f"Unsupported HF model_type: {args.model_type}") - profiler.start('load HF model') - model_class = getattr(__import__('transformers'), - SUPPORTED_MODEL_TYPES[args.model_type]) - hf_model = model_class.from_pretrained(args.hf_model_dir, - dtype=torch.float16, - device_map="cuda:0", - trust_remote_code=True) - profiler.stop('load HF model') - - logger.info( - f'Load HF model takes: {profiler.elapsed_time_in_sec("load HF model")} sec' - ) - return hf_model - - -def load_trtllm_model(args): - profiler.start('load TensorRT LLM model') - trtllm_model = MultimodalModelRunner(args) - profiler.stop('load TensorRT LLM model') - logger.info( - f'Load TensorRT LLM model takes: {profiler.elapsed_time_in_sec("load TensorRT LLM model")} sec' - ) - return trtllm_model - - -def prepare_prompts(task, data, model_type, processor) -> str: - prompts = None - question = data['question'] - if question[-1] != '?': - question += '?' - - if task == 'lmms-lab/ai2d': - for j, option in enumerate(data['options']): - question += f" ({j}) {option}" - - if model_type in ['blip2', 'neva']: - prompts = f"Question: {question} Answer: " - elif model_type == 'fuyu': - prompts = f"Answer the following {task} question based on the image: {question}" - elif model_type == 'kosmos-2': - prompts = f" Question: {question} Answer: " - elif model_type == 'cogvlm': - prompts = f"A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: {question} ASSISTANT:" - elif model_type in ['llava', 'llava_next', 'llava_onevision']: - conversation = [ - { - "role": - "user", - "content": [ - { - "type": "image" - }, - { - "type": "text", - "text": question - }, - ], - }, - ] - prompts = processor.apply_chat_template(conversation, - add_generation_prompt=True) - elif model_type in ['vila', 'internvl']: - prompts = f"\n{question}" - elif model_type == 'phi-3-vision': - messages = [ - { - "role": "user", - "content": f"<|image_1|>\n{question}" - }, - ] - prompts = processor.tokenizer.apply_chat_template( - messages, tokenize=False, add_generation_prompt=True) - elif model_type == 'qwen2_vl': - conversation = [{ - "role": - "user", - "content": [{ - "type": "image", - }, { - "type": "text", - "text": question - }] - }] - prompts = processor.apply_chat_template(conversation, - add_generation_prompt=True) - elif model_type == 'mllama': - prompts = processor.apply_chat_template(images=data['image'], - text=question + "; answer: ") - else: - raise ValueError(f"Unsupported model_type: {model_type}") - - return prompts - - -def eval(output, task, data) -> bool: - output = output.strip().lower() - if task == 'lmms-lab/VQAv2': - return any(answer['answer'] in output for answer in data['answers']) - else: - return data['answer'].lower() in output - - -os.environ["TOKENIZERS_PARALLELISM"] = "false" -args = parse_arguments() -if args.model_type not in SUPPORTED_MODEL_TYPES: - raise ValueError(f"Unsupported model_type: {args.model_type}") - -logger.set_level(args.log_level) - -runtime_rank = tensorrt_llm.mpi_rank() -dataset = load_dataset(args) -hf_model = load_hf_model(args) if args.test_hf else None -trtllm_model = load_trtllm_model(args) if args.test_trtllm else None -if SUPPORTED_MODEL_TYPES[args.model_type] is None: - hf_processor = None -else: - hf_processor = AutoProcessor.from_pretrained(args.hf_model_dir, - trust_remote_code=True) -hf_correct = trtllm_correct = 0 - -if args.model_type == 'mllama': - from tensorrt_llm.runtime.processor_wrapper import MllamaProcessorWrapper - hf_processor = MllamaProcessorWrapper(hf_processor, logger) - -torch.random.manual_seed(0) -profiler.start('evaluation') -if args.test_trtllm or args.test_hf: - for i in range(args.max_ite): - logger.debug(f"Ite: {i:3d}") - data = dataset[i] - if i > len(dataset): - break - prompts = prepare_prompts(args.eval_task, data, args.model_type, - hf_processor) - image = data['image'] - - if args.test_hf: - assert hf_model is not None, f"Unsupported HF model_type: {args.model_type}" - profiler.start('hf') - inputs = hf_processor( - images=image, - text=prompts, - return_tensors="pt", - ).to(hf_model.device, - torch.float16) # add torch.float16 for llava-onevision - input_length = inputs.input_ids.shape[-1] - hf_output = hf_model.generate(**inputs, - max_new_tokens=args.max_new_tokens) - hf_result = (hf_processor.batch_decode( - hf_output, skip_special_tokens=True)[0] if args.model_type in [ - 'blip2' - ] else hf_processor.decode(hf_output[0][input_length:], - skip_special_tokens=True)) - hf_correct += eval(hf_result, args.eval_task, data) - profiler.stop('hf') - - if args.test_trtllm: - profiler.start('tensorrt_llm') - _, output_text = trtllm_model.run( - input_text=prompts, - input_image=image, - input_audio=None, - max_new_tokens=args.max_new_tokens) - if runtime_rank == 0: - trtllm_result = output_text[0][0] - trtllm_correct += eval(trtllm_result, args.eval_task, data) - profiler.stop('tensorrt_llm') - - if runtime_rank == 0: - if args.eval_task == 'lmms-lab/VQAv2': - answer = data['answers'] - else: - answer = data['answer'] - logger.debug(f"prompts: {prompts}") - logger.debug(f"reference answer: {answer}") - if args.test_hf: - logger.debug(f"HF's answer: {hf_result}") - if args.test_trtllm: - logger.debug(f"TRT-LLM's answer: {trtllm_result}") - - if runtime_rank == 0: - logger.info(f"total iterations: {args.max_ite}") - if args.test_hf: - logger.info( - f"HF's accuracy: {100 * hf_correct / args.max_ite:4.2f}%") - if args.test_trtllm: - logger.info( - f"TRT-LLM's accuracy: {100 * trtllm_correct / args.max_ite:4.2f}%" - ) - # check if the accuracy is above the threshold - if args.accuracy_threshold is not None and args.test_trtllm: - assert trtllm_correct / args.max_ite >= args.accuracy_threshold / 100, \ - f"TRT-LLM's accuracy is below the threshold: {args.accuracy_threshold}%." -else: - logger.info("Neither enable test_trtllm nor enable test_hf") - -profiler.stop('evaluation') -logger.info( - f'Evaluation takes: {profiler.elapsed_time_in_sec("evaluation")} sec') diff --git a/examples/models/core/multimodal/requirements-eclair.txt b/examples/models/core/multimodal/requirements-eclair.txt deleted file mode 100644 index 281c8ae93a91..000000000000 --- a/examples/models/core/multimodal/requirements-eclair.txt +++ /dev/null @@ -1 +0,0 @@ -timm diff --git a/examples/models/core/multimodal/requirements-internlm-xcomposer2.txt b/examples/models/core/multimodal/requirements-internlm-xcomposer2.txt deleted file mode 100644 index 5d27a97ebb68..000000000000 --- a/examples/models/core/multimodal/requirements-internlm-xcomposer2.txt +++ /dev/null @@ -1 +0,0 @@ -transformers==4.56.0 diff --git a/examples/models/core/multimodal/requirements-llava_onevision.txt b/examples/models/core/multimodal/requirements-llava_onevision.txt deleted file mode 100644 index 126cbda08e15..000000000000 --- a/examples/models/core/multimodal/requirements-llava_onevision.txt +++ /dev/null @@ -1,4 +0,0 @@ -git+https://github.com/LLaVA-VL/LLaVA-NeXT.git -transformers>=4.56.0 -einops -av diff --git a/examples/models/core/multimodal/requirements-qwen2vl.txt b/examples/models/core/multimodal/requirements-qwen2vl.txt deleted file mode 100644 index 50f14d1d8095..000000000000 --- a/examples/models/core/multimodal/requirements-qwen2vl.txt +++ /dev/null @@ -1,2 +0,0 @@ -accelerate -qwen-vl-utils==0.0.8 # 0.0.9 has bug https://github.com/QwenLM/Qwen2-VL/pull/673, rollback until a newer version is released diff --git a/examples/models/core/multimodal/requirements-vila.txt b/examples/models/core/multimodal/requirements-vila.txt deleted file mode 100644 index 00775aa0cdfa..000000000000 --- a/examples/models/core/multimodal/requirements-vila.txt +++ /dev/null @@ -1,2 +0,0 @@ -git+https://github.com/bfshi/scaling_on_scales.git -transformers==4.56.0 diff --git a/examples/models/core/multimodal/run.py b/examples/models/core/multimodal/run.py deleted file mode 100644 index fb13554848e7..000000000000 --- a/examples/models/core/multimodal/run.py +++ /dev/null @@ -1,128 +0,0 @@ -import argparse -import os - -from utils import add_common_args, compute_str_match_rate - -import tensorrt_llm -import tensorrt_llm.profiler as profiler -from tensorrt_llm import logger -from tensorrt_llm.runtime import MultimodalModelRunner - - -def print_result(model, input_text, output_text, args): - logger.info("---------------------------------------------------------") - if model.model_type != 'nougat': - logger.info(f"\n[Q] {input_text}") - for i in range(len(output_text)): - logger.info(f"\n[A]: {output_text[i]}") - - if args.num_beams == 1: - output_ids = model.tokenizer(output_text[0][0], - add_special_tokens=False)['input_ids'] - logger.info(f"Generated {len(output_ids)} tokens") - - if args.check_accuracy: - if model.model_type != 'nougat': - if model.model_type == "vila": - for i in range(len(args.image_path.split(args.path_sep))): - if i % 2 == 0: - assert output_text[i][0].lower( - ) == "the image captures a bustling city intersection teeming with life. from the perspective of a car's dashboard camera, we see" - else: - assert output_text[i][0].lower( - ) == "the image captures the iconic merlion statue in singapore, a renowned worldwide landmark. the merlion, a mythical" - elif model.model_type == "llava": - for i in range(len(args.image_path.split(args.path_sep))): - assert output_text[i][0].lower() == 'singapore' - elif model.model_type == 'fuyu': - assert output_text[0][0].lower() == '4' - elif model.model_type == "pix2struct": - assert "characteristic | cat food, day | cat food, wet | cat treats" in output_text[ - 0][0].lower() - elif model.model_type in [ - 'blip2', 'neva', 'phi-3-vision', 'llava_next', - 'phi-4-multimodal', 'pixtral' - ]: - assert 'singapore' in output_text[0][0].lower() - elif model.model_type == 'video-neva': - assert 'robot' in output_text[0][0].lower() - elif model.model_type == 'kosmos-2': - assert 'snowman' in output_text[0][0].lower() - elif model.model_type == "mllama": - if "If I had to write a haiku for this one" in input_text: - ref_1 = ", it would be:.\\nPeter Rabbit is a rabbit.\\nHe lives in a cozy little house.\\nHe's a very good rabbit.\\" - ref_2 = "Here is a haiku for the image:\n\n" - - elif "Answer:" in input_text: - ref_1 = "2,173. A 1 2 3 4 5 6 Date Income 2005-12-17" - ref_2 = "Answer: 2,173. 1 2 3 4 5 6 Date Income 2005-12-17" - - elif "The key to life is" in input_text: - ref_1 = "to find your passion and pursue it with all your heart. For me, that passion is photography. I love capturing the beauty of the world around me" - ref_2 = "not to be found in the external world," - output = output_text[0][0] - match_rate = max(compute_str_match_rate(ref_1, output), - compute_str_match_rate(ref_2, output)) - logger.info(f"match rate: {match_rate}") - assert match_rate >= 50, \ - f"expected results: '{ref_1}' or '{ref_2}', generated results: '{output}'" - - elif model.model_type == 'llava_onevision': - if args.video_path is None: - assert 'singapore' in output_text[0][0].lower() - else: - assert 'the video is funny because the child\'s actions are' in output_text[ - 0][0].lower() - elif model.model_type == "qwen2_vl": - assert 'dog' in output_text[0][0].lower() - else: - assert output_text[0][0].lower() == 'singapore' - - if args.run_profiling: - msec_per_batch = lambda name: 1000 * profiler.elapsed_time_in_sec( - name) / args.profiling_iterations - logger.info('Latencies per batch (msec)') - logger.info('e2e generation: %.1f' % (msec_per_batch('Generate'))) - logger.info(' ' * 2 + 'Preprocessing: %.1f' % - (msec_per_batch('Preprocess'))) - logger.info(' ' * 4 + 'Vision encoder: %.1f' % - (msec_per_batch('Vision encoder'))) - if profiler.elapsed_time_in_sec('Feature transform') is not None: - logger.info(' ' * 4 + 'Feature transform: %.1f' % - (msec_per_batch('Feature transform'))) - logger.info(' ' * 2 + 'LLM generate: %.1f' % (msec_per_batch('LLM'))) - logger.info(' ' * 2 + 'Tokenizer decode: %.1f' % - (msec_per_batch('Tokenizer decode'))) - - logger.info("---------------------------------------------------------") - - -if __name__ == '__main__': - os.environ["TOKENIZERS_PARALLELISM"] = "false" - parser = argparse.ArgumentParser() - parser = add_common_args(parser) - args = parser.parse_args() - logger.set_level(args.log_level) - - model = MultimodalModelRunner(args) - visual_data = model.load_test_data(args.image_path, args.video_path) - audio_data = model.load_test_audio(args.audio_path) - - if args.run_profiling: - num_warmup_iters = 3 # Multiple iterations to load both vision and LLM engines into memory - for _ in range(num_warmup_iters): - input_text, output_text = model.run(args.input_text, visual_data, - audio_data, args.max_new_tokens) - profiler.reset() - - num_iters = args.profiling_iterations if args.run_profiling else 1 - - for _ in range(num_iters): - input_text, output_text = model.run(args.input_text, visual_data, - audio_data, args.max_new_tokens) - - runtime_rank = tensorrt_llm.mpi_rank() - if runtime_rank == 0: - print_result(model, input_text, output_text, args) - -# TODO: raise error if VILA mode 1 with C++ runtime diff --git a/examples/models/core/multimodal/utils.py b/examples/models/core/multimodal/utils.py deleted file mode 100644 index 6d13fb600f40..000000000000 --- a/examples/models/core/multimodal/utils.py +++ /dev/null @@ -1,163 +0,0 @@ -def add_common_args(parser): - parser.add_argument('--max_new_tokens', type=int, default=128) - parser.add_argument('--batch_size', type=int, default=1) - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument('--engine_dir', - type=str, - default=None, - help='Directory containing visual and LLM TRT engines') - parser.add_argument('--visual_engine_name', - type=str, - default='model.engine', - help='Name of visual TRT engine') - parser.add_argument('--audio_engine_name', - type=str, - default='model.engine', - help='Name of audio TRT engine') - parser.add_argument('--hf_model_dir', - type=str, - default=None, - help="Directory containing tokenizer") - parser.add_argument('--input_text', - type=str, - nargs='+', - default=None, - help='Text prompt to LLM') - parser.add_argument('--num_beams', - type=int, - help="Use beam search if num_beams >1", - default=1) - parser.add_argument('--top_k', type=int, default=1) - parser.add_argument('--top_p', type=float, default=0.0) - parser.add_argument('--temperature', type=float, default=1.0) - parser.add_argument('--repetition_penalty', type=float, default=1.0) - parser.add_argument('--run_profiling', - action='store_true', - help='Profile runtime over several iterations') - parser.add_argument('--profiling_iterations', - type=int, - help="Number of iterations to run profiling", - default=20) - parser.add_argument('--check_accuracy', - action='store_true', - help='Check correctness of text output') - parser.add_argument( - '--video_path', - type=str, - default=None, - help= - 'Path to your local video file, using \'llava-onevision-accuracy\' to check the Llava-OneVision model accuracy' - ) - parser.add_argument( - '--video_num_frames', - type=int, - help= - "The number of frames sampled from the video in the Llava-OneVision model.", - default=None) - parser.add_argument("--image_path", - type=str, - nargs='+', - default=None, - help='List of input image paths, separated by symbol') - parser.add_argument("--audio_path", - type=str, - default=None, - help='input audio path') - parser.add_argument("--path_sep", - type=str, - default=",", - help='Path separator symbol') - parser.add_argument("--prompt_sep", - type=str, - default=",", - help="Prompt separator symbol") - parser.add_argument('--enable_context_fmha_fp32_acc', - action='store_true', - default=None, - help="Enable FMHA runner FP32 accumulation.") - parser.add_argument( - '--enable_chunked_context', - action='store_true', - help='Enables chunked context (only available with cpp session).', - ) - parser.add_argument( - '--mm_embedding_offloading', - type=lambda s: s.lower() == "true", - default=None, - help= - 'Enable position table offloading. When not specified, defaults to True if using a multimodal model with chunked context.' - ) - parser.add_argument( - '--session', - default='cpp_llm_only', - type=str, - choices=['python', 'cpp_llm_only', 'cpp'], - help= - 'Rumtime used to run the models. \n`cpp_llm_only`: vision engine run in python runtime, but LLM in pybind cpp runtime\n`python`: everything runs in python runtime\n`cpp`: everything runs in C++ runtime' - ) - parser.add_argument( - '--kv_cache_free_gpu_memory_fraction', - default=0.7, - type=float, - help='Specify the free gpu memory fraction.', - ) - parser.add_argument( - '--cross_kv_cache_fraction', - default=0.5, - type=float, - help= - 'Specify the kv cache fraction reserved for cross attention. Only applicable for encoder-decoder models. By default 0.5 for self and 0.5 for cross.', - ) - parser.add_argument( - '--multi_block_mode', - type=lambda s: s.lower() in - ("yes", "true", "t", "1" - ), # custom boolean function to convert input string to boolean - default=True, - help= - "Distribute the work across multiple CUDA thread-blocks on the GPU for masked MHA kernel." - ) - parser.add_argument( - '--lora_task_uids', - type=str, - default=None, - nargs="+", - help="The list of LoRA task uids; use -1 to disable the LoRA module") - parser.add_argument('--debug_mode', - default=False, - action='store_true', - help="Whether or not to turn on the debug mode") - parser.add_argument( - '--trust_remote_code', - action='store_true', - default=False, - help='Allow loading models with custom remote code from HuggingFace Hub. ' - 'Only enable this for models from trusted sources.') - return parser - - -def levenshtein_distance(s1, s2): - if len(s1) < len(s2): - return levenshtein_distance(s2, s1) - - if len(s2) == 0: - return len(s1) - - previous_row = range(len(s2) + 1) - for i, c1 in enumerate(s1): - current_row = [i + 1] - for j, c2 in enumerate(s2): - insertions = previous_row[j + 1] + 1 - deletions = current_row[j] + 1 - substitutions = previous_row[j] + (c1 != c2) - current_row.append(min(insertions, deletions, substitutions)) - previous_row = current_row - - return previous_row[-1] - - -def compute_str_match_rate(s1, s2): - distance = levenshtein_distance(s1, s2) - max_length = max(len(s1), len(s2)) - match_rate = (1 - distance / max_length) * 100 - return match_rate diff --git a/examples/models/core/nemotron/README_nemotron-3.md b/examples/models/core/nemotron/README_nemotron-3.md deleted file mode 100644 index 0816eb995fd8..000000000000 --- a/examples/models/core/nemotron/README_nemotron-3.md +++ /dev/null @@ -1,214 +0,0 @@ -# Nemotron-3 - -This document demonstrates how to build the Nemotron models using TensorRT LLM and run on a single GPU or multiple GPUs. - -- [Nemotron](#nemotron) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [Download weights from HuggingFace Transformers](#download-weights-from-huggingface-transformers) - - [Build TensorRT engine(s)](#build-tensorrt-engines) - - [FP8 Quantization](#fp8-quantization) - - [INT4 AWQ Quantization](#int4-awq-quantization) - - [Run Inference](#run-inference) - -## Overview - -The TensorRT LLM Nemotron implementation is based on the GPT model, which can be found in [`tensorrt_llm/models/gpt/model.py`](../../../../tensorrt_llm/models/gpt/model.py). The TensorRT LLM Nemotron example is located in [`examples/models/core/nemotron`](./). - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`run.py`](../../../run.py) to run the inference on an input text; -* [`summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16/BF16 - * FP8 - * INT4 AWQ - * Tensor Parallel - * Pipeline Parallel - * Inflight Batching - * PAGED_KV_CACHE - * STRONGLY TYPED - * checkpoint type: Nemo, Huggingface (HF) - -## Nemo checkpoint - Usage - -### Download weights from HuggingFace Transformers - - -Install the dependencies and setup `git-lfs`. - -```bash -# Install dependencies -pip install -r requirements.txt - -# Setup git-lfs -git lfs install -``` - -Download one or more Nemotron models that you would like to build to TensorRT LLM engines. You can download from the [HuggingFace](https://huggingface.co) hub: - -```bash -# Download nemotron-3-8b-base-4k -git clone https://huggingface.co/nvidia/nemotron-3-8b-base-4k - -# Download nemotron-3-8b-chat-4k-sft -git clone https://huggingface.co/nvidia/nemotron-3-8b-chat-4k-sft - -# Download nemotron-3-8b-chat-4k-rlhf -git clone https://huggingface.co/nvidia/nemotron-3-8b-chat-4k-rlhf -``` - -### Build TensorRT engine(s) -The [`examples/quantization/quantize.py`](../../../quantization/quantize.py) script can quantize the Nemotron models and export to TensorRT LLM checkpoints. You may optionally skip the quantization step by specifying `--qformat full_prec` and thus export float16 or bfloat16 TensorRT LLM checkpoints. - -The `trtllm-build` command builds TensorRT LLM engines from TensorRT LLM checkpoints. The number of engine files is same to the number of GPUs used to run inference. Normally, `trtllm-build` uses one GPU by default, but if you have already more GPUs available at build time, you may enable parallel builds to make the engine building process faster by adding the `--workers` argument. - -Here are some examples: - -```bash -# single gpu, dtype bfloat16 -python3 ../../../quantization/quantize.py \ - --nemo_ckpt_path nemotron-3-8b-base-4k/Nemotron-3-8B-Base-4k.nemo \ - --dtype bfloat16 \ - --batch_size 64 \ - --qformat full_prec \ - --output_dir nemotron-3-8b/trt_ckpt/bf16/1-gpu - -trtllm-build --checkpoint_dir nemotron-3-8b/trt_ckpt/bf16/1-gpu \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 \ - --output_dir nemotron-3-8b/trt_engines/bf16/1-gpu -``` - -```bash -# 2-way tensor parallelism -python3 ../../../quantization/quantize.py \ - --nemo_ckpt_path nemotron-3-8b-base-4k/Nemotron-3-8B-Base-4k.nemo \ - --dtype bfloat16 \ - --batch_size 64 \ - --qformat full_prec \ - --tp_size 2 \ - --output_dir nemotron-3-8b/trt_ckpt/bf16/tp2 - -trtllm-build --checkpoint_dir nemotron-3-8b/trt_ckpt/bf16/tp2 \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 \ - --workers 2 \ - --output_dir nemotron-3-8b/trt_engines/bf16/tp2 -``` - -```bash -# 2-way tensor parallelism for both calibration and inference -mpirun -np 2 \ - python3 ../../../quantization/quantize.py \ - --nemo_ckpt_path nemotron-3-8b-base-4k/Nemotron-3-8B-Base-4k.nemo \ - --dtype bfloat16 \ - --batch_size 64 \ - --qformat full_prec \ - --calib_tp_size 2 \ - --tp_size 2 \ - --output_dir nemotron-3-8b/trt_ckpt/bf16/tp2 - -trtllm-build --checkpoint_dir nemotron-3-8b/trt_ckpt/bf16/tp2 \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 \ - --workers 2 \ - --output_dir nemotron-3-8b/trt_engines/bf16/tp2 -``` - -#### FP8 Quantization - -Quantize the Nemotron models to FP8 by specifying `--qformat fp8` to `quantize.py`. - -```bash -# single gpu, fp8 quantization -python3 ../../../quantization/quantize.py \ - --nemo_ckpt_path nemotron-3-8b-base-4k/Nemotron-3-8B-Base-4k.nemo \ - --dtype bfloat16 \ - --batch_size 64 \ - --qformat fp8 \ - --output_dir nemotron-3-8b/trt_ckpt/fp8/1-gpu - -trtllm-build --checkpoint_dir nemotron-3-8b/trt_ckpt/fp8/1-gpu \ - --gpt_attention_plugin bfloat16 \ - --output_dir nemotron-3-8b/trt_engines/fp8/1-gpu -``` - -#### INT4 AWQ Quantization - -Quantize the Nemotron models using INT4 AWQ by specifying `--qformat int4_awq` to `quantize.py`. - -```bash -# single gpu, int4 awq quantization -python3 ../../../quantization/quantize.py \ - --nemo_ckpt_path nemotron-3-8b-base-4k/Nemotron-3-8B-Base-4k.nemo \ - --dtype bfloat16 \ - --batch_size 64 \ - --qformat int4_awq \ - --output_dir nemotron-3-8b/trt_ckpt/int4_awq/1-gpu - -trtllm-build --checkpoint_dir nemotron-3-8b/trt_ckpt/int4_awq/1-gpu \ - --gpt_attention_plugin bfloat16 \ - --output_dir nemotron-3-8b/trt_engines/int4_awq/1-gpu -``` - -### Run Inference - -The `summarize.py` script can run the built engines to summarize the articles from the -[cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -```bash -# single gpu -python3 ../../../summarize.py --test_trt_llm \ - --no_add_special_tokens \ - --engine_dir nemotron-3-8b/trt_engines/bf16/1-gpu \ - --vocab_file nemotron-3-8b/trt_ckpt/bf16/1-gpu/tokenizer.model - -# multiple gpus -mpirun -np 2 \ - python3 ../../../summarize.py --test_trt_llm \ - --no_add_special_tokens \ - --engine_dir nemotron-3-8b/trt_engines/bf16/tp2 \ - --vocab_file nemotron-3-8b/trt_ckpt/bf16/tp2/tokenizer.model -``` - -If the engines are run successfully, you will see output like: -``` -...... -[04/23/2024-09:55:54] [TRT-LLM] [I] TensorRT LLM (total latency: 14.926485538482666 sec) -[04/23/2024-09:55:54] [TRT-LLM] [I] TensorRT LLM (total output tokens: 2000) -[04/23/2024-09:55:54] [TRT-LLM] [I] TensorRT LLM (tokens per second: 133.99001357980129) -[04/23/2024-09:55:54] [TRT-LLM] [I] TensorRT LLM beam 0 result -[04/23/2024-09:55:54] [TRT-LLM] [I] rouge1 : 19.48743720965424 -[04/23/2024-09:55:54] [TRT-LLM] [I] rouge2 : 6.272381295466071 -[04/23/2024-09:55:54] [TRT-LLM] [I] rougeL : 15.011005943152721 -[04/23/2024-09:55:54] [TRT-LLM] [I] rougeLsum : 17.76145734406502 -``` - -## HF checkpoint - Usage -Support for Nemotron models was added with transformers 4.44.0 release. - -```bash -# install transformers library -pip install transformers>=4.44.0 -# Download hf minitron model -git clone https://huggingface.co/nvidia/Minitron-4B-Base - -# Convert to TensorRT LLM checkpoint -python3 ../gpt/convert_checkpoint.py --model_dir Minitron-4B-Base \ - --dtype bfloat16 \ - --output_dir minitron/trt_ckpt/bf16/1-gpu - -# Build TensorRT LLM engines -trtllm-build --checkpoint_dir minitron/trt_ckpt/bf16/1-gpu \ - --gemm_plugin auto \ - --output_dir minitron/trt_engines/bf16/1-gpu - -# Run inference -python3 ../../../run.py --engine_dir minitron/trt_engines/bf16/1-gpu \ - --tokenizer_dir Minitron-4B-Base \ - --input_text "def print_hello_world():" \ - --max_output_len 20 -``` diff --git a/examples/models/core/nemotron/requirements.txt b/examples/models/core/nemotron/requirements.txt deleted file mode 100644 index f8d1f8e1abf4..000000000000 --- a/examples/models/core/nemotron/requirements.txt +++ /dev/null @@ -1,7 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -nemo-toolkit[all]==2.0.0rc1 -megatron-core @ git+https://github.com/NVIDIA/Megatron-LM@core_r0.8.0 -datasets==3.1.0 -evaluate -rouge_score diff --git a/examples/models/core/qwen2audio/README.md b/examples/models/core/qwen2audio/README.md deleted file mode 100644 index 92115345077e..000000000000 --- a/examples/models/core/qwen2audio/README.md +++ /dev/null @@ -1,80 +0,0 @@ -# Guide to Qwen2-Audio deployment pipeline - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -1. Download the Qwen2-Audio model. - ```bash - git lfs install - export MODEL_PATH="tmp/Qwen2-Audio-7B-Instruct" - git clone https://huggingface.co/Qwen/Qwen2-Audio-7B-Instruct $MODEL_PATH - ``` -2. Generate the TensorRT engine of audio encoder. - ```bash - export ENGINE_DIR="./trt_engines/qwen2audio/fp16" - python3 ../multimodal/build_multimodal_engine.py --model_type qwen2_audio --model_path $MODEL_PATH --max_batch_size 32 --output_dir ${ENGINE_DIR}/audio - ``` - - The TensorRT engine will be generated under `${ENGINE_DIR}/audio`. - -3. Build Qwen2 LLM TensorRT engine. -- Convert checkpoint - 1. Install packages - ```bash - pip install -r requirements.txt - ``` - 2. Convert - 2.1 FP16 checkpoint - ```bash - python3 ../qwen/convert_checkpoint.py --model_dir=$MODEL_PATH \ - --dtype=float16 \ - --output_dir=./tllm_checkpoint_1gpu_fp16 - ``` - 2.2 (Optional) INT8 Weight Only checkpoint - ```bash - python3 ../qwen/convert_checkpoint.py --model_dir=$MODEL_PATH \ - --dtype=float16 \ - --use_weight_only \ - --weight_only_precision=int8 \ - --output_dir=./tllm_checkpoint_1gpu_fp16_wo8 - ``` - -- Build TensorRT LLM engine - - NOTE: `max_prompt_embedding_table_size = query_token_num * max_batch_size`, therefore, if you change `max_batch_size`, `--max_prompt_embedding_table_size` must be reset accordingly. - ```bash - trtllm-build --checkpoint_dir=./tllm_checkpoint_1gpu_fp16 \ - --gemm_plugin=float16 --gpt_attention_plugin=float16 \ - --max_batch_size=1 --max_prompt_embedding_table_size=4096 \ - --output_dir=${ENGINE_DIR}/llm - ``` - The built Qwen engines are located in `${ENGINE_DIR}/llm`. - - You can replace the `--checkpoint_dir` with INT8 Weight Only checkpoint to build INT8 Weight Only engine as well. - For more information about Qwen, refer to the README.md in [`example/models/core/qwen`](../qwen). - -4. Assemble everything into the Qwen2-Audio pipeline. - - 4.1 Run with FP16 LLM engine - ```bash - python3 run.py \ - --tokenizer_dir=$MODEL_PATH \ - --engine_dir=${ENGINE_DIR}/llm \ - --audio_engine_path=${ENGINE_DIR}/audio/model.engine \ - --audio_url='./audio/glass-breaking-151256.mp3' - ``` - 4.2 (Optional) For multiple rounds of dialogue, you can run: - ```bash - python3 run_chat.py \ - --tokenizer_dir=$MODEL_PATH \ - --engine_dir=${ENGINE_DIR}/llm \ - --audio_engine_path=${ENGINE_DIR}/audio/model.engine \ - --max_new_tokens=256 - ``` - - Note: - - This example supports reusing the KV Cache for audio segments by assigning unique audio IDs. - - To further optimize performance, users can also cache the audio features (encoder output) to bypass the audio encoder if the original audio data remains unchanged. diff --git a/examples/models/core/qwen2audio/audio/glass-breaking-151256.mp3 b/examples/models/core/qwen2audio/audio/glass-breaking-151256.mp3 deleted file mode 100644 index 150e1080f2e5..000000000000 Binary files a/examples/models/core/qwen2audio/audio/glass-breaking-151256.mp3 and /dev/null differ diff --git a/examples/models/core/qwen2audio/requirements.txt b/examples/models/core/qwen2audio/requirements.txt deleted file mode 100644 index 1d6d844e4b64..000000000000 --- a/examples/models/core/qwen2audio/requirements.txt +++ /dev/null @@ -1,10 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score -transformers>=4.45.0 -transformers-stream-generator -sentencepiece>=0.1.99 -tiktoken -einops diff --git a/examples/models/core/qwen2audio/run.py b/examples/models/core/qwen2audio/run.py deleted file mode 100644 index a0b0a68fb1c9..000000000000 --- a/examples/models/core/qwen2audio/run.py +++ /dev/null @@ -1,640 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 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. -import argparse -import json -import os -from io import BytesIO -from urllib.request import urlopen - -import librosa -import tensorrt as trt -import torch -from transformers import AutoConfig, AutoProcessor, AutoTokenizer -from utils import add_common_args - -import tensorrt_llm -import tensorrt_llm.profiler as profiler -from tensorrt_llm import logger -from tensorrt_llm.llmapi.kv_cache_type import KVCacheType -from tensorrt_llm.quantization import QuantMode -from tensorrt_llm.runtime import (PYTHON_BINDINGS, ModelConfig, ModelRunner, - SamplingConfig, Session, TensorInfo) - -if PYTHON_BINDINGS: - from tensorrt_llm.runtime import ModelRunnerCpp - - -def get_engine_name(rank): - return "rank{}.engine".format(rank) - - -def trt_dtype_to_torch(dtype): - if dtype == trt.float16: - return torch.float16 - elif dtype == trt.float32: - return torch.float32 - elif dtype == trt.int32: - return torch.int32 - else: - raise TypeError("%s is not supported" % dtype) - - -class QWenInfer(object): - - def __init__(self, - audio_engine_path, - tokenizer_dir, - engine_dir, - log_level, - output_csv, - output_npy, - num_beams, - gpu_id=0): - self.audio_engine_path = audio_engine_path - self.tokenizer_dir = tokenizer_dir - self.engine_dir = engine_dir - self.log_level = log_level - self.max_seq_len = 0 - self.runner = None - self.hf_audio_tower = None - self.tokenizer = None - self.config = None - self.sampling_config = None - self.output_csv = output_csv - self.output_npy = output_npy - self.num_beams = num_beams - self.model_config = None - self.gpu_device = torch.device("cuda", gpu_id) - - def get_model(self): - # --load the tokenizer and engines # - tokenizer = AutoTokenizer.from_pretrained( - self.tokenizer_dir, - legacy=False, - trust_remote_code=True, - ) - processor = AutoProcessor.from_pretrained(self.tokenizer_dir, - trust_remote_code=True) - config_path = os.path.join(self.engine_dir, "config.json") - with open(config_path, "r") as f: - config = json.load(f) - self.max_seq_len = config["build_config"]["max_seq_len"] - assert self.max_seq_len > 0, "max_seq_len must be positive" - - gen_config_path = os.path.join(self.tokenizer_dir, - "generation_config.json") - with open(gen_config_path, "r") as f: - gen_config = json.load(f) - top_k = gen_config["top_k"] - top_p = gen_config["top_p"] - eos_token_id = tokenizer.pad_token_id - pad_token_id = tokenizer.pad_token_id - - use_gpt_attention_plugin = config["build_config"]["plugin_config"][ - "gpt_attention_plugin"] - remove_input_padding = config["build_config"]["plugin_config"][ - "remove_input_padding"] - dtype = config["pretrained_config"]["dtype"] - tp_size = config["pretrained_config"]["mapping"]["tp_size"] - pp_size = config["pretrained_config"]["mapping"]["pp_size"] - world_size = tp_size * pp_size - assert ( - world_size == tensorrt_llm.mpi_world_size() - ), f"Engine world size ({world_size}) != Runtime world size ({tensorrt_llm.mpi_world_size()})" - num_heads = config["pretrained_config"][ - "num_attention_heads"] // world_size - max_batch_size = config["build_config"]["max_batch_size"] - hidden_size = config["pretrained_config"]["hidden_size"] // world_size - vocab_size = config["pretrained_config"]["vocab_size"] - num_layers = config["pretrained_config"]["num_hidden_layers"] - num_kv_heads = config["pretrained_config"].get("num_key_value_heads", - num_heads) - if "kv_cache_type" in config["build_config"]: - kv_cache_type = KVCacheType(config["build_config"]["kv_cache_type"]) - else: - kv_cache_type = KVCacheType.CONTINUOUS - - tokens_per_block = config["build_config"]["plugin_config"][ - "tokens_per_block"] - max_prompt_embedding_table_size = config["build_config"].get( - "max_prompt_embedding_table_size", 0) - quant_mode = QuantMode.from_quant_algo( - config["pretrained_config"]["quantization"]["quant_algo"], - config["pretrained_config"]["quantization"]["kv_cache_quant_algo"], - ) - if config["pretrained_config"].get("multi_query_mode", False): - tensorrt_llm.logger.warning( - "`multi_query_mode` config is deprecated. Please rebuild the engine." - ) - num_kv_heads = 1 - - runtime_rank = tensorrt_llm.mpi_rank() - runtime_mapping = tensorrt_llm.Mapping(world_size=world_size, - rank=runtime_rank, - tp_size=tp_size, - pp_size=pp_size) - - model_config = ModelConfig( - max_batch_size=max_batch_size, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - hidden_size=hidden_size, - vocab_size=vocab_size, - num_layers=num_layers, - gpt_attention_plugin=use_gpt_attention_plugin, - kv_cache_type=kv_cache_type, - tokens_per_block=tokens_per_block, - remove_input_padding=remove_input_padding, - dtype=dtype, - quant_mode=quant_mode, - max_prompt_embedding_table_size=max_prompt_embedding_table_size, - max_beam_width=self.num_beams, - ) - sampling_config = SamplingConfig( - end_id=eos_token_id, - pad_id=pad_token_id, - num_beams=self.num_beams, - top_k=top_k, - top_p=top_p, - temperature=1.0, - ) - - engine_name = get_engine_name(runtime_rank) - serialize_path = os.path.join(self.engine_dir, engine_name) - print(f"Loading engine from {serialize_path}") - return ( - model_config, - sampling_config, - runtime_mapping, - runtime_rank, - serialize_path, - tokenizer, - processor, - eos_token_id, - pad_token_id, - ) - - def qwen_model_init(self, args): - logger.info(f"Loading audio engine from {self.audio_engine_path}") - with open(self.audio_engine_path, "rb") as f: - engine_buffer = f.read() - logger.info(f"Creating session from engine {self.audio_engine_path}") - self.session_audio = Session.from_serialized_engine(engine_buffer) - - self.config, _ = AutoConfig.from_pretrained( - self.tokenizer_dir, - return_unused_kwargs=True, - trust_remote_code=True, - ) - - ( - model_config, - sampling_config, - runtime_mapping, - runtime_rank, - serialize_path, - tokenizer, - processor, - eos_token_id, - pad_token_id, - ) = self.get_model() - runner_cls = ModelRunner if args.use_py_session else ModelRunnerCpp - runner_kwargs = dict( - engine_dir=args.engine_dir, - lora_dir=args.lora_dir, - rank=runtime_rank, - debug_mode=args.debug_mode, - lora_ckpt_source=args.lora_ckpt_source, - gpu_weights_percent=args.gpu_weights_percent, - max_output_len=args.max_new_tokens, - ) - if not args.use_py_session: - runner_kwargs.update( - is_enc_dec=False, - max_batch_size=model_config.max_batch_size, - max_input_len=self.max_seq_len - args.max_new_tokens, - max_beam_width=model_config.max_beam_width, - max_attention_window_size=args.max_attention_window_size, - sink_token_length=args.sink_token_length, - max_tokens_in_paged_kv_cache=args.max_tokens_in_paged_kv_cache, - kv_cache_enable_block_reuse=args.kv_cache_enable_block_reuse, - kv_cache_free_gpu_memory_fraction=args. - kv_cache_free_gpu_memory_fraction, - cross_kv_cache_fraction=None, - enable_chunked_context=args.enable_chunked_context, - multi_block_mode=args.multi_block_mode, - cuda_graph_mode=args.cuda_graph_mode, - device_ids=[args.gpu_id]) - runner_kwargs.update( - enable_context_fmha_fp32_acc=args.enable_context_fmha_fp32_acc) - self.runner = runner_cls.from_dir(**runner_kwargs) - self.tokenizer = tokenizer - self.processor = processor - self.sampling_config = sampling_config - self.model_config = model_config - - def ptuning_setup(self, prompt_table, dtype, hidden_size, tasks, input_ids): - if prompt_table is not None: - task_vocab_size = torch.tensor([prompt_table.shape[0]], - dtype=torch.int32, - device=self.gpu_device) - prompt_table = prompt_table.to( - dtype=tensorrt_llm._utils.str_dtype_to_torch(dtype), - device=self.gpu_device) - else: - prompt_table = torch.empty([1, hidden_size], device=self.gpu_device) - task_vocab_size = torch.zeros([1], device=self.gpu_device) - - if tasks is not None: - tasks = torch.tensor([int(t) for t in tasks.split(",")], - dtype=torch.int32, - device=self.gpu_device) - assert (tasks.shape[0] == input_ids.shape[0] - ), "Number of supplied tasks must match input batch size" - else: - tasks = torch.zeros([input_ids.size(0)], - dtype=torch.int32, - device=self.gpu_device) - - return [prompt_table, tasks, task_vocab_size] - - def build_user_input(self, audio=None, text=None): - assert isinstance(audio, str) or isinstance( - text, str), "audio or text must be provided as user input" - content = [] - if audio: - content.append({'type': 'audio', 'audio_url': audio}) - if text: - content.append({'type': 'text', 'text': text}) - user_input = {'role': 'user', 'content': content} - return user_input - - def get_raw_audios(self, audio_url): - audios = [] - for url in audio_url: - if os.path.isfile(url): - audio_data, _ = librosa.load( - url, sr=self.processor.feature_extractor.sampling_rate) - else: - audio_data, _ = librosa.load( - BytesIO(urlopen(url).read()), - sr=self.processor.feature_extractor.sampling_rate) - audios.append(audio_data) - return audios - - def audio_tower(self, audios, mask, stream, run_time=1): - audios = audios.to(self.gpu_device) - mask = mask.to(self.gpu_device) - audio_inputs = {"input": audios.float(), "mask": mask} - audio_output_info = self.session_audio.infer_shapes([ - TensorInfo("input", trt.DataType.FLOAT, audios.shape), - TensorInfo("mask", trt.DataType.HALF, mask.shape) - ]) - audio_outputs = { - t.name: - torch.empty(tuple(t.shape), - dtype=trt_dtype_to_torch(t.dtype), - device=self.gpu_device) - for t in audio_output_info - } - profiler.start("Audio") - for _ in range(run_time): - ok = self.session_audio.run(audio_inputs, audio_outputs, - stream.cuda_stream) - stream.synchronize() - audio_time = profiler.stop("Audio") / run_time - logger.info(f"TensorRT LLM Audio latency: {audio_time:3f} sec ") - - assert ok, "Runtime execution failed for audio session" - - audio_features = audio_outputs["output"] - - return audio_features - - def generate_for_qwen_audio( - self, - input_tokens, - args, - prompt_table=None, - extra_ids=None, - run_time=1, - ): - input_ids = torch.as_tensor(input_tokens, - device=self.gpu_device, - dtype=torch.int32) - input_lengths = torch.tensor([input_ids.size(1)], - device=self.gpu_device, - dtype=torch.int32) - max_input_length = torch.max(input_lengths).item() - max_new_tokens = min(args.max_new_tokens, - self.max_seq_len - max_input_length) - - prompt_table = prompt_table.unsqueeze(0) - profiler.start("QWen") - for _ in range(run_time): - outputs = self.runner.generate( - batch_input_ids=input_ids, - max_new_tokens=max_new_tokens, - max_attention_window_size=args.max_attention_window_size, - sink_token_length=args.sink_token_length, - end_id=self.sampling_config.end_id, - pad_id=self.sampling_config.pad_id, - temperature=args.temperature, - top_k=args.top_k, - top_p=args.top_p, - num_beams=args.num_beams, - num_return_sequences=args.num_return_sequences, - length_penalty=args.length_penalty, - early_stopping=args.early_stopping, - repetition_penalty=args.repetition_penalty, - presence_penalty=args.presence_penalty, - frequency_penalty=args.frequency_penalty, - stop_words_list=[[[151643], [151645]]], - bad_words_list=self.sampling_config.bad_words_list, - random_seed=args.random_seed, - lora_uids=args.lora_task_uids, - prompt_table=prompt_table, - prompt_tasks="0", - output_sequence_lengths=True, - no_repeat_ngram_size=args.no_repeat_ngram_size, - return_dict=True, - return_all_generated_tokens=False, - input_token_extra_ids=extra_ids) - output_ids = outputs['output_ids'] - torch.cuda.synchronize() - Qwen_time = profiler.stop("QWen") / run_time - - return output_ids, Qwen_time - - def get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor): - """ - Computes the output length of the convolutional layers and the output length of the audio encoder - """ - input_lengths = (input_lengths - 1) // 2 + 1 - output_lengths = (input_lengths - 2) // 2 + 1 - return input_lengths, output_lengths - - def qwen_infer(self, - input_text, - audios, - audio_ids, - args, - stream, - history=None, - past_audio_features=None, - run_time=1): - assert input_text, "input_text must be provided" - assert torch.cuda.is_available(), "no gpu available" - # preprocess on CPU maybe faster - device = torch.device("cpu") - if isinstance(history, list): - history.append(input_text) - full_text = self.processor.apply_chat_template( - history, add_generation_prompt=True, tokenize=False) - else: - full_text = input_text - inputs = self.processor( - text=full_text, - audios=audios, - return_tensors="pt", - padding=True, - sampling_rate=self.processor.feature_extractor.sampling_rate) - inputs = inputs.to(device) - input_ids = inputs.input_ids - - if hasattr(inputs, - 'input_features') and inputs.input_features is not None: - # audio tower - batch_size, _, max_mel_seq_len = inputs.input_features.shape - feature_attention_mask = inputs.feature_attention_mask - - audio_feat_lengths, num_audio_tokens = self.get_feat_extract_output_lengths( - feature_attention_mask.sum(-1)) - - max_seq_len = (max_mel_seq_len - 2) // 2 + 1 - # Create a sequence tensor of shape (batch_size, max_seq_len) - seq_range = (torch.arange(0, - max_seq_len, - dtype=audio_feat_lengths.dtype, - device=device).unsqueeze(0).expand( - batch_size, max_seq_len)) - lengths_expand = audio_feat_lengths.unsqueeze(1).expand( - batch_size, max_seq_len) - # Create mask - padding_mask = seq_range >= lengths_expand - - audio_attention_mask_ = padding_mask.view( - batch_size, 1, 1, max_seq_len).expand(batch_size, 1, - max_seq_len, max_seq_len) - audio_attention_mask = audio_attention_mask_.to(dtype=torch.float16, - device=device) - audio_attention_mask[audio_attention_mask_] = float("-inf") - - audio_features = self.audio_tower(inputs.input_features, - audio_attention_mask, stream, - run_time) - - # merge audio features and input ids - num_audios, max_audio_tokens, embed_dim = audio_features.shape - audio_features_mask = torch.arange( - max_audio_tokens, device=device).expand( - num_audios, - max_audio_tokens) < num_audio_tokens.unsqueeze(1) - masked_audio_features = audio_features[audio_features_mask].view( - -1, embed_dim) - batch_size, _ = input_ids.shape - - # 1. Create a mask to know where special audio tokens are - special_audio_token_mask = input_ids == self.config.audio_token_index - special_audio_token_num = special_audio_token_mask.sum().item() - if past_audio_features is not None: - assert isinstance(past_audio_features, - list), f'past_audio_features should be a list' - assert ( - special_audio_token_num == len(past_audio_features) + - num_audios - ), f'special_audio_token_num {special_audio_token_num} should be equal to len(past_audio_features) + num_audios ({len(past_audio_features)} + {num_audios})' - # split to get current audio features - cur_audio_features = torch.split(masked_audio_features, - num_audio_tokens.tolist()) - if len(past_audio_features) > 0: - # concat past and current audio features - masked_audio_features = torch.cat( - (torch.cat(past_audio_features).to( - masked_audio_features.device), - masked_audio_features)) - # get past audio tokens number - past_num_audio_tokens = torch.tensor([ - past_feat.size(0) for past_feat in past_audio_features - ]) - # concat past and current audio tokens number - num_audio_tokens = torch.cat( - (past_num_audio_tokens.to(num_audio_tokens.device), - num_audio_tokens)) - # extend past audio features, cache them in CPU memory - past_audio_features.extend( - [cur_feat.cpu() for cur_feat in cur_audio_features]) - - batch_indices, non_audio_indices = torch.where( - input_ids != self.config.audio_token_index) - - # 2. Fill the final input ids based on the mask. - batch_indices, audio_indices = torch.where( - input_ids == self.config.audio_token_index) - - vocab_size = self.config.vocab_size - fake_prompt_id = torch.arange(vocab_size, - vocab_size + num_audio_tokens.sum(), - device=device) - - input_ids[batch_indices, audio_indices] = fake_prompt_id - input_lengths = torch.tensor(input_ids.size(1), - dtype=torch.int32, - device=self.gpu_device) - dtype = self.model_config.dtype - prompt_table, tasks, task_vocab_size = self.ptuning_setup( - masked_audio_features, dtype, embed_dim, None, input_ids) - - # build extra ids - assert isinstance(audio_ids, list), "audio_ids must be a list" - assert ( - len(audio_ids) == num_audio_tokens.size(0) - ), f"audio_ids length doesn't match with num_audio_tokens ({len(audio_ids)} != {num_audio_tokens.size(0)})" - for i in audio_ids: - assert isinstance( - i, int - ) and i > 0, "audio_id should be an integer greater than 0" - extra_ids = torch.zeros_like(input_ids, - dtype=torch.int64, - device=device) - seq_extra_ids = torch.cat([ - torch.full((n, ), audio_ids[i], dtype=torch.int64) - for i, n in enumerate(num_audio_tokens) - ]).to(device) - extra_ids[batch_indices, audio_indices] = seq_extra_ids - extra_ids = extra_ids.tolist() - else: - input_ids = input_ids.to(dtype=torch.int32, device=self.gpu_device) - input_lengths = torch.tensor(input_ids.size(1), - dtype=torch.int32, - device=self.gpu_device) - dtype = self.model_config.dtype - prompt_table, tasks, task_vocab_size = self.ptuning_setup( - None, dtype, self.model_config.hidden_size, None, input_ids) - extra_ids = torch.zeros_like(input_ids, dtype=torch.int64).tolist() - - # print(f"extra_ids: {extra_ids}") - output_ids, Qwen_time = self.generate_for_qwen_audio( - input_ids, args, prompt_table, extra_ids, run_time) - - runtime_rank = tensorrt_llm.mpi_rank() - input_lengths = torch.tensor([input_ids.size(1)], - device=self.gpu_device, - dtype=torch.int32) - effective_output_token = 0 - if runtime_rank == 0: - if self.output_csv is None and self.output_npy is None: - for b in range(input_lengths.size(0)): - inputs = input_ids[b] - if self.num_beams <= 1: - outputs = output_ids[b][0, len(inputs):].tolist() - try: - effective_output_token = (effective_output_token + - outputs.index(151643)) - except: - effective_output_token = 1 - output_text = self.tokenizer.decode( - outputs, skip_special_tokens=True) - print(f'Output: "{output_text}"') - else: - for beam in range(self.num_beams): - outputs = output_ids[b][beam, len(inputs):].tolist() - output_text = self.tokenizer.decode( - outputs, skip_special_tokens=True) - print(f'Output(beam: {beam}): "{output_text}"') - logger.info(f"Input length={input_lengths[b]}") - logger.info(f"Output length={output_ids.shape}") - logger.info(f"TensorRT LLM QWen time: {Qwen_time:3f} sec ") - if isinstance(history, list): - history.append({'role': 'assistant', 'content': output_text}) - return output_text, past_audio_features - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument("--max_new_tokens", type=int, default=10) - parser.add_argument( - "--audio_engine_path", - type=str, - default="plan/audio_encoder/audio_encoder_fp16.plan", - ) - parser.add_argument( - "--input_text", - type=str, - default= - "<|audio_bos|><|AUDIO|><|audio_eos|>Generate the caption in English:") - parser.add_argument( - "--audio_url", - nargs="+", - type=str, - default=["./audio/glass-breaking-151256.mp3"], - ) - parser.add_argument( - "--input_tokens", - dest="input_file", - type=str, - help= - "CSV or Numpy file containing tokenized input. Alternative to text input.", - default=None, - ) - parser.add_argument( - "--output_csv", - type=str, - help="CSV file where the tokenized output is stored.", - default=None, - ) - parser.add_argument( - "--output_npy", - type=str, - help="Numpy file where the tokenized output is stored.", - default=None, - ) - parser.add_argument( - "--gpu_id", - type=int, - help= - "Specify GPU device index for running. Should be the index seen by torch, not original index", - default=0, - ) - parser = add_common_args(parser) - - return parser.parse_args() - - -if __name__ == "__main__": - args = parse_arguments() - tensorrt_llm.logger.set_level(args.log_level) - - # use cudaSetDevice before loading audio engine - torch.cuda.set_device(args.gpu_id) - qinfer = QWenInfer(args.audio_engine_path, args.tokenizer_dir, - args.engine_dir, args.log_level, args.output_csv, - args.output_npy, args.num_beams, args.gpu_id) - qinfer.qwen_model_init(args) - - audios = qinfer.get_raw_audios(args.audio_url) - gpu_device = torch.device("cuda", args.gpu_id) - stream = torch.cuda.current_stream(device=gpu_device) - qinfer.qwen_infer(args.input_text, audios, [1], args, stream, None, None, 1) diff --git a/examples/models/core/qwen2audio/run_chat.py b/examples/models/core/qwen2audio/run_chat.py deleted file mode 100644 index 00d58cdc862b..000000000000 --- a/examples/models/core/qwen2audio/run_chat.py +++ /dev/null @@ -1,81 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 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. - -# isort: off -import torch -from run import QWenInfer, parse_arguments - -import tensorrt_llm -# isort: on - -if __name__ == '__main__': - args = parse_arguments() - stream = torch.cuda.current_stream() - tensorrt_llm.logger.set_level(args.log_level) - qinfer = QWenInfer( - args.audio_engine_path, - args.tokenizer_dir, - args.engine_dir, - args.log_level, - args.output_csv, - args.output_npy, - args.num_beams, - ) - qinfer.qwen_model_init(args) - - run_i = 0 - history = [] - audios = None - global_audio_id = 1 - audio_ids = [] - - while True: - input_text = None - try: - input_text = input( - "Text (type 'q' to quit, or 'audio_url:[url]' to input audio): " - ) - except: - continue - - if input_text == "clear history": - history = [] - audios = None - continue - - if input_text.lower() == 'q': - break - print('\n') - - if input_text.startswith('audio_url:'): - audio_url = input_text[len('audio_url:'):].strip() - if isinstance(audios, list): - audios.extend(qinfer.get_raw_audios([audio_url])) - else: - audios = qinfer.get_raw_audios([audio_url]) - user_input = qinfer.build_user_input(audio=audio_url) - audio_ids.append(global_audio_id) - global_audio_id += 1 - else: - user_input = qinfer.build_user_input(text=input_text) - - qinfer.qwen_infer( - user_input, - audios, - audio_ids, - args, - stream, - history, - ) diff --git a/examples/models/core/qwen2audio/utils.py b/examples/models/core/qwen2audio/utils.py deleted file mode 100644 index 3252beebbf7d..000000000000 --- a/examples/models/core/qwen2audio/utils.py +++ /dev/null @@ -1,130 +0,0 @@ -from argparse import BooleanOptionalAction - - -def add_common_args(parser): - # sampling arguments - parser.add_argument('--num_beams', - type=int, - help="Use beam search if num_beams > 1", - default=1) - parser.add_argument('--num_return_sequences', - type=int, - help="Number of sequences to generate for each input.", - default=None) - parser.add_argument('--temperature', type=float, default=1.0) - parser.add_argument('--top_k', type=int, default=1) - parser.add_argument('--top_p', type=float, default=0.0) - parser.add_argument('--length_penalty', type=float, default=1.0) - parser.add_argument('--repetition_penalty', type=float, default=1.0) - parser.add_argument('--presence_penalty', type=float, default=0.0) - parser.add_argument('--frequency_penalty', type=float, default=0.0) - parser.add_argument('--random_seed', type=int, default=0) - parser.add_argument('--early_stopping', - type=int, - help='Use early stopping if num_beams > 1, ' - '1 for early-stopping, 0 for non-early-stopping' - 'other values for stopping by length', - default=1) - parser.add_argument('--no_repeat_ngram_size', type=int, default=None) - - # common runtime arguments - parser.add_argument('--sink_token_length', - type=int, - default=None, - help='The sink token length.') - parser.add_argument( - '--max_attention_window_size', - type=int, - default=None, - nargs="+", - help= - 'The attention window size that controls the sliding window attention kv cache behavior' - ) - parser.add_argument( - '--multi_block_mode', - type=lambda s: s.lower() in - ("yes", "true", "t", "1" - ), # custom boolean function to convert input string to boolean - default=True, - help= - "Distribute the work across multiple CUDA thread-blocks on the GPU for masked MHA kernel." - ) - parser.add_argument('--enable_context_fmha_fp32_acc', - action='store_true', - help="Enable FMHA runner FP32 accumulation.") - parser.add_argument('--cuda_graph_mode', - action='store_true', - help="Enable cuda graphs in the inference.") - parser.add_argument( - '--log_level', - type=str, - choices=['verbose', 'info', 'warning', 'error', 'internal_error'], - default='info') - parser.add_argument('--use_py_session', - default=False, - action='store_true', - help="Whether or not to use Python runtime session") - parser.add_argument('--debug_mode', - default=False, - action='store_true', - help="Whether or not to turn on the debug mode") - parser.add_argument('--lora_dir', - type=str, - default=None, - nargs="+", - help="The directory of LoRA weights") - parser.add_argument('--lora_ckpt_source', - type=str, - default="hf", - choices=["hf", "nemo"], - help="The source of lora checkpoint.") - parser.add_argument( - '--lora_task_uids', - type=str, - default=None, - nargs="+", - help="The list of LoRA task uids; use -1 to disable the LoRA module") - - # model arguments - parser.add_argument('--engine_dir', type=str, default='engine_outputs') - parser.add_argument('--hf_model_dir', '--model_dir', type=str, default=None) - parser.add_argument( - '--tokenizer_dir', - default=None, - help='tokenizer path; defaults to hf_model_dir if left unspecified') - - # memory argument - parser.add_argument( - '--gpu_weights_percent', - default=1, - type=float, - help= - 'Specify the percentage of weights that reside on GPU instead of CPU and streaming load during runtime.', - ) - parser.add_argument( - '--max_tokens_in_paged_kv_cache', - default=None, - type=int, - help= - 'Specify the maximum number of tokens in a kv cache page (only available with cpp session).', - ) - parser.add_argument( - '--kv_cache_enable_block_reuse', - default=True, - action=BooleanOptionalAction, - help= - 'Enables block reuse in kv cache (only available with cpp session).', - ) - parser.add_argument( - '--kv_cache_free_gpu_memory_fraction', - default=0.9, - type=float, - help='Specify the free gpu memory fraction.', - ) - parser.add_argument( - '--enable_chunked_context', - action='store_true', - help='Enables chunked context (only available with cpp session).', - ) - - return parser diff --git a/examples/models/core/qwenvl/README.md b/examples/models/core/qwenvl/README.md deleted file mode 100644 index 27fc5d1dfce0..000000000000 --- a/examples/models/core/qwenvl/README.md +++ /dev/null @@ -1,114 +0,0 @@ -# Guide to Qwen-VL deployment pipeline - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -1. Download the Qwen vision-language model (Qwen-VL). - ```bash - git lfs install - git clone https://huggingface.co/Qwen/Qwen-VL-Chat - ``` -2. Generate the Vision Transformer (ViT) ONNX model and the TensorRT engine. -- If you don't have ONNX file, run: - ```bash - python3 vit_onnx_trt.py --pretrained_model_path ./Qwen-VL-Chat - ``` - The ONNX and TensorRT engine will be generated under `./onnx/visual_encoder` and `./plan/visual_encoder` respectively. - -- If you already have an ONNX file under `./onnx/visual_encoder` and want to build a TensorRT engine with it, run: - ```bash - python3 vit_onnx_trt.py --pretrained_model_path ./Qwen-VL-Chat --only_trt - ``` - This command saves the test image tensor to `image.pt` for later pipeline inference. - -3. Build Qwen TensorRT engine. -- Convert checkpoint - 1. Install packages - ```bash - pip install -r requirements.txt - ``` - 2. Convert - ```bash - python3 ./examples/models/core/qwen/convert_checkpoint.py --model_dir=./Qwen-VL-Chat \ - --output_dir=./tllm_checkpoint_1gpu \ - --dtype float16 - ``` - -- Build TensorRT LLM engine - - NOTE: `max_prompt_embedding_table_size = query_token_num * max_batch_size`, therefore, if you change `max_batch_size`, `--max_prompt_embedding_table_size` must be reset accordingly. - ```bash - trtllm-build --checkpoint_dir=./tllm_checkpoint_1gpu \ - --gemm_plugin=float16 --gpt_attention_plugin=float16 \ - --max_input_len=2048 --max_seq_len=3072 \ - --max_batch_size=8 --max_prompt_embedding_table_size=2048 \ - --remove_input_padding=enable \ - --output_dir=./trt_engines/Qwen-VL-7B-Chat - ``` - The built Qwen engines are located in `./trt_engines/Qwen-VL-7B-Chat`. - For more information about Qwen, refer to the README.md in [`example/qwen`](../qwen). - -4. Assemble everything into the Qwen-VL pipeline. - - 4.1 Run with INT4 GPTQ weight-only quantization engine - ```bash - python3 run.py \ - --tokenizer_dir=./Qwen-VL-Chat \ - --qwen_engine_dir=./trt_engines/Qwen-VL-7B-Chat \ - --vit_engine_path=./plan/visual_encoder/visual_encoder_fp16.plan \ - --images_path='{"image": "./pics/demo.jpeg"}' - ``` - 4.2 (Optional) For multiple rounds of dialogue, you can run: - ```bash - python3 run_chat.py \ - --tokenizer_dir=./Qwen-VL-Chat \ - --qwen_engine_dir=./trt_engines/Qwen-VL-7B-Chat \ - --vit_engine_path=./plan/visual_encoder/visual_encoder_fp16.plan \ - --images_path='{"image": "./pics/demo.jpeg"}' - ``` - 4.3 (Optional) To show the bounding box result in the demo picture, install OpenCV, ZMQ, and request: - ```bash - pip install opencv-python==4.5.5.64 - pip install opencv-python-headless==4.5.5.64 - pip install zmq - pip install request - ``` - -   4.3.1 If the current program is executed on a remote machine, run the following command on a local machine: - - ```bash - python3 show_pic.py --ip=127.0.0.1 --port=8006 - ``` - -   Replace the `ip` and `port` values, where `ip` is your remote machine IP address. - -   Run the following command on the remote machine: - - ```bash - python3 run_chat.py \ - --tokenizer_dir=./Qwen-VL-Chat \ - --qwen_engine_dir=./trt_engines/Qwen-VL-7B-Chat \ - --vit_engine_path=./plan/visual_encoder/visual_encoder_fp16.plan \ - --display \ - --port=8006 - ``` - -   Replace the `port` value. - -   4.3.2 If the current program is executed on the local machine, run the following command: - - ```bash - python3 run_chat.py \ - --tokenizer_dir=./Qwen-VL-Chat \ - --qwen_engine_dir=./trt_engines/Qwen-VL-7B-Chat \ - --vit_engine_path=./plan/visual_encoder/visual_encoder_fp16.plan \ - --display \ - --local_machine - ``` - -   The question "Print the bounding box of the girl" is displayed. You should see the following image: - - ![image](./pics/1.png) diff --git a/examples/models/core/qwenvl/pics/1.png b/examples/models/core/qwenvl/pics/1.png deleted file mode 100644 index 09e872ac90da..000000000000 Binary files a/examples/models/core/qwenvl/pics/1.png and /dev/null differ diff --git a/examples/models/core/qwenvl/pics/demo.jpeg b/examples/models/core/qwenvl/pics/demo.jpeg deleted file mode 100644 index 9fdc04005062..000000000000 Binary files a/examples/models/core/qwenvl/pics/demo.jpeg and /dev/null differ diff --git a/examples/models/core/qwenvl/requirements.txt b/examples/models/core/qwenvl/requirements.txt deleted file mode 100644 index 9257ddb9545a..000000000000 --- a/examples/models/core/qwenvl/requirements.txt +++ /dev/null @@ -1,11 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score -transformers-stream-generator -sentencepiece>=0.1.99 -tiktoken -einops -matplotlib -torchvision diff --git a/examples/models/core/qwenvl/run.py b/examples/models/core/qwenvl/run.py deleted file mode 100644 index c996b20fd643..000000000000 --- a/examples/models/core/qwenvl/run.py +++ /dev/null @@ -1,549 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2023 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. -import argparse -import json -import os -from typing import List, Tuple - -import tensorrt as trt -import torch -from transformers import AutoConfig, AutoTokenizer -from vit_onnx_trt import Preprocss - -import tensorrt_llm -import tensorrt_llm.profiler as profiler -from tensorrt_llm import logger -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm.llmapi.kv_cache_type import KVCacheType -from tensorrt_llm.quantization import QuantMode -from tensorrt_llm.runtime import (ModelConfig, SamplingConfig, Session, - TensorInfo) - - -def get_engine_name(rank): - return "rank{}.engine".format(rank) - - -def trt_dtype_to_torch(dtype): - if dtype == trt.float16: - return torch.float16 - elif dtype == trt.float32: - return torch.float32 - elif dtype == trt.int32: - return torch.int32 - else: - raise TypeError("%s is not supported" % dtype) - - -class QWenInfer(object): - - def __init__( - self, - tokenizer_dir, - qwen_engine_dir, - log_level, - output_csv, - output_npy, - num_beams, - ): - self.tokenizer_dir = tokenizer_dir - self.qwen_engine_dir = qwen_engine_dir - self.log_level = log_level - self.global_max_input_len = 2048 - self.decoder = None - self.tokenizer = None - self.config = None - self.sampling_config = None - self.output_csv = output_csv - self.output_npy = output_npy - self.num_beams = num_beams - self.model_config = None - - def get_model(self): - # --load the tokenizer and engine # - tokenizer = AutoTokenizer.from_pretrained( - self.tokenizer_dir, - legacy=False, - trust_remote_code=True, - ) - config_path = os.path.join(self.qwen_engine_dir, "config.json") - with open(config_path, "r") as f: - config = json.load(f) - gen_config_path = os.path.join(self.tokenizer_dir, - "generation_config.json") - with open(gen_config_path, "r") as f: - gen_config = json.load(f) - top_k = gen_config["top_k"] - top_p = gen_config["top_p"] - chat_format = gen_config["chat_format"] - if chat_format == "raw": - eos_token_id = gen_config["eos_token_id"] - pad_token_id = gen_config["pad_token_id"] - elif chat_format == "chatml": - pad_token_id = eos_token_id = tokenizer.im_end_id - else: - raise Exception("unknown chat format ", chat_format) - - use_gpt_attention_plugin = config["build_config"]["plugin_config"][ - "gpt_attention_plugin"] - gemm_allreduce_plugin = config["build_config"]["plugin_config"][ - "gemm_allreduce_plugin"] - - remove_input_padding = config["build_config"]["plugin_config"][ - "remove_input_padding"] - dtype = config["pretrained_config"]["dtype"] - tp_size = config["pretrained_config"]["mapping"]["tp_size"] - pp_size = config["pretrained_config"]["mapping"]["pp_size"] - world_size = tp_size * pp_size - assert ( - world_size == tensorrt_llm.mpi_world_size() - ), f"Engine world size ({world_size}) != Runtime world size ({tensorrt_llm.mpi_world_size()})" - num_heads = config["pretrained_config"][ - "num_attention_heads"] // world_size - max_batch_size = config["build_config"]["max_batch_size"] - hidden_size = config["pretrained_config"]["hidden_size"] // world_size - vocab_size = config["pretrained_config"]["vocab_size"] - num_layers = config["pretrained_config"]["num_hidden_layers"] - num_kv_heads = config["pretrained_config"].get("num_key_value_heads", - num_heads) - if "kv_cache_type" in config["build_config"]: - kv_cache_type = KVCacheType(config["build_config"]["kv_cache_type"]) - else: - kv_cache_type = KVCacheType.CONTINUOUS - - tokens_per_block = config["build_config"]["plugin_config"][ - "tokens_per_block"] - max_prompt_embedding_table_size = config["build_config"].get( - "max_prompt_embedding_table_size", 0) - quant_mode = QuantMode.from_quant_algo( - config["pretrained_config"]["quantization"]["quant_algo"], - config["pretrained_config"]["quantization"]["kv_cache_quant_algo"], - ) - if config["pretrained_config"].get("multi_query_mode", False): - tensorrt_llm.logger.warning( - "`multi_query_mode` config is deprecated. Please rebuild the engine." - ) - num_kv_heads = 1 - - runtime_rank = tensorrt_llm.mpi_rank() - runtime_mapping = tensorrt_llm.Mapping(world_size=world_size, - rank=runtime_rank, - tp_size=tp_size, - pp_size=pp_size) - torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node) - - model_config = ModelConfig( - max_batch_size=max_batch_size, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - hidden_size=hidden_size, - vocab_size=vocab_size, - num_layers=num_layers, - gpt_attention_plugin=use_gpt_attention_plugin, - gemm_allreduce_plugin=gemm_allreduce_plugin, - kv_cache_type=kv_cache_type, - tokens_per_block=tokens_per_block, - remove_input_padding=remove_input_padding, - dtype=dtype, - quant_mode=quant_mode, - max_prompt_embedding_table_size=max_prompt_embedding_table_size, - max_beam_width=self.num_beams, - ) - sampling_config = SamplingConfig( - end_id=eos_token_id, - pad_id=pad_token_id, - num_beams=self.num_beams, - top_k=top_k, - top_p=top_p, - temperature=1.0, - ) - - engine_name = get_engine_name(runtime_rank) - serialize_path = os.path.join(self.qwen_engine_dir, engine_name) - print(f"Loading engine from {serialize_path}") - return ( - model_config, - sampling_config, - runtime_mapping, - runtime_rank, - serialize_path, - tokenizer, - eos_token_id, - pad_token_id, - ) - - def qwen_model_init(self): - ( - model_config, - sampling_config, - runtime_mapping, - runtime_rank, - serialize_path, - tokenizer, - eos_token_id, - pad_token_id, - ) = self.get_model() - with open(serialize_path, "rb") as f: - engine_buffer = f.read() - self.decoder = tensorrt_llm.runtime.GenerationSession( - model_config, - engine_buffer, - runtime_mapping, - ) - self.tokenizer = tokenizer - self.sampling_config = sampling_config - self.model_config = model_config - self.config, _ = AutoConfig.from_pretrained( - self.tokenizer_dir, - return_unused_kwargs=True, - trust_remote_code=True, - ) - - def ptuning_setup(self, prompt_table, dtype, hidden_size, tasks, input_ids): - if prompt_table is not None: - task_vocab_size = torch.tensor([prompt_table.shape[1]], - dtype=torch.int32, - device="cuda") - prompt_table = prompt_table.view( - (prompt_table.shape[0] * prompt_table.shape[1], - prompt_table.shape[2])) - prompt_table = prompt_table.cuda().to( - dtype=tensorrt_llm._utils.str_dtype_to_torch(dtype)) - else: - prompt_table = torch.empty([1, hidden_size]).cuda() - task_vocab_size = torch.zeros([1]).cuda() - - if tasks is not None: - tasks = torch.tensor([int(t) for t in tasks.split(",")], - dtype=torch.int32, - device="cuda") - assert (tasks.shape[0] == input_ids.shape[0] - ), "Number of supplied tasks must match input batch size" - else: - tasks = torch.zeros([input_ids.size(0)], dtype=torch.int32).cuda() - - return [prompt_table, tasks, task_vocab_size] - - def make_context( - self, - query: str, - history: List[Tuple[str, str]] = None, - system: str = "You are a helpful assistant.", - max_window_size: int = 6144, - ): - if history is None: - history = [] - - im_start, im_end = "<|im_start|>", "<|im_end|>" - im_start_tokens = [self.tokenizer.im_start_id] # 151644 - im_end_tokens = [self.tokenizer.im_end_id] # [151645] - nl_tokens = self.tokenizer.encode("\n") - - def _tokenize_str(role, content): - return f"{role}\n{content}", self.tokenizer.encode( - role, allowed_special=set(self.tokenizer.IMAGE_ST) - ) + nl_tokens + self.tokenizer.encode( - content, allowed_special=set(self.tokenizer.IMAGE_ST)) - - system_text, system_tokens_part = _tokenize_str("system", system) - system_tokens = im_start_tokens + system_tokens_part + im_end_tokens - - raw_text = "" - context_tokens = [] - - for turn_query, turn_response in reversed(history): - query_text, query_tokens_part = _tokenize_str("user", turn_query) - query_tokens = im_start_tokens + query_tokens_part + im_end_tokens - if turn_response is not None: - response_text, response_tokens_part = _tokenize_str( - "assistant", turn_response) - response_tokens = im_start_tokens + response_tokens_part + im_end_tokens - - next_context_tokens = (nl_tokens + query_tokens + nl_tokens + - response_tokens) - prev_chat = f"\n{im_start}{query_text}{im_end}\n{im_start}{response_text}{im_end}" - else: - next_context_tokens = nl_tokens + query_tokens + nl_tokens - prev_chat = f"\n{im_start}{query_text}{im_end}\n" - - current_context_size = (len(system_tokens) + - len(next_context_tokens) + - len(context_tokens)) - if current_context_size < max_window_size: - context_tokens = next_context_tokens + context_tokens - raw_text = prev_chat + raw_text - else: - break - - context_tokens = system_tokens + context_tokens - raw_text = f"{im_start}{system_text}{im_end}" + raw_text - context_tokens += (nl_tokens + im_start_tokens + - _tokenize_str("user", query)[1] + im_end_tokens + - nl_tokens + im_start_tokens + - self.tokenizer.encode("assistant") + nl_tokens) - raw_text += f"\n{im_start}user\n{query}{im_end}\n{im_start}assistant\n" - - return raw_text, context_tokens - - def generate_for_qwenvl( - self, - input_tokens, - max_new_tokens: int, - prompt_table=None, - tasks=None, - task_vocab_size=None, - num_beams=1, - ): - input_ids = None - input_lengths = None - input_ids = torch.as_tensor(input_tokens, - device="cuda", - dtype=torch.int32) - input_lengths = torch.tensor([input_ids.size(1)], - device="cuda", - dtype=torch.int32) - max_input_length = torch.max(input_lengths).item() - max_new_tokens = min(max_new_tokens, - self.global_max_input_len - max_input_length) - - profiler.start("QWen") - run_time = 10 - for _ in range(run_time): - self.decoder.setup( - batch_size=input_lengths.size(0), - max_context_length=max_input_length, - max_new_tokens=max_new_tokens, - beam_width=num_beams, - ) - output_ids = self.decoder.decode( - input_ids, - input_lengths, - self.sampling_config, - prompt_table, - tasks, - task_vocab_size, - ) - torch.cuda.synchronize() - profiler.stop("QWen") - Qwen_time = profiler.elapsed_time_in_sec("QWen") / run_time - - return output_ids, Qwen_time - - def qwen_infer( - self, - input_vit, - images_path, - input_text, - max_new_tokens, - num_beams=1, - history=None, - ): - if images_path is None: - content_list = [] - else: - content_list = images_path - if history is None: - history = [] - content_list.append({"text": input_text}) - query = self.tokenizer.from_list_format(content_list) - raw_text, context_tokens = self.make_context(query, history=history) - # context_tokens = self.tokenizer.encode(query) - input_ids = torch.tensor([context_tokens]).to("cuda") - bos_pos = torch.where(input_ids == self.config.visual["image_start_id"]) - eos_pos = torch.where( - input_ids == self.config.visual["image_start_id"] + 1) - assert (bos_pos[0] == eos_pos[0]).all() - img_pos = torch.stack((bos_pos[0], bos_pos[1], eos_pos[1]), dim=1) - vocab_size = self.config.vocab_size - fake_prompt_id = torch.arange( - vocab_size, - vocab_size + input_vit.shape[0] * input_vit.shape[1], - device="cuda", - ) - fake_prompt_id = fake_prompt_id.reshape(input_vit.shape[0], - input_vit.shape[1]) - for idx, (i, a, b) in enumerate(img_pos): - input_ids[i][a + 1:b] = fake_prompt_id[idx] - input_ids = input_ids.contiguous().to(torch.int32).cuda() - input_lengths = torch.tensor(input_ids.size(1), - dtype=torch.int32).cuda() - dtype = self.model_config.dtype - prompt_table, tasks, task_vocab_size = self.ptuning_setup( - input_vit, dtype, self.config.hidden_size, None, input_ids) - - output_ids, Qwen_time = self.generate_for_qwenvl( - input_ids, max_new_tokens, prompt_table, tasks, task_vocab_size, - num_beams) - - runtime_rank = tensorrt_llm.mpi_rank() - input_lengths = torch.tensor([input_ids.size(1)], - device="cuda", - dtype=torch.int32) - effective_output_token = 0 - if runtime_rank == 0: - if self.output_csv is None and self.output_npy is None: - for b in range(input_lengths.size(0)): - inputs = input_ids[b] - if content_list is not None: - print(f'Input: "{content_list}"') - print("\n") - if self.num_beams <= 1: - outputs = output_ids[b][0, len(inputs):].tolist() - try: - effective_output_token = (effective_output_token + - outputs.index(151643)) - except: - effective_output_token = 1 - output_text = self.tokenizer.decode( - outputs, skip_special_tokens=True) - print(f'Output: "{output_text}"') - print("\n") - else: - for beam in range(self.num_beams): - outputs = output_ids[b][beam, len(inputs):].tolist() - output_text = self.tokenizer.decode( - outputs, skip_special_tokens=True) - print(f'Output(beam: {beam}): "{output_text}"') - logger.info(f"Input length={input_lengths[b]}") - logger.info(f"Output length={output_ids.shape}") - logger.info(f"TensorRT LLM QWen time: {Qwen_time:3f} sec ") - history.append((query, output_text)) - return output_text - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument("--max_new_tokens", type=int, default=200) - parser.add_argument("--log_level", type=str, default="info") - parser.add_argument( - "--vit_engine_path", - type=str, - default="plan/visual_encoder/visual_encoder_fp16.plan", - ) - parser.add_argument( - "--qwen_engine_dir", - type=str, - default="qwen_outputs", - ) - parser.add_argument( - "--tokenizer_dir", - type=str, - default=".", - help="Directory containing the tokenizer.model.", - ) - parser.add_argument("--input_text", - type=str, - default="Describe the picture") - parser.add_argument( - "--images_path", - nargs="+", - type=json.loads, - default=[{ - "image": "./pics/demo.jpeg" - }], - ) - parser.add_argument( - "--input_tokens", - dest="input_file", - type=str, - help= - "CSV or Numpy file containing tokenized input. Alternative to text input.", - default=None, - ) - parser.add_argument( - "--output_csv", - type=str, - help="CSV file where the tokenized output is stored.", - default=None, - ) - parser.add_argument( - "--output_npy", - type=str, - help="Numpy file where the tokenized output is stored.", - default=None, - ) - parser.add_argument("--num_beams", - type=int, - help="Use beam search if num_beams >1", - default=1) - parser.add_argument("--display", default=False, action='store_true') - parser.add_argument('--port', type=str, default='8006') - parser.add_argument("--local_machine", default=False, action='store_true') - - return parser.parse_args() - - -def vit_process(image_path, vit_engine_path, stream): - img_processor = Preprocss(448) - logger.info(f"Loading engine from {vit_engine_path}") - with open(vit_engine_path, "rb") as f: - engine_buffer = f.read() - logger.info(f"Creating session from engine {vit_engine_path}") - session_vit = Session.from_serialized_engine(engine_buffer) - device = torch.device("cuda") if torch.cuda.is_available() else "cpu" - image_path_list = [] - for item in image_path: - image_path_list.append(next(iter(item.values()))) - images = img_processor.encode(image_path_list).to(device) - batch_size = images.size(0) - images = images.expand(batch_size, -1, -1, -1).contiguous() - visual_inputs = {"input": images.float()} - visual_output_info = session_vit.infer_shapes( - [TensorInfo("input", trt.DataType.FLOAT, images.shape)]) - visual_outputs = { - t.name: - torch.empty(tuple(t.shape), - dtype=trt_dtype_to_torch(t.dtype), - device="cuda") - for t in visual_output_info - } - profiler.start("ViT") - - run_time = 10 - for _ in range(run_time): - ok = session_vit.run(visual_inputs, visual_outputs, stream) - profiler.stop("ViT") - Vit_time = profiler.elapsed_time_in_sec("ViT") / run_time - logger.info(f"TensorRT LLM ViT latency: {Vit_time:3f} sec ") - - assert ok, "Runtime execution failed for vit session" - - image_embeds = visual_outputs["output"] - return image_embeds - - -if __name__ == "__main__": - emit_engine_arch_deprecation("run.py") - args = parse_arguments() - stream = torch.cuda.current_stream().cuda_stream - tensorrt_llm.logger.set_level(args.log_level) - image_embeds = vit_process(args.images_path, args.vit_engine_path, stream) - qinfer = QWenInfer( - args.tokenizer_dir, - args.qwen_engine_dir, - args.log_level, - args.output_csv, - args.output_npy, - args.num_beams, - ) - qinfer.qwen_model_init() - qinfer.qwen_infer( - image_embeds, - args.images_path, - args.input_text, - args.max_new_tokens, - args.num_beams, - history=[], - ) diff --git a/examples/models/core/qwenvl/run_chat.py b/examples/models/core/qwenvl/run_chat.py deleted file mode 100644 index 1f1ba6fb6faf..000000000000 --- a/examples/models/core/qwenvl/run_chat.py +++ /dev/null @@ -1,128 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2023 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. -import re - -# isort: off -import torch -from run import QWenInfer, parse_arguments, vit_process -# isort: on - - -def make_display(port=8006): - import cv2 - import zmq - context = zmq.Context() - socket = context.socket(zmq.REP) - socket.bind(f"tcp://*:{port}") - - def func(image): - data = cv2.imencode(".jpg", image)[1].tobytes() - socket.recv() - socket.send(data) - - return func - - -def show_pic(image_path, port): - import cv2 - image = cv2.imread(image_path) - display_obj = make_display(port) - display_obj(image) - - -def show_pic_local(image_path): - import cv2 - import matplotlib.pyplot as plt - image = cv2.imread(image_path) - image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) - - plt.imshow(image_rgb) - plt.pause(0.1) - - -def cooridinate_extract_show(input, history, tokenizer, local_machine, port): - pattern = r"\((\d+),(\d+)\)" - coordinates = re.findall(pattern, input) - result = "Box({},{})".format(coordinates[0][0], - coordinates[0][1]) - result += ",({},{})".format(coordinates[1][0], coordinates[1][1]) - - image = tokenizer.draw_bbox_on_latest_picture(result, history) - if image: - image.save('1.png') - if local_machine: - show_pic_local('1.png') - else: - show_pic('1.png', port) - else: - print("======No bounding boxes are detected!") - - -def exist_cooridinate(input): - pattern = r"\((\d+),(\d+)\)" - match = re.search(pattern, input) - if match: - return True - else: - return False - - -if __name__ == '__main__': - args = parse_arguments() - stream = torch.cuda.current_stream().cuda_stream - image_embeds = vit_process(args.images_path, args.vit_engine_path, stream) - qinfer = QWenInfer(args.tokenizer_dir, args.qwen_engine_dir, args.log_level, - args.output_csv, args.output_npy, args.num_beams) - qinfer.qwen_model_init() - - run_i = 0 - history = [] - if args.display: - if args.local_machine: - show_pic_local("./pics/demo.jpeg") - else: - show_pic("./pics/demo.jpeg", args.port) - - while True: - input_text = None - try: - input_text = input("Text (or 'q' to quit): ") - except: - continue - - if input_text == "clear history": - history = [] - continue - - if input_text.lower() == 'q': - break - print('\n') - - content_list = args.images_path - content_list.append({'text': input_text}) - - if run_i == 0: - query = qinfer.tokenizer.from_list_format(content_list) - else: - query = input_text - - run_i = run_i + 1 - output_text = qinfer.qwen_infer(image_embeds, None, query, - args.max_new_tokens, args.num_beams, - history) - if args.display: - if exist_cooridinate(output_text): - cooridinate_extract_show(output_text, history, qinfer.tokenizer, - args.local_machine, args.port) diff --git a/examples/models/core/qwenvl/show_pic.py b/examples/models/core/qwenvl/show_pic.py deleted file mode 100644 index f390ec860f58..000000000000 --- a/examples/models/core/qwenvl/show_pic.py +++ /dev/null @@ -1,34 +0,0 @@ -import argparse - -import cv2 -import numpy as np -import zmq - -context = zmq.Context() -socket = context.socket(zmq.REQ) - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--ip', type=str, default='127.0.0.1') - parser.add_argument('--port', type=str, default='8006') - return parser.parse_args() - - -args = parse_arguments() -ip_addr = "tcp://" + args.ip + ":" + args.port -socket.connect(ip_addr) - -while True: - socket.send(b"a") - message = socket.recv() - if len(message) == 1 and message == b'x': - break - image = np.frombuffer(message, dtype=np.uint8) - image = cv2.imdecode(image, 1) - image = cv2.resize(image, dsize=(512, 384)) - cv2.imshow("image", image) - key = cv2.waitKey(1) & 0xFF - - if key == ord('q'): - break diff --git a/examples/models/core/qwenvl/vit_onnx_trt.py b/examples/models/core/qwenvl/vit_onnx_trt.py deleted file mode 100644 index ba21fc93ef0a..000000000000 --- a/examples/models/core/qwenvl/vit_onnx_trt.py +++ /dev/null @@ -1,196 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2023 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. -import argparse -import os -import time -from typing import List - -import requests -import tensorrt as trt -import torch -from PIL import Image -from torchvision import transforms -from torchvision.transforms import InterpolationMode -from transformers import AutoModelForCausalLM - -from tensorrt_llm._utils import release_gc, str_dtype_to_torch - - -class Preprocss: - - def __init__(self, image_size: int): - mean = (0.48145466, 0.4578275, 0.40821073) - std = (0.26862954, 0.26130258, 0.27577711) - self.image_transform = transforms.Compose([ - transforms.Resize((image_size, image_size), - interpolation=InterpolationMode.BICUBIC), - transforms.ToTensor(), - transforms.Normalize(mean=mean, std=std), - ]) - - def encode(self, image_paths: List[str]): - images = [] - for image_path in image_paths: - if image_path.startswith("http://") or image_path.startswith( - "https://"): - image = Image.open(requests.get(image_path, stream=True).raw) - else: - image = Image.open(image_path) - image = image.convert("RGB") - images.append(self.image_transform(image)) - images = torch.stack(images, dim=0) - return images - - -class ONNX_TRT: - - def __init__(self, image_size): - self.image_size = image_size - - def export_onnx(self, onnx_file_path, pretrained_model_path, image_url): - print("Start converting ONNX model!") - image_pre_obj = Preprocss(self.image_size) - torch_dtype = str_dtype_to_torch("float16") - model = AutoModelForCausalLM.from_pretrained( - pretrained_model_path, - device_map="cuda", - dtype=torch_dtype, - fp16=True, - trust_remote_code=True, - ).eval() - device = torch.device("cuda") if torch.cuda.is_available() else "cpu" - image = image_pre_obj.encode(image_url).to(device) - if not os.path.exists("image.pt"): - torch.save(image, "image.pt") - - model_visual = model.transformer.visual - model_visual.eval() - del model # To save GPU memory - - torch.onnx.export( - model_visual, - image.to("cuda"), - onnx_file_path, - opset_version=17, - input_names=["input"], - output_names=["output"], - dynamic_axes={"input": { - 0: "batch" - }}, - # Required for pytorch>=2.9.0 as dynamo becomes the default and introduces bugs as it does not support opset_version=17 natively - dynamo=False) - release_gc() # Further release memory - print( - f"Export to ONNX file successfully! The ONNX file stays in {onnx_file_path}" - ) - - def generate_trt_engine(self, - onnxFile, - planFile, - minBS=1, - optBS=2, - maxBS=4): - print("Start converting TRT engine!") - logger = trt.Logger(trt.Logger.VERBOSE) - builder = trt.Builder(logger) - network = builder.create_network( - 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) - profile = builder.create_optimization_profile() - config = builder.create_builder_config() - config.set_flag(trt.BuilderFlag.FP16) - parser = trt.OnnxParser(network, logger) - - with open(onnxFile, "rb") as model: - if not parser.parse(model.read(), "/".join(onnxFile.split("/"))): - print("Failed parsing %s" % onnxFile) - for error in range(parser.num_errors): - print(parser.get_error(error)) - print("Succeeded parsing %s" % onnxFile) - - nBS = -1 - nMinBS = minBS - nOptBS = optBS - nMaxBS = maxBS - inputT = network.get_input(0) - inputT.shape = [nBS, 3, self.image_size, self.image_size] - profile.set_shape( - inputT.name, - [nMinBS, 3, self.image_size, self.image_size], - [nOptBS, 3, self.image_size, self.image_size], - [nMaxBS, 3, self.image_size, self.image_size], - ) - - config.add_optimization_profile(profile) - - t0 = time.time() - engineString = builder.build_serialized_network(network, config) - t1 = time.time() - if engineString is None: - print("Failed building %s" % planFile) - else: - print("Succeeded building %s in %d s" % (planFile, t1 - t0)) - with open(planFile, "wb") as f: - f.write(engineString) - - -def parse_arguments(): - parser = argparse.ArgumentParser() - # onnx/visual_encoder - parser.add_argument("--onnxFile", - type=str, - default="visual_encoder/visual_encoder.onnx", - help="") - parser.add_argument("--pretrained_model_path", - type=str, - default="Qwen-VL-Chat", - help="") - parser.add_argument( - "--planFile", - type=str, - default="plan/visual_encoder/visual_encoder_fp16.plan", - help="", - ) - parser.add_argument( - "--only_trt", - action="store_true", - help="Run only convert the onnx to TRT engine.", - ) - parser.add_argument("--minBS", type=int, default=1) - parser.add_argument("--optBS", type=int, default=1) - parser.add_argument("--maxBS", type=int, default=4) - parser.add_argument("--image_url", nargs="+", default=["./pics/demo.jpeg"]) - args = parser.parse_args() - return args - - -if __name__ == "__main__": - args = parse_arguments() - onnx_file_dir = os.path.dirname(args.onnxFile) - if not onnx_file_dir == "" and not os.path.exists(onnx_file_dir): - os.makedirs(onnx_file_dir) - plan_file_dir = os.path.dirname(args.planFile) - if not os.path.exists(plan_file_dir): - os.makedirs(plan_file_dir) - - onnx_trt_obj = ONNX_TRT(448) # or ONNX_TRT(config.visual['image_size']) - - if args.only_trt: - onnx_trt_obj.generate_trt_engine(args.onnxFile, args.planFile, - args.minBS, args.optBS, args.maxBS) - else: - onnx_trt_obj.export_onnx(args.onnxFile, args.pretrained_model_path, - args.image_url) - onnx_trt_obj.generate_trt_engine(args.onnxFile, args.planFile, - args.minBS, args.optBS, args.maxBS) diff --git a/examples/ngram/README.md b/examples/ngram/README.md index 1e4dc8792b08..4b5a00ee43ac 100644 --- a/examples/ngram/README.md +++ b/examples/ngram/README.md @@ -1,39 +1,32 @@ # NGram Speculative Decoding -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a model using NGram speculative decoding (supported as `ASSISTED_GENERATION` in transformers and vLLM, source: [GitHub](https://github.com/apoorvumang/prompt-lookup-decoding/tree/main)) in TensorRT LLM on single GPU, or single node multiple GPU. +This document shows how to run a model with NGram speculative decoding +(supported as `ASSISTED_GENERATION` in transformers and vLLM, source: +[GitHub](https://github.com/apoorvumang/prompt-lookup-decoding/tree/main)) +in TensorRT LLM. ## Overview -We provide two styles of workflow to run NGram (named V1 and V2 respectively) now. V1 is in TRT workflow and similar to the Draft-Target-Model workflow, running in orchestrator mode and calling `runner.generate()` multiple times to get outputs, which is more flexible for customizing but slightly more overhead. V2 is in pytorch workflow and similar to the Look-Ahead workflow, running in leader mode and calling `runner.generate()` only one time to get outputs, which provides higher performance but fixed process. - -The NGram has 3 additional hyperparameters that you need to specify to control the process of generation: -- `max_draft_len`: the maximum number of tokens provided as draft tokens in one iteration, which is usually from 4 to 10 in common usage (default value: 4). Empirically, the larger the value is, the higher acceptance rate but higher overhead is expected at the same time, so the right balance based on the models and application scenarios needs to be found. -- `max_matching_ngram_size`: the maximum number of tokens extracted from the tail of the input prompt or generated output as a pattern, which is used to search corresponding draft tokens (default value: 2). Empirically, the larger the value is, the more precise context can be matched from the existed sequence, indicating higher acceptance rate, but the higher probability of miss-match and higher overhead appear, which fall back to normal generation (one token per iteration). -- `device_list`: the index list of device(s) to run the model in V1 workflow. The length of it must be the same as the TP size of the draft model engine. For instances, `device_list=[0]` means using tp_size=1 and GPU 0 for the model, `device_list=[4,5,6,7]` means using tp=4 and GPU from 4 to 7 for the model. This parameter is neddless in V2 workflow. - -+ For example, the process of getting draft tokens using `max_draft_len=2` and `max_matching_ngram_size=4` with a sentence `prefix=[..., t1, t2, t3, t4]` is like below: - -```Python -pattern = prefix[:-2] # pattern=[t3, t4] (length=2) -if pattern in pool and len(pool[pattern]) == 4: # assuming it is {(t3, t4): (t5, t6, t7, t8)} - return pool[pattern] # draft token = [t5, t6, t7, t8] -elif pattern in pool and len(pool[pattern]) == <4: # assuming it is {(t3, t4): (t9, t10, t11)} - return pool[pattern] # draft token = [t9, t10, t11] -pattern = prefix[:-1] # Try shorter pattern if no candidate of length=2 exists, pattern=[t4] (length=1) -if pattern in pool and len(pool[pattern]) == 4: # The same process as above - return pool[pattern] -elif pattern in pool and len(pool[pattern]) == <4: - return pool[pattern] -return None # No any candidate exists -``` +NGram builds a pattern pool from the prompt and previously generated tokens +and proposes draft tokens by matching the tail of the current sequence +against that pool. It has 2 hyperparameters that control the process of +generation: + +- `max_draft_len`: the maximum number of tokens provided as draft tokens in + one iteration, which is usually from 4 to 10 in common usage (default + value: 4). Empirically, the larger the value is, the higher acceptance rate + but higher overhead is expected at the same time, so the right balance + based on the models and application scenarios needs to be found. +- `max_matching_ngram_size`: the maximum number of tokens extracted from the + tail of the input prompt or generated output as a pattern, which is used to + search corresponding draft tokens (default value: 2). Empirically, the + larger the value is, the more precise context can be matched from the + existed sequence, indicating higher acceptance rate, but the higher + probability of miss-match and higher overhead appear, which fall back to + normal generation (one token per iteration). ## Support Matrix + * GPU Compute Capability >= 8.0 (Ampere or newer) * FP16 / BF16 / FP8 * Paged KV Cache @@ -41,59 +34,6 @@ return None # No any candidate exists ## Usage -### V1 workflow - -+ We use an open-source `llama-v2-13B` models in this example. -+ `--use_paged_context_fmha=enable` must be specified since we need KVcache reuse in this approach. -+ `--speculative_decoding_mode=draft_tokens_external` must be specified. -+ `--max_draft_len` must be specified as the length maximum of the draft tokens. -+ `--ngram_config` is corresponding configuration of NGram, we can see its usage in [util.py](../util.py). - + As an example, `[10,2,[0]]` means `max_draft_len=10`, `max_matching_ngram_size=2`, and device of target model is `GPU0`. -+ `--kv_cache_enable_block_reuse` must be specified for this approach. -+ Only CPP session is supported, so `--use_py_session` must not be specified. -+ `--num_beams` can not be specified as larger than 1 since beam search is not supported in this approach yet. - -```bash -# Build engine -python3 examples/models/core/llama/convert_checkpoint.py \ - --model_dir \ - --output_dir ./ckpt-target \ - --dtype float16 - -trtllm-build \ - --checkpoint_dir ./ckpt-target \ - --output_dir ./target-engine \ - --gemm_plugin float16 \ - --use_paged_context_fmha enable \ - --speculative_decoding_mode draft_tokens_external \ - --max_draft_len 10 \ - --max_batch_size 4 \ - --max_input_len 3200 \ - --max_seq_len 4800 - -# Run decoding -python3 examples/run.py \ - --tokenizer_dir \ - --engine_dir ./target-engine \ - --ngram_config "[10,2,[0]]" \ - --max_output_len 256 \ - --kv_cache_enable_block_reuse \ - --input_text "How does Draft-Sampling work?" - -# Run summarization tasks -python examples/summarize.py \ - --test_hf \ - --test_trt_llm \ - --check_accuracy \ - --hf_model_dir \ - --engine_dir ./target-engine \ - --batch_size 1 \ - --ngram_config "[10,2,[0]]" \ - --kv_cache_enable_block_reuse -``` - -### V2 workflow - ```bash python3 examples/llm-api/quickstart_advanced.py \ --spec_decode_max_draft_len 4 \ @@ -101,3 +41,8 @@ python3 examples/llm-api/quickstart_advanced.py \ --disable_overlap_scheduler \ --disable_kv_cache_reuse ``` + +With the LLM API, configure NGram through `NGramDecodingConfig` +(`speculative_config`). See the +[speculative decoding documentation](https://nvidia.github.io/TensorRT-LLM/features/speculative-decoding.html) +for details. diff --git a/examples/ngram/run_dtm_ngram.py b/examples/ngram/run_dtm_ngram.py deleted file mode 100644 index d0cd8687ef86..000000000000 --- a/examples/ngram/run_dtm_ngram.py +++ /dev/null @@ -1,382 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. - -import ast - -import numpy as np -import torch -from ordered_set import OrderedSet - -from tensorrt_llm.logger import logger -from tensorrt_llm.runtime import ModelRunnerCpp - - -class NgramPool: # Ngrams pool for Ngram - - def __init__( - self, - input_batch_size: int, - max_draft_len: int, - max_matching_ngram_size: int, - end_id: int, - max_seq_len: list[int], - is_keep_all: bool = True, - is_use_oldest: bool = True, - ): - self.input_batch_size = input_batch_size - self.max_draft_len = max_draft_len - self.max_matching_ngram_size = max_matching_ngram_size - self.end_id = end_id - self.max_seq_len = max_seq_len - self.is_keep_all = is_keep_all - self.is_use_oldest = is_use_oldest - self.pool = [{} for _ in range(input_batch_size)] - self.start_index = [0 for _ in range(input_batch_size)] - - assert self.max_draft_len > 0, f"max_draft_len must be greater than 0, but got {self.max_draft_len}" - assert self.max_matching_ngram_size > 0, f"max_matching_ngram_size must be greater than 0, but got {self.max_matching_ngram_size}" - - def print_pool(self): - """ - For debug - """ - logger.info(f"Batch size = {self.input_batch_size}") - for i, map in enumerate(self.pool): - logger.info(f"Slot {i}, size = {len(map)}") - for key, values in map.items(): - logger.info(f" {key}->{values}") - - def get_draft_tokens(self, prefix: list[torch.Tensor], - batch_slot: list[int]): - """ - Get draft tokens from a batch of requests - modified from `transformers/generation/candidate_generator.py` - """ - batch_size = len(prefix) - prefix_len = [len(prefix[bi]) for bi in range(batch_size)] - draft_tokens = [] # `logits` is useless yet - for bi in range(batch_size): - gbi = batch_slot[bi] # Global index in the input batch - chosen_ids = [self.end_id] - # Skip search if prefix is length of `max_length - 1` - if prefix_len[bi] >= self.max_seq_len[gbi] - 1: - draft_tokens.append(chosen_ids) - continue - - # Update pool - sequence = prefix[bi][self.start_index[gbi]:].tolist() - for size in range( - min(self.max_matching_ngram_size, prefix_len[bi] - 1), 0, - -1): - # Find each possible key-value combination, and use tuple for hash - for l in range(len(sequence) - size): - r = min(l + size + self.max_draft_len, len(sequence)) - key = tuple(sequence[l:l + size]) - value = tuple(sequence[l + size:r]) - if key not in self.pool[gbi] or not self.is_keep_all or \ - len(self.pool[gbi][key][0]) < self.max_draft_len: - # Update the value if - # 1. the key does not exist - # 2. we only keep the newest one value for each key (MRU) - # 3. the length of the value saved before is less than `max_draft_len` - self.pool[gbi][key] = OrderedSet((value, )) - elif value not in self.pool[gbi][key]: - # Extend the value if the key is already existed but count of values is not enough - self.pool[gbi][key].add(value) - - # Find match - for size in range( - min(self.max_matching_ngram_size, prefix_len[bi] - 1), 0, - -1): - pattern = tuple(prefix[bi][-size:].tolist()) - if pattern not in self.pool[gbi]: - continue - if self.is_use_oldest: - # Always choose the oldest match, aligned with HF - chosen_ids = self.pool[gbi][pattern][0] - else: - # Always choose the newest match - chosen_ids = self.pool[gbi][pattern][-1] - break - draft_tokens.append(chosen_ids) - self.start_index[gbi] = max( - 0, prefix_len[bi] - - (self.max_draft_len + self.max_matching_ngram_size - 1)) - - return draft_tokens, None - - -def run_dtm_ngram(batch_input_ids, - args, - runtime_rank, - end_id, - pad_id, - stop_words_list, - bad_words_list, - vocab_size, - *, - target_runner=None): - # `dtm` for Draft-Target-Model, `ngram` for NGram - is_dtm = (args.draft_target_model_config is not None) - is_ngram = (args.ngram_config is not None) - assert is_dtm ^ is_ngram, "`--draft_target_model_config` and `--ngram_config` can not be specified at the same time." - if is_dtm: - assert args.draft_engine_dir is not None, "`--draft_engine_dir` must be specified in Draft-Target-Model." - draft_len, draft_device_list, target_device_list, use_logits = ast.literal_eval( - args.draft_target_model_config) - logger.info(f"Using Draft-Target-Model speculative decoding") - logger.info(f"draft_len: {draft_len}") - logger.info(f"Device(s) for draft model: {draft_device_list}") - logger.info(f"Device(s) for target model: {target_device_list}") - logger.info(f"Use logits to accept tokens: {use_logits}") - if is_ngram: - logger.info(f"Using NGram speculative decoding V1 workflow") - max_draft_len, max_matching_ngram_size, target_device_list = ast.literal_eval( - args.ngram_config) - logger.info(f"max_draft_len: {max_draft_len}") - logger.info(f"max_matching_ngram_size: {max_matching_ngram_size}") - logger.info(f"Device(s) for the model: {target_device_list}") - use_logits = False # `logits` is useless in this approach yet - - # Variables keeping constant during decoding - input_batch_size = len(batch_input_ids) # Note as `BS` - beam_width = args.num_beams # Note as `BW` - is_compute_acceptance_ratio = logger.level == 'verbose' # Only for verbose - input_len = [len(p) for p in batch_input_ids] - max_seq_len = [i + args.max_output_len for i in input_len] - # Variables changing during decoding - n_iteration = 0 - prefix = batch_input_ids # Input for each iteration - batch_slot = list(range(input_batch_size)) # Index of requests - if is_compute_acceptance_ratio: - n_draft_token = [0 for _ in range(input_batch_size)] - n_accept_token = [0 for _ in range(input_batch_size)] - - if is_ngram: - ngram_pool = NgramPool(input_batch_size, max_draft_len, - max_matching_ngram_size, end_id, max_seq_len) - - # Repack the output like the output of function `generate` - outputs = {} - outputs["output_ids"] = torch.full( - [input_batch_size, beam_width, - max(max_seq_len)], - end_id, - dtype=torch.int32) - for bi in range(input_batch_size): - outputs["output_ids"][bi, :, :input_len[bi]] = batch_input_ids[bi] - outputs["sequence_lengths"] = torch.full([input_batch_size, beam_width], - 0, - dtype=torch.int32) - outputs["context_logits"] = None - outputs["generation_logits"] = torch.full( - [input_batch_size, beam_width, - max(max_seq_len), vocab_size], - 0, - dtype=torch.float16) - outputs['cum_log_probs'] = None - outputs['log_probs'] = None - - # Model runner - common_runner_kwargs = dict( - lora_dir=args.lora_dir, - rank=runtime_rank, - debug_mode=args.debug_mode, - lora_ckpt_source=args.lora_ckpt_source, - gpu_weights_percent=args.gpu_weights_percent, - max_output_len=args.max_output_len, - is_enc_dec=False, - max_batch_size=input_batch_size, - max_input_len=max(input_len) + args.max_output_len, - max_beam_width=beam_width, - max_attention_window_size=args.max_attention_window_size, - sink_token_length=args.sink_token_length, - max_tokens_in_paged_kv_cache=args.max_tokens_in_paged_kv_cache, - kv_cache_enable_block_reuse=args.kv_cache_enable_block_reuse, - kv_cache_free_gpu_memory_fraction=args. - kv_cache_free_gpu_memory_fraction, - cross_kv_cache_fraction=None, - enable_chunked_context=args.enable_chunked_context, - multi_block_mode=args.multi_block_mode, - cuda_graph_mode=args.cuda_graph_mode, - enable_context_fmha_fp32_acc=args.enable_context_fmha_fp32_acc, - is_orchestrator_mode=True, - ) - - if is_dtm: - draft_runner_kwargs = common_runner_kwargs.copy() - draft_runner_kwargs.update(engine_dir=args.draft_engine_dir, - device_ids=draft_device_list) - draft_runner = ModelRunnerCpp.from_dir(**draft_runner_kwargs) - - if target_runner is None: # Skip this constructor if we have prepared the runner before - target_runner_kwargs = common_runner_kwargs.copy() - target_runner_kwargs.update(engine_dir=args.engine_dir, - device_ids=target_device_list) - target_runner = ModelRunnerCpp.from_dir(**target_runner_kwargs) - - if is_dtm and use_logits: - assert draft_runner.gather_generation_logits and target_runner.gather_generation_logits, "`--gather_generation_logits` must be specified while building draft/target models for using logits to accept" - - common_generaion_kwargs = dict( - max_attention_window_size=args.max_attention_window_size, - sink_token_length=args.sink_token_length, - end_id=end_id, - pad_id=pad_id, - temperature=args.temperature, - top_k=args.top_k, - top_p=args.top_p, - num_beams=beam_width, - num_return_sequences=args.num_return_sequences, - length_penalty=args.length_penalty, - early_stopping=args.early_stopping, - beam_width_array=None, - repetition_penalty=args.repetition_penalty, - presence_penalty=args.presence_penalty, - frequency_penalty=args.frequency_penalty, - min_p=args.min_p, - stop_words_list=stop_words_list, - bad_words_list=bad_words_list, - random_seed=args.random_seed, - lora_uids=args.lora_task_uids, - prompt_table=args.prompt_table_path, - prompt_tasks=args.prompt_tasks, - streaming=False, - output_sequence_lengths=True, - no_repeat_ngram_size=args.no_repeat_ngram_size, - return_dict=True, - return_all_generated_tokens=args.return_all_generated_tokens, - ) - - while True: - n_iteration += 1 - # Dynamic batch_size, decreases if some requests finish - batch_size = len(prefix) - prefix_len = [len(prefix[i]) for i in range(batch_size)] - # Get draft tokens - # `d_*` means variables from draft - # `d_seq_len` includes input part, but `d_len` doesn't - if is_dtm: - draft_generation_kwargs = common_generaion_kwargs.copy() - draft_generation_kwargs.update( - batch_input_ids=prefix, - max_new_tokens=draft_len, - streaming=False, - output_sequence_lengths=True, - return_dict=True, - ) - draft = draft_runner.generate(**draft_generation_kwargs) - torch.cuda.synchronize() - - # draft["output_ids"].shape -> [BS, BW, maxSL] - # draft["sequence_lengths"].shape -> [BS, BW] - # draft["generation_logits"].shape -> [BS, BW, draft_len, vocab_size] - d_ids = [[end_id]] * batch_size - d_logits = [None] * batch_size if use_logits else None - d_seq_len = draft["sequence_lengths"][:, 0].tolist() - d_len = [d_seq_len[bi] - prefix_len[bi] for bi in range(batch_size)] - for bi in range(batch_size): - l, r = prefix_len[bi], d_seq_len[bi] - if l >= r: # No useful draft tokens - continue - d_ids[bi] = draft["output_ids"][bi, 0, l:r].tolist() - if use_logits: - d_logits[bi] = draft["generation_logits"][bi, 0, - -d_len[bi]:, :] - if is_ngram: - d_ids, d_logits = ngram_pool.get_draft_tokens(prefix, batch_slot) - d_len = [len(i) for i in d_ids] - - # Run target model - # `t_*` means variables from target model - # `t_seq_len` and `t_seq_ids` include input part, but `t_len` or `t_ids` don't - target_generation_kwargs = common_generaion_kwargs.copy() - target_generation_kwargs.update(batch_input_ids=prefix, - draft_tokens_list=d_ids, - draft_logits_list=d_logits) - if is_dtm: - max_new_tokens = draft_len + 1 - if is_ngram: - max_new_tokens = max_draft_len + 1 - target_generation_kwargs.update(max_new_tokens=max_new_tokens) - target = target_runner.generate(**target_generation_kwargs) - torch.cuda.synchronize() - - t_ids = [None] * batch_size - t_seq_ids = [None] * batch_size - t_seq_len = target["sequence_lengths"][:, 0].tolist() - t_len = [t_seq_len[bi] - prefix_len[bi] for bi in range(batch_size)] - - # Update output and tokens for next iteration - for bi in range(batch_size): - gbi = batch_slot[bi] # Global index in the input batch - l = prefix_len[bi] - r = min(t_seq_len[bi], max_seq_len[gbi]) - t_ids[bi] = target["output_ids"][bi, 0, l:r].tolist() - t_seq_ids[bi] = target["output_ids"][bi, 0, :r] - outputs["output_ids"][gbi, 0, l:r] = torch.IntTensor(t_ids[bi]) - outputs["sequence_lengths"][gbi, 0] = r - if use_logits: - outputs["generation_logits"][gbi, 0, (l - input_len[bi]):(r - input_len[bi])] = \ - target["generation_logits"][bi][0,:(r-l)].detach().cpu() - if is_compute_acceptance_ratio: - n_draft_token[gbi] += d_len[bi] - length = min(d_len[bi], t_len[bi], - max_seq_len[gbi] - prefix_len[bi]) - res = [d_ids[bi][i] == t_ids[bi][i] for i in range(length)] - n_accept_token[gbi] += \ - ((~torch.BoolTensor(res)).cumsum(axis=-1) < 1).sum() - - # Yield output if using streaming - if args.streaming and not n_iteration % args.streaming_interval: - yield outputs - - # Evaluate stop criteria and prepare inputs for next iteration - prefix_next = [] - batch_slot_next = [] - for bi in range(batch_size): - gbi = batch_slot[bi] # Global index in the input batch - # Stop due to output length - if len(t_seq_ids[bi]) >= max_seq_len[gbi]: - continue # No need to update for the stopped requests - # Stop due to the same output. Normally target should return 1 more token. - # if (d_ids is not None and np.array_equal(d_ids[bi], t_ids[bi])): - # continue - # Stop due to no change (hit early stopping) - if np.array_equal(t_seq_ids[bi].cpu().numpy(), - prefix[bi].cpu().numpy()): - continue - # Stop due to end words - if end_id in t_seq_ids[bi][prefix_len[bi]:]: - continue - # TODO: Check bad words and stop words criteria - prefix_next.append(t_seq_ids[bi]) - batch_slot_next.append(gbi) - prefix = prefix_next - batch_slot = batch_slot_next - if len(prefix) == 0: # Leave while loop if no request remained - break - - if is_compute_acceptance_ratio: - logger.debug(f"Count of iteration(s): {n_iteration}") - logger.debug(f"Acceptance ratio:") - for i, (a, d) in enumerate(zip(n_accept_token, n_draft_token)): - logger.debug(f"Request {i}: {a / d * 100 :6.2f}%") - - # Return runner in No-Streaming mode - if args.streaming: - yield outputs - else: - yield outputs, target_runner diff --git a/examples/openai_triton/README.md b/examples/openai_triton/README.md deleted file mode 100644 index b5f39d105974..000000000000 --- a/examples/openai_triton/README.md +++ /dev/null @@ -1,7 +0,0 @@ -# Integration for OpenAI Triton - -The typical approach to integrate a kernel into TensorRT LLM is to create TensorRT plugins. -Specially for integrating OpenAI Triton kernels, there are two methods: - -1. Creating TensorRT plugin manually, you can refer to [manual plugin example](./manual_plugin/) for details, -2. Generate the TensorRT plugins automatically, please refer to [automatic plugin example](./plugin_autogen/) for details. diff --git a/examples/openai_triton/manual_plugin/CMakeLists.txt b/examples/openai_triton/manual_plugin/CMakeLists.txt deleted file mode 100644 index bec14231511e..000000000000 --- a/examples/openai_triton/manual_plugin/CMakeLists.txt +++ /dev/null @@ -1,113 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -cmake_minimum_required(VERSION 3.1) - -# Enable C++ -set(CMAKE_CXX_STANDARD 17) -set(CMAKE_CXX_STANDARD_REQUIRED TRUE) - -# Define project name -set(TARGET_NAME trt_llm_custom_plugins) -project(${TARGET_NAME}) - -set(CMAKE_VERBOSE_MAKEFILE 1) - -# Compile options -set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -Wall -pthread ") -set(CMAKE_C_FLAGS_DEBUG "-g -O0") -set(CMAKE_C_FLAGS_RELEASE "-O2") -set(CMAKE_CXX_FLAGS "${CMAKE_C_FLAGS} -lstdc++") -set(CMAKE_CXX_FLAGS_DEBUG ${CMAKE_C_FLAGS_DEBUG}) -set(CMAKE_CXX_FLAGS_RELEASE ${CMAKE_C_FLAGS_RELEASE}) - -set(CMAKE_BUILD_TYPE release) - -find_package(CUDA REQUIRED) -message(STATUS "CUDA library status:") -message(STATUS " config: ${CUDA_DIR}") -message(STATUS " version: ${CUDA_VERSION}") -message(STATUS " libraries: ${CUDA_LIBRARIES}") -message(STATUS " include path: ${CUDA_INCLUDE_DIRS}") - -if(NOT DEFINED TRT_INCLUDE_DIR) - set(TRT_INCLUDE_DIR "/usr/local/tensorrt/include") - if(NOT EXISTS ${TRT_INCLUDE_DIR}) - # In case of TensorRT installed from a deb package. - set(TRT_INCLUDE_DIR "/usr/include/x86_64-linux-gnu") - endif() -endif() -message(STATUS "tensorrt include path: ${TRT_INCLUDE_DIR}") -if(DEFINED TRT_LLM_INCLUDE_DIR) - message( - STATUS "openai_triton/manual_plugin example has been self-contained " - "and TRT_LLM_INCLUDE_DIR is now unnecessary to specify the path of " - "C++ runtime source files.") -endif() - -if(NOT DEFINED TRT_LIB_DIR) - set(TRT_LIB_DIR "/usr/local/tensorrt/lib") - if(NOT EXISTS ${TRT_INCLUDE_DIR}) - # In case of TensorRT installed from a deb package. - set(TRT_LIB_DIR "/lib/${CMAKE_SYSTEM_PROCESSOR}-linux-gnu") - endif() -endif() -find_library( - TRT_LIB_PATH nvinfer - HINTS ${TRT_LIB_DIR} - NO_DEFAULT_PATH) -find_library(TRT_LIB_PATH nvinfer REQUIRED) -message(STATUS "TRT_LIB_DIR: ${TRT_LIB_DIR}") -message(STATUS "Found nvinfer library: ${TRT_LIB_PATH}") - -if(NOT DEFINED TRT_LLM_LIB_DIR) - # Find at tensorrt_llm/libs. - execute_process( - COMMAND ${CMAKE_COMMAND} -E env "PYTHONPATH=${PYTHONPATH}" "python" "-c" - "import tensorrt_llm; print(f'{tensorrt_llm.__path__[0]}/libs')" - OUTPUT_VARIABLE TRT_LLM_LIB_DIR - OUTPUT_STRIP_TRAILING_WHITESPACE) - # Find /tensorrt_llm/libs. - list(APPEND TRT_LLM_LIB_DIR "../../../tensorrt_llm/libs") -endif() -find_library(TRT_LLM_LIB_PATH nvinfer_plugin_tensorrt_llm - HINTS ${TRT_LLM_LIB_DIR} NO_DEFAULT_PATH) -find_library(TRT_LLM_LIB_PATH nvinfer_plugin_tensorrt_llm REQUIRED) -message(STATUS "Found nvinfer_plugin_tensorrt_llm library: ${TRT_LLM_LIB_PATH}") - -find_library(TRT_LLM_COMMON_LIB_PATH th_common HINTS ${TRT_LLM_LIB_DIR} - NO_DEFAULT_PATH) -find_library(TRT_LLM_COMMON_LIB_PATH th_common REQUIRED) -message(STATUS "Found th_common library: ${TRT_LLM_COMMON_LIB_PATH}") - -# Declare the target library. -add_library( - ${TARGET_NAME} SHARED - tritonPlugins.cpp - TritonFlashAttentionPlugin.cpp - aot/fmha_kernel_fp16.c - aot/fmha_kernel_fp32.c - aot/fp16/fmha_kernel_d64_fp16.fbf0f274_0d1d2d3d4d5d6789.c - aot/fp32/fmha_kernel_d64_fp32.f30323ef_0d1d2d3d4d5d6789.c) - -target_link_libraries( - ${TARGET_NAME} PUBLIC cuda ${CUDA_LIBRARIES} ${TRT_LLM_LIB_PATH} - ${TRT_LLM_COMMON_LIB_PATH} ${TRT_LIB_PATH}) - -if(NOT MSVC) - set_property(TARGET ${TARGET_NAME} PROPERTY LINK_FLAGS "-Wl,--no-undefined") -endif() - -target_include_directories(${TARGET_NAME} PUBLIC /usr/local/cuda/include) -target_include_directories(${TARGET_NAME} PUBLIC ${TRT_INCLUDE_DIR}) diff --git a/examples/openai_triton/manual_plugin/README.md b/examples/openai_triton/manual_plugin/README.md deleted file mode 100644 index 5c8b5d481d52..000000000000 --- a/examples/openai_triton/manual_plugin/README.md +++ /dev/null @@ -1,177 +0,0 @@ -# OpenAI Triton Plugin in TensorRT-LLM - -This document describes how to build and run a custom plugin leveraging [OpenAI Triton](https://github.com/openai/triton) in TensorRT-LLM. -The workflow can be summarized as follows. - 1. Implement a kernel using Triton in Python. - 2. Compile that kernel using Triton AoT (Ahead-of-Time) compilation tool to generate C files. - 3. Implement a custom TensorRT LLM plugin to execute the compiled kernel. - 4. Build the TensorRT engine. - 5. It is ready to be executed by TensorRT. - -In this example, we show how to create a TensorRT LLM plugin to wrap a [Fused Attention]((fmha_triton.py)) kernel implemented in OpenAI Triton. -As a prerequisite, it is necessary to have the TensorRT LLM C++ runtime library. -The instructions to build that library can be found [here](../../README.md#build-from-source). - -## 1. Triton AoT Preparation - -OpenAI Triton offers an Ahead-of-Time (AoT) compilation tool to generate C files that wrap compiled GPU kernel. -To use the AoT feature, you need a Triton version posterior to the [d0c35b3](https://github.com/openai/triton/commit/d0c35b3b7d6badf0c0d56a821dddab7ace73b4de) commit -and this example has been tested on the [b43c28f](https://github.com/openai/triton/tree/b43c28fdd7a2f95b2e87180cba5d984732120d5c) commit. -```bash -git clone https://github.com/openai/triton -cd triton/python/ -git checkout d4644d6cb3ae674e1f15932cac1f28104795744f -pip install cmake && pip install . -cd - -``` - -For AoT compilation, it is necessary to provide a kernel signature and specify the values of `tl.constexpr` parameters in a comma-separated format. -Details can be found in the [compile.py](https://github.com/openai/triton/blob/main/python/triton/tools/compile.py) file in the Triton project. - -Here are examples of kernel AOT compilations for the [Fused Attention](fmha_triton.py) kernel. -```bash -# Kernel for data type=float16, BLOCK_M=128, BLOCK_DMODEL=64, BLOCK_N=128 -export TRITON_ROOT=$(pip show triton | grep Location | cut -d' ' -f2) -rm -rf aot -mkdir -p aot/fp16 -python ${TRITON_ROOT}/triton/tools/compile.py \ - fmha_triton.py \ - -n fused_attention_kernel \ - -o aot/fp16/fmha_kernel_d64_fp16 \ - --out-name fmha_d64_fp16 \ - -w 4 \ - -ns 2 \ - -s "*fp16:16, *fp32:16, *fp32:16, *fp16:16, *fp16:16, *fp16:16, fp32, i32, i32, i32, 128, 64, 128" \ - -g "(seq_len + 127) / 128, batch_size * num_heads, 1" -# Kernel for data type=float32, BLOCK_M=64, BLOCK_DMODEL=64, BLOCK_N=64 -mkdir -p aot/fp32 -python ${TRITON_ROOT}/triton/tools/compile.py \ - fmha_triton.py \ - -n fused_attention_kernel \ - -o aot/fp32/fmha_kernel_d64_fp32 \ - --out-name fmha_d64_fp32 \ - -w 4 \ - -ns 2 \ - -s "*fp32:16, *fp32:16, *fp32:16, *fp32:16, *fp32:16, *fp32:16, fp32, i32, i32, i32, 64, 64, 64" \ - -g "(seq_len + 63) / 64, batch_size * num_heads, 1" - -# Link generated headers and create dispatchers. -python ${TRITON_ROOT}/triton/tools/link.py aot/fp16/*.h -o aot/fmha_kernel_fp16 -python ${TRITON_ROOT}/triton/tools/link.py aot/fp32/*.h -o aot/fmha_kernel_fp32 -``` -The tool will generate .c and .h files to launch the GPU kernel. -Note that it is necessary to specify the kernel name using the --out-name option, it allows to define dispatcher names for the different data types. -The above invocations will generate `aot/fmha_kernel_{fp16|fp32}.{c|h}` files that contain three functions: - - the `load_fmha_d64_{fp16|fp32}` function to load the code of the GPU kernel, - - the `fmha_d64_{fp16|fp32}` function to launch the kernel, - - the `unload_fmha_d64_{fp16|fp32}` function to unload the GPU kernel. - -If GPU resources are limited, it is recommended to adjust the number of stages or warps accordingly. For example, on the V100, the aforementioned arguments might fail due to insufficient shared memory of the GPU. This can be mitigated by reducing the number of stages by one, using `-ns 1`. - - -## 2. Implement a Custom TensorRT Plugin - -This section describes how to implement a custom plugin for TensorRT LLM to execute the Triton kernel created in the previous section. -We provide an example of plugin implementation. - - TritonFlashAttentionPlugin([.cpp](TritonFlashAttentionPlugin.cpp), [.h](TritonFlashAttentionPlugin.h)): TensorRT plugin. - - [plugin.py](plugin.py): Python wrapper. - -`TritonFlashAttentionPlugin` is a TensorRT plugin that integrates a Triton kernel generated with the AoT compiler. -The `initialize` and `terminate` functions show how to initialize and terminate the TensorRT plugin. -The `enqueue` member function shows how to call the generated Triton kernel on the GPU. -Note that the name of the Triton kernel depends on the function's signature, meaning that different types or specialization leads a different kernel name. -Thus, if you change an option during AoT compilation like `-s `, you also have to update file names in CMakeLists.txt in order to match the names generated by the AoT compiler. - -To build a shared library for the custom Triton plugin, run: -```bash -mkdir -p build && cd build -cmake .. && make -cd .. -``` -As mentioned in the previous section, it is necessary to have the TensorRT LLM C++ runtime library. -If you want to specify the library paths, run: -```bash -cmake -DTRT_LIB_DIR= -DTRT_INCLUDE_DIR= -DTRT_LLM_LIB_DIR= .. -``` -If the build is successful, you should be able to find a shared library for the custom plugin at `build/libtrt_llm_custom_plugins.so`. - -A Python wrapper of the Fused Multihead Attention (FMHA) operator and the corresponding TensorRT LLM layer are implemented in [plugin.py](plugin.py). -It is similar to other TensorRT LLM operators and layers implemented in [functional.py](../../tensorrt_llm/functional.py) and [layers](../../tensorrt_llm/layers), respectively. -That FMHA operator uses the custom plugin that wraps the functions generated from the Triton kernel. - -## 3. Build and Run the TensorRT Engine - -We are now ready to build and run the TensorRT engine that uses the Triton kernel. -Here are the two commands to build and run the engine: -```bash -python build.py --num_heads 32 --head_size 64 --max_batch_size 8 --max_seq_len 512 --dtype float16 -python run.py --num_heads 32 --head_size 64 --batch_size 8 --seq_len 512 --log_level verbose --benchmark -``` - -## 4. Known Issues - -### 1. A generated dispatcher might not execute a kernel without raising an error due to a missing branch. - -The kernel dispatcher written by `link.py` has a missing branch, which can result in returning without executing a kernel. -For instance, in our example, the generated dispatcher looks like this: -```c++ -CUresult fmha_d64_fp16(CUstream stream, unsigned int gX, unsigned int gY, unsigned int gZ, CUdeviceptr Out, CUdeviceptr L, CUdeviceptr M, CUdeviceptr Q, CUdeviceptr K, CUdeviceptr V, float sm_scale, int32_t seq_len){ - if ((Out % 16 == 0) && (L % 16 == 0) && (M % 16 == 0) && (Q % 16 == 0) && (K % 16 == 0) && (V % 16 == 0)) - return fmha_d64_fp16_0eb6b090_0d1d2d3d4d5d67(stream, gX, gY, gZ, Out, L, M, Q, K, V, sm_scale, seq_len); -} -``` -It is recommended to manually update the generated functions by `link.py` to return a proper error for proper error handling. - - -### 2. The shared memory required by a generated kernel may exceed the hardware limitation. - -The AoT compiler does not verify the limitations of shared memory size during compilation time, which could potentially lead to the out-of-resource errors during runtime. -It would be helpful to verify if the requirement of the dynamic shared memory size in a generated kernel exceeds the hardware limitation. -You can find the number at the line of `cuLaunchKernel` call in the generated `.c` file. -For instance, the shared memory size is 114690 bytes in our example. -```c++ -CUresult fmha_d64_fp16_0eb6b090_0d1d2d3d4d5d67(CUstream stream, unsigned int gX, unsigned int gY, unsigned int gZ, CUdeviceptr Out, CUdeviceptr L, CUdeviceptr M, CUdeviceptr Q, CUdeviceptr K, CUdeviceptr V, float sm_scale, int32_t seq_len) { - if (fmha_d64_fp16_0eb6b090_0d1d2d3d4d5d67_func == NULL) - load_fmha_d64_fp16_0eb6b090_0d1d2d3d4d5d67(); - void *args[8] = { &Out, &L, &M, &Q, &K, &V, &sm_scale, &seq_len }; - // TODO: shared memory - if(gX * gY * gZ > 0) - return cuLaunchKernel(fmha_d64_fp16_0eb6b090_0d1d2d3d4d5d67_func, gX, gY, gZ, 4 * 32, 1, 1, 114690, stream, args, NULL); -} -``` -It may be resolved by reduing the block size. - - -### 3. AttributeError: module 'triton' has no attribute 'jit' - -This problem may arise if Triton is installed in editable mode. To resolve this issue, please install Triton using the non-editable mode. Refer https://github.com/openai/triton/issues/1693. - -### 4. Unload the same module more than once while building the engine -When the plugin is used more than once within a model, the function cuModuleUnload() will be invoked multiple times during the engine building stage. Related code is generated by Openai Triton and can be found in the folder `examples/openai_triton/manual_plugin/aot/`. One example is: - -```c++ -void unload_fmha_d64_fp32_f30323ef_0d1d2d3d4d5d6789(void) { - CUDA_CHECK(cuModuleUnload(fmha_d64_fp32_f30323ef_0d1d2d3d4d5d6789_mod)); -} -``` - -As the generated code didn't check the value of the module object, this function might unload the same module multiple times, which will cause an error as follows: - -``` -Triton Error [CUDA]: invalid resource handle\n/opt/rapids/src/cudf/cpp/build/_deps/arrow-src/cpp/src/arrow/filesystem/s3fs.cc:2904:  arrow::fs::FinalizeS3 was not called even though S3 was initialized.  This could lead to a segmentation fault at exit -``` - -The error message is ambiguous. If we use compute-sanitizer to help debug, we can get the following information: -``` -========= Program hit CUDA_ERROR_INVALID_HANDLE (error 400) due to "invalid resource handle" on CUDA API call to cuModuleUnload. -``` - -So we need to modify the above generated code as follows to avoid the above error. -```c++ -void unload_fmha_d64_fp32_f30323ef_0d1d2d3d4d5d6789(void) { - if(fmha_d64_fp32_f30323ef_0d1d2d3d4d5d6789_mod){ - CUDA_CHECK(cuModuleUnload(fmha_d64_fp32_f30323ef_0d1d2d3d4d5d6789_mod)); - } - fmha_d64_fp32_f30323ef_0d1d2d3d4d5d6789_mod=NULL; -} -``` diff --git a/examples/openai_triton/manual_plugin/TritonFlashAttentionPlugin.cpp b/examples/openai_triton/manual_plugin/TritonFlashAttentionPlugin.cpp deleted file mode 100644 index 198d8be1ca16..000000000000 --- a/examples/openai_triton/manual_plugin/TritonFlashAttentionPlugin.cpp +++ /dev/null @@ -1,386 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-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. - */ -#include "TritonFlashAttentionPlugin.h" - -// Import a generated header to use generated triton kernels. -extern "C" -{ -#include "aot/fmha_kernel_fp16.h" -#include "aot/fmha_kernel_fp32.h" -} - -#include -#include -#include -#include - -using namespace nvinfer1; -using openai_triton::plugin::TritonFlashAttentionPluginCreator; -using openai_triton::plugin::TritonFlashAttentionPlugin; - -static char const* TRITON_FLASH_ATTENTION_PLUGIN_VERSION{"1"}; -static char const* TRITON_FLASH_ATTENTION_PLUGIN_NAME{"TritonFlashAttention"}; -PluginFieldCollection TritonFlashAttentionPluginCreator::mFC{}; -std::vector TritonFlashAttentionPluginCreator::mPluginAttributes; - -namespace openai_triton::plugin -{ - -// Write values into buffer -template -void writeArg(char*& buffer, T const& val) -{ - std::memcpy(buffer, &val, sizeof(T)); - buffer += sizeof(T); -} - -// Read values from buffer -template -void readArg(char const*& buffer, T& val) -{ - std::memcpy(&val, buffer, sizeof(T)); - buffer += sizeof(T); -} - -std::uintptr_t constexpr kCudaMemAlign = 128; - -int8_t* nextWorkspacePtr(int8_t* ptr, uintptr_t previousWorkspaceSize) -{ - uintptr_t addr = (uintptr_t) ptr; - addr += previousWorkspaceSize; - if (addr % kCudaMemAlign) - { - addr += kCudaMemAlign - addr % kCudaMemAlign; - } - return (int8_t*) addr; -} - -TritonFlashAttentionPlugin::TritonFlashAttentionPlugin( - int numHeads, int headSize, float softmaxScale, nvinfer1::DataType type) - : mNumHeads(numHeads) - , mHeadSize(headSize) - , mSoftmaxScale(softmaxScale) - , mType(type) -{ -} - -// Parameterized constructor -TritonFlashAttentionPlugin::TritonFlashAttentionPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - readArg(d, mNumHeads); - readArg(d, mHeadSize); - readArg(d, mSoftmaxScale); - readArg(d, mType); - TLLM_CHECK(d == a + length); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* TritonFlashAttentionPlugin::clone() const noexcept -{ - auto* plugin = new TritonFlashAttentionPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs TritonFlashAttentionPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - // Output shape. - // output tensor [batchSize, seqLen, mNumHeads, head_size] - assert(outputIndex == 0); - return inputs[outputIndex]; -} - -bool TritonFlashAttentionPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - // In this example, inputs: Q, K, V, outputs: Out - assert(nbInputs + nbOutputs == 4); - assert(0 <= pos && pos < nbInputs + nbOutputs); - - bool is_valid = false; - if (0 <= pos && pos < 3) // Q, K, V - { - is_valid = inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (pos == nbInputs) // Out - { - is_valid = inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - } - return is_valid; -} - -void TritonFlashAttentionPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t TritonFlashAttentionPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - // Set workspace size if needed. In this example, we need for L and m buffers. - auto const Q = inputs[0]; - int const batchSize = Q.dims.d[0]; - int const seqLen = Q.dims.d[2]; - int const numBuffers = 2; - size_t workspaces[numBuffers]; - workspaces[0] = sizeof(float) * batchSize * mNumHeads * seqLen; - workspaces[1] = sizeof(float) * batchSize * mNumHeads * seqLen; - - size_t total = 0; - for (int i = 0; i < numBuffers; i++) - { - total += workspaces[i]; - if (workspaces[i] % kCudaMemAlign) - { - total += kCudaMemAlign - (workspaces[i] % kCudaMemAlign); - } - } - return total; -} - -template -int TritonFlashAttentionPlugin::enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) -{ - assert(inputDesc[0].dims.d[1] == mNumHeads && inputDesc[0].dims.d[3] == mHeadSize); - assert(inputDesc[1].dims.d[1] == mNumHeads && inputDesc[1].dims.d[3] == mHeadSize); - assert(inputDesc[2].dims.d[1] == mNumHeads && inputDesc[2].dims.d[3] == mHeadSize); - - int batchSize = inputDesc[0].dims.d[0]; - int seqLen = inputDesc[0].dims.d[2]; - - T* Out = reinterpret_cast(outputs[0]); - - const size_t bufSize = sizeof(float) * batchSize * mNumHeads * seqLen; - float* L = reinterpret_cast(workspace); - float* M = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(L), bufSize)); - - T const* Q = reinterpret_cast(inputs[0]); - T const* K = reinterpret_cast(inputs[1]); - T const* V = reinterpret_cast(inputs[2]); - - // Launch a cuda kernel generated by Triton AoT. - int res = 0; - if (std::is_same::value) - { - res = fmha_d64_fp32_default(stream, reinterpret_cast(Out), reinterpret_cast(L), - reinterpret_cast(M), reinterpret_cast(Q), reinterpret_cast(K), - reinterpret_cast(V), mSoftmaxScale, batchSize, mNumHeads, seqLen); - } - else - { - res = fmha_d64_fp16_default(stream, reinterpret_cast(Out), reinterpret_cast(L), - reinterpret_cast(M), reinterpret_cast(Q), reinterpret_cast(K), - reinterpret_cast(V), mSoftmaxScale, batchSize, mNumHeads, seqLen); - } - return res; -} - -int TritonFlashAttentionPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - int res = 1; - if (mType == DataType::kHALF) - { - res = enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (mType == DataType::kFLOAT) - { - res = enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - sync_check_cuda_error(); - return res; -} - -// IPluginV2Ext Methods -nvinfer1::DataType TritonFlashAttentionPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - assert(index == 0); - return inputTypes[0]; -} - -// IPluginV2 Methods - -char const* TritonFlashAttentionPlugin::getPluginType() const noexcept -{ - return TRITON_FLASH_ATTENTION_PLUGIN_NAME; -} - -char const* TritonFlashAttentionPlugin::getPluginVersion() const noexcept -{ - return TRITON_FLASH_ATTENTION_PLUGIN_VERSION; -} - -int TritonFlashAttentionPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int TritonFlashAttentionPlugin::initialize() noexcept -{ - // Load kernels generated by Triton AoT. - load_fmha_d64_fp32(); - load_fmha_d64_fp16(); - return 0; -} - -void TritonFlashAttentionPlugin::terminate() noexcept -{ - // Unload kernels generated by Triton AoT. - unload_fmha_d64_fp32(); - unload_fmha_d64_fp16(); -} - -size_t TritonFlashAttentionPlugin::getSerializationSize() const noexcept -{ - return sizeof(mNumHeads) + sizeof(mHeadSize) + sizeof(mSoftmaxScale) + sizeof(mType); -} - -void TritonFlashAttentionPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - writeArg(d, mNumHeads); - writeArg(d, mHeadSize); - writeArg(d, mSoftmaxScale); - writeArg(d, mType); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void TritonFlashAttentionPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -void TritonFlashAttentionPlugin::setPluginNamespace(char const* libNamespace) noexcept -{ - mNamespace = libNamespace; -} - -char const* TritonFlashAttentionPlugin::getPluginNamespace() const noexcept -{ - return mNamespace.c_str(); -} - -/////////////// - -TritonFlashAttentionPluginCreator::TritonFlashAttentionPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("num_heads", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("head_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("softmax_scale", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* TritonFlashAttentionPluginCreator::getPluginName() const noexcept -{ - return TRITON_FLASH_ATTENTION_PLUGIN_NAME; -} - -char const* TritonFlashAttentionPluginCreator::getPluginVersion() const noexcept -{ - return TRITON_FLASH_ATTENTION_PLUGIN_VERSION; -} - -PluginFieldCollection const* TritonFlashAttentionPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* TritonFlashAttentionPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int numHeads = 0; - int headSize = 0; - float softmaxScale = 1.0f; - nvinfer1::DataType type; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "num_heads")) - { - assert(fields[i].type == PluginFieldType::kINT32); - numHeads = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "head_size")) - { - assert(fields[i].type == PluginFieldType::kINT32); - headSize = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "softmax_scale")) - { - assert(fields[i].type == PluginFieldType::kFLOAT32); - softmaxScale = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - assert(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - auto* obj = new TritonFlashAttentionPlugin(numHeads, headSize, softmaxScale, type); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - std::cerr << "Caught exception: " << e.what() << std::endl; - } - return nullptr; -} - -IPluginV2* TritonFlashAttentionPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call TritonFlashAttentionPlugin::destroy() - try - { - auto* obj = new TritonFlashAttentionPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - std::cerr << "Caught exception: " << e.what() << std::endl; - } - return nullptr; -} - -void TritonFlashAttentionPluginCreator::setPluginNamespace(char const* libNamespace) noexcept -{ - mNamespace = libNamespace; -} - -char const* TritonFlashAttentionPluginCreator::getPluginNamespace() const noexcept -{ - return mNamespace.c_str(); -} - -} // namespace openai_triton::plugin diff --git a/examples/openai_triton/manual_plugin/TritonFlashAttentionPlugin.h b/examples/openai_triton/manual_plugin/TritonFlashAttentionPlugin.h deleted file mode 100644 index cd95eb48ec2a..000000000000 --- a/examples/openai_triton/manual_plugin/TritonFlashAttentionPlugin.h +++ /dev/null @@ -1,113 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-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. - */ -#pragma once - -#include - -#include -#include -#include -#include - -#include -#include - -namespace openai_triton::plugin -{ - -class TritonFlashAttentionPlugin : public nvinfer1::IPluginV2DynamicExt -{ -public: - TritonFlashAttentionPlugin(int numHeads, int headSize, float softmaxScale, nvinfer1::DataType type); - - TritonFlashAttentionPlugin(void const* data, size_t length); - - ~TritonFlashAttentionPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - template - int enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - void setPluginNamespace(char const* pluginNamespace) noexcept override; - char const* getPluginNamespace() const noexcept override; - -private: - const std::string mLayerName; - std::string mNamespace; - - int mNumHeads; - int mHeadSize; - float mSoftmaxScale; - nvinfer1::DataType mType; - - CUmodule mModule; - CUfunction mKernel; -}; - -class TritonFlashAttentionPluginCreator : public nvinfer1::IPluginCreator -{ -public: - TritonFlashAttentionPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - - void setPluginNamespace(char const* pluginNamespace) noexcept override; - - char const* getPluginNamespace() const noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; - std::string mNamespace; -}; - -} // namespace openai_triton::plugin diff --git a/examples/openai_triton/manual_plugin/build.py b/examples/openai_triton/manual_plugin/build.py deleted file mode 100644 index 12b3ca883e7c..000000000000 --- a/examples/openai_triton/manual_plugin/build.py +++ /dev/null @@ -1,137 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -import argparse -import math -import time -from pathlib import Path - -import tensorrt as trt -from plugin import LAYER_NAME, FmhaLayer, get_engine_name - -import tensorrt_llm -from tensorrt_llm.builder import Builder, BuilderConfig -from tensorrt_llm.logger import logger -from tensorrt_llm.network import net_guard - - -def build_engine(builder: Builder, builder_config: BuilderConfig, - engine_name: str, args: argparse.Namespace) -> trt.IHostMemory: - ''' - - @brief: Build a TensorRT engine. - @param args: The cmd line arguments. - @return: The built or refitted engine. - ''' - - # Initialize Module - softmax_scale = 1.0 / math.sqrt(args.head_size) - layer = FmhaLayer(args.num_heads, args.head_size, softmax_scale, args.dtype) - - # Module -> Network - network = builder.create_network() - network.trt_network.name = engine_name - network.plugin_config.to_legacy_setting() - with net_guard(network): - # Prepare - inputs = layer.prepare_inputs(args.max_batch_size, args.max_seq_len) - # Forward - logger.debug(f'model inputs: {inputs}') - out = layer(*inputs) - out.trt_tensor.name = 'out' - - # Network -> Engine - engine = builder.build_engine(network, builder_config) - config_path = Path(args.output_dir) / 'config.json' - builder.save_config(builder_config, str(config_path)) - return engine - - -def build(args): - tensorrt_llm.logger.set_level(args.log_level) - output_dir = Path(args.output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - - builder = Builder() - cache = None - builder_config = builder.create_builder_config( - name=LAYER_NAME, - precision=args.dtype, - timing_cache=args.timing_cache if cache is None else cache, - profiling_verbosity=args.profiling_verbosity) - - engine_name = get_engine_name(args.head_size, args.dtype) - engine = build_engine(builder, builder_config, engine_name, args) - assert engine is not None - - engine_path = output_dir / engine_name - logger.info(f'Serializing engine to {str(engine_path)}...') - tik = time.time() - with engine_path.open('wb') as f: - f.write(engine) - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Engine serialized. Total time: {t}') - - ok = builder.save_timing_cache(builder_config, - Path(args.output_dir) / "model.cache") - assert ok, "Failed to save timing cache." - - -if __name__ == '__main__': - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter) - parser.add_argument('--max_batch_size', type=int, default=4) - parser.add_argument('--max_seq_len', type=int, default=256) - parser.add_argument('--num_heads', type=int, default=8) - parser.add_argument('--head_size', type=int, default=64) - parser.add_argument('--dtype', - type=str, - default='float16', - choices=['float16', 'float32']) - parser.add_argument( - '--timing_cache', - type=str, - default='model.cache', - help='The path of to read timing cache from, will be ignored ' - 'if the file does not exist') - parser.add_argument( - '--profiling_verbosity', - type=str, - default='layer_names_only', - choices=['layer_names_only', 'detailed', 'none'], - help= - 'The profiling verbosity for the generated TRT engine. Set to detailed can inspect tactic choices and kernel parameters.' - ) - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument( - '--output_dir', - type=str, - default='outputs', - help='The path to save the serialized engine files, timing cache ' - 'file and model configs') - args = parser.parse_args() - - logger.set_level(args.log_level) - logger.info('Parameters'.center(40, '=')) - for k, v in vars(args).items(): - logger.info(f' - {k.ljust(15, ".")}: {v}') - logger.info(''.center(40, '=')) - - tik = time.time() - logger.info('Build TensorRT engine.') - build(args) - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Total time of building TRT engine: {t}') diff --git a/examples/openai_triton/manual_plugin/fmha_triton.py b/examples/openai_triton/manual_plugin/fmha_triton.py deleted file mode 100644 index 3e47dff263f0..000000000000 --- a/examples/openai_triton/manual_plugin/fmha_triton.py +++ /dev/null @@ -1,135 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -""" -Fused attention from triton tutorial. -Modified from the original implementation - - https://github.com/openai/triton/blob/main/python/tutorials/06-fused-attention.py -=============== - -This is a Triton implementation of the Flash Attention algorithm -(see: Dao et al., https://arxiv.org/pdf/2205.14135v2.pdf; Rabe and Staats https://arxiv.org/pdf/2112.05682v2.pdf) -""" - -import torch -import triton -import triton.language as tl - - -# yapf: disable -@triton.jit -def fused_attention_kernel( - Out, L, M, # outputs - Q, K, V, - sm_scale, - batch_size, num_heads, seq_len, - BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, - BLOCK_N: tl.constexpr, -): - start_m = tl.program_id(0) - off_hz = tl.program_id(1) - stride_h = BLOCK_DMODEL * seq_len - - # initialize offsets - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - offs_n = tl.arange(0, BLOCK_N) - offs_d = tl.arange(0, BLOCK_DMODEL) - off_q = off_hz * stride_h + offs_m[:, None] * BLOCK_DMODEL + offs_d[None, :] - off_k = off_hz * stride_h + offs_n[None, :] * BLOCK_DMODEL + offs_d[:, None] - off_v = off_hz * stride_h + offs_n[:, None] * BLOCK_DMODEL + offs_d[None, :] - # Initialize pointers to Q, K, V - q_ptrs = Q + off_q - k_ptrs = K + off_k - v_ptrs = V + off_v - # initialize pointer to m and l - m_prev = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") - l_prev = tl.zeros([BLOCK_M], dtype=tl.float32) - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) - # load q: it will stay in SRAM throughout - q = tl.load(q_ptrs) - # loop over k, v and update accumulator - for start_n in range(0, (start_m + 1) * BLOCK_M, BLOCK_N): - # -- compute qk ---- - k = tl.load(k_ptrs) - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) - qk += tl.dot(q, k) - qk *= sm_scale - qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, float("-inf")) - # compute new m - m_curr = tl.maximum(tl.max(qk, 1), m_prev) - # correct old l - l_prev *= tl.exp(m_prev - m_curr) - # attention weights - p = tl.exp(qk - m_curr[:, None]) - l_curr = tl.sum(p, 1) + l_prev - # rescale operands of matmuls - l_rcp = 1. / l_curr - p *= l_rcp[:, None] - acc *= (l_prev * l_rcp)[:, None] - # update acc - p = p.to(Q.dtype.element_ty) - v = tl.load(v_ptrs) - acc += tl.dot(p, v) - # update m_i and l_i - l_prev = l_curr - m_prev = m_curr - # update pointers - k_ptrs += BLOCK_N * BLOCK_DMODEL - v_ptrs += BLOCK_N * BLOCK_DMODEL - # rematerialize offsets to save registers - start_m = tl.program_id(0) - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - # write back l and m - l_ptrs = L + off_hz * seq_len + offs_m - m_ptrs = M + off_hz * seq_len + offs_m - tl.store(l_ptrs, l_prev) - tl.store(m_ptrs, m_prev) - # initialize pointers to output - offs_n = tl.arange(0, BLOCK_DMODEL) - off_o = off_hz * stride_h + offs_m[:, None] * BLOCK_DMODEL + offs_n[None, :] - out_ptrs = Out + off_o - tl.store(out_ptrs, acc) - - -def fused_attention(q, k, v, sm_scale, o_buf=None, l_buf=None, m_buf=None): - BLOCK = 128 if q.dtype == torch.float16 else 64 - # shape constraints - Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] - assert Lq == Lk and Lk == Lv - assert Lk in {16, 32, 64, 128} - o = torch.empty_like(q) if o_buf is None else o_buf - grid = (triton.cdiv(q.shape[2], BLOCK), q.shape[0] * q.shape[1], 1) - shape = (q.shape[0] * q.shape[1], q.shape[2]) - L = torch.empty(shape, device=q.device, dtype=torch.float32) if l_buf is None else l_buf - m = torch.empty(shape, device=q.device, dtype=torch.float32) if m_buf is None else m_buf - - num_warps = 4 if Lk <= 64 else 8 - # Adjust num_stages for limited resource cases. - num_stages = 2 if torch.cuda.get_device_capability() >= (8, 0) else 1 - - fused_attention_kernel[grid]( - o, L, m, - q, k, v, - sm_scale, - q.shape[0], q.shape[1], q.shape[2], - # tl.constexpr - BLOCK_M=BLOCK, - BLOCK_N=BLOCK, - BLOCK_DMODEL=Lk, - num_warps=num_warps, - num_stages=num_stages, - ) - - return o -# yapf: enable diff --git a/examples/openai_triton/manual_plugin/plugin.py b/examples/openai_triton/manual_plugin/plugin.py deleted file mode 100644 index 7009caaeb680..000000000000 --- a/examples/openai_triton/manual_plugin/plugin.py +++ /dev/null @@ -1,133 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -import ctypes -from collections import OrderedDict -from pathlib import Path -from typing import List - -import numpy as np -import tensorrt as trt - -from tensorrt_llm._common import default_trtnet -from tensorrt_llm._utils import str_dtype_to_trt -from tensorrt_llm.functional import Tensor, _create_tensor -from tensorrt_llm.module import Module - -TRT_LLM_PLUGIN_NAMESPACE = 'tensorrt_llm' -LAYER_NAME = 'TritonFlashAttentionLayer' -FMHA_KERNEL_BLOCK_SIZE = 128 - - -def _load_triton_plugin_lib(): - triton_plugin_dir = Path(__file__).parent.absolute() - plugin_lib = triton_plugin_dir / 'build/libtrt_llm_custom_plugins.so' - handle = ctypes.CDLL(plugin_lib, mode=ctypes.RTLD_GLOBAL) - if handle is None: - raise ImportError('TensorRT LLM Triton Plugin is unavailable') - handle.initOpenAiTritonPlugins.argtypes = [ctypes.c_void_p, ctypes.c_char_p] - handle.initOpenAiTritonPlugins.restype = ctypes.c_bool - assert handle.initOpenAiTritonPlugins( - None, TRT_LLM_PLUGIN_NAMESPACE.encode('utf-8')) - - -_load_triton_plugin_lib() - - -def flash_attention_op(num_heads: int, head_size: int, softmax_scale: float, - inputs: List[trt.ITensor]) -> Tensor: - # Create a plugin instance. - plugin_creator = trt.get_plugin_registry().get_plugin_creator( - 'TritonFlashAttention', '1', TRT_LLM_PLUGIN_NAMESPACE) - assert plugin_creator is not None - - pfc = trt.PluginFieldCollection([ - trt.PluginField("num_heads", np.array([num_heads], np.int32), - trt.PluginFieldType.INT32), - trt.PluginField("head_size", np.array([head_size], np.int32), - trt.PluginFieldType.INT32), - trt.PluginField("softmax_scale", np.array([softmax_scale], np.float32), - trt.PluginFieldType.FLOAT32), - trt.PluginField("type_id", np.array([int(inputs[0].dtype)], np.int32), - trt.PluginFieldType.INT32) - ]) - plugin = plugin_creator.create_plugin("flash_attention", pfc) - layer = default_trtnet().add_plugin_v2(inputs, plugin) - return _create_tensor(layer.get_output(0), layer) - - -class FmhaLayer(Module): - - def __init__(self, num_heads: int, head_size: int, softmax_scale: float, - dtype: str): - super().__init__() - self.num_heads = num_heads - self.head_size = head_size - self.softmax_scale = softmax_scale - self.dtype = str_dtype_to_trt(dtype) - - def forward(self, Q: Tensor, K: Tensor, V: Tensor): - inputs = [Q, K, V] - out = flash_attention_op(num_heads=self.num_heads, - head_size=self.head_size, - softmax_scale=self.softmax_scale, - inputs=[p.trt_tensor for p in inputs]) - out.mark_output('out', self.dtype) - return out - - def prepare_inputs(self, max_batch_size: int, max_len: int) -> List[Tensor]: - ''' - - @brief: Prepare inputs Tensors for the model, the given sizes are used to - determine the ranges of the dimensions of when using TRT dynamic shapes. - - @return: a list contains values which can be fed into the self.forward() - ''' - - bs_range = [1, (max_batch_size + 1) // 2, max_batch_size] - max_len_range = [1, (max_len + 1) // 2, max_len] - - dynamic_shape = [-1, self.num_heads, -1, self.head_size] - Q = Tensor(name='Q', - dtype=self.dtype, - shape=dynamic_shape, - dim_range=OrderedDict([ - ('batch_size', [bs_range]), - ('num_heads', [self.num_heads]), - ('seq_len', [max_len_range]), - ('head_size', [self.head_size]), - ])) - K = Tensor(name='K', - dtype=self.dtype, - shape=dynamic_shape, - dim_range=OrderedDict([ - ('batch_size', [bs_range]), - ('num_heads', [self.num_heads]), - ('seq_len', [max_len_range]), - ('head_size', [self.head_size]), - ])) - V = Tensor(name='V', - dtype=self.dtype, - shape=dynamic_shape, - dim_range=OrderedDict([ - ('batch_size', [bs_range]), - ('num_heads', [self.num_heads]), - ('seq_len', [max_len_range]), - ('head_size', [self.head_size]), - ])) - return [Q, K, V] - - -def get_engine_name(head_size, dtype): - return f'{LAYER_NAME}_{FMHA_KERNEL_BLOCK_SIZE}_d{head_size}_{dtype}.engine' diff --git a/examples/openai_triton/manual_plugin/run.py b/examples/openai_triton/manual_plugin/run.py deleted file mode 100644 index ec7cf4dd5600..000000000000 --- a/examples/openai_triton/manual_plugin/run.py +++ /dev/null @@ -1,170 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -import argparse -import json -import math -from pathlib import Path - -import torch -from fmha_triton import fused_attention -from plugin import get_engine_name - -from tensorrt_llm import profiler -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import (str_dtype_to_torch, str_dtype_to_trt, - trt_dtype_to_torch) -from tensorrt_llm.logger import logger -from tensorrt_llm.runtime.session import Session, TensorInfo - - -def run(engine_dir, - batch_size, - seq_len, - num_heads, - head_size, - do_benchmark=False): - # Load trt engine. - engine_dir = Path(engine_dir) - config_path = engine_dir / 'config.json' - with config_path.open('r') as f: - config = json.load(f) - dtype = config['builder_config']['precision'] - serialize_path = engine_dir / get_engine_name(head_size, dtype) - - with open(serialize_path, 'rb') as f: - session = Session.from_serialized_engine(f.read()) - - # Prepare input tensors. - torch_dtype = str_dtype_to_torch(dtype) if isinstance(dtype, str) else dtype - shape = (batch_size, num_heads, seq_len, head_size) - q = torch.normal(mean=0.1, - std=0.2, - size=shape, - dtype=torch_dtype, - device='cuda') - k = torch.normal(mean=0.4, - std=0.2, - size=shape, - dtype=torch_dtype, - device='cuda') - v = torch.normal(mean=0.3, - std=0.2, - size=shape, - dtype=torch_dtype, - device='cuda') - inputs = {'Q': q, 'K': k, 'V': v} - - # Prepare output tensors. - output_info = session.infer_shapes([ - TensorInfo(name, str_dtype_to_trt(dtype), tensor.shape) - for name, tensor in inputs.items() - ]) - logger.debug(f'output info {output_info}') - outputs = { - t.name: - torch.empty(tuple(t.shape), - dtype=trt_dtype_to_torch(t.dtype), - device='cuda') - for t in output_info - } - - # Execute model inference - stream = torch.cuda.Stream() - ok = session.run(inputs=inputs, outputs=outputs, stream=stream.cuda_stream) - assert ok, 'Engine execution failed' - - # Sanity check - stream.synchronize() - sm_scale = 1.0 / math.sqrt(head_size) - ref = fused_attention(q, k, v, sm_scale) - out = outputs["out"] - logger.debug( - f'Out: vals: {out.view(1, -1)} abs_sum: {out.float().abs().sum()}') - logger.debug( - f'Ref: vals: {ref.view(1, -1)} abs_sum: {ref.float().abs().sum()}') - torch.testing.assert_close(out, ref) - - if do_benchmark: - n_repeats = 10 - - # For fair comparison, pre-allocate buffers as trt plugin does. - shape = (q.shape[0] * q.shape[1], q.shape[2]) - L = torch.empty(shape, device=q.device, dtype=torch.float32) - m = torch.empty(shape, device=q.device, dtype=torch.float32) - o = torch.empty_like(q) - - # Triton warm-up - fused_attention(q, k, v, sm_scale, l_buf=L, m_buf=m, o_buf=o) - stream.synchronize() - for _ in range(n_repeats): - profiler.start('Triton') - fused_attention(q, k, v, sm_scale, l_buf=L, m_buf=m, o_buf=o) - stream.synchronize() - profiler.stop('Triton') - - # TRT warm-up - stream.synchronize() - ok = session.run(inputs=inputs, - outputs=outputs, - stream=stream.cuda_stream) - stream.synchronize() - for _ in range(n_repeats): - profiler.start('TRT Plugin') - ok = session.run(inputs=inputs, - outputs=outputs, - stream=stream.cuda_stream) - stream.synchronize() - profiler.stop('TRT Plugin') - assert ok - profiler.summary() - - -if __name__ == '__main__': - emit_engine_arch_deprecation("run.py") - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter) - parser.add_argument('--batch_size', type=int, default=4) - parser.add_argument('--seq_len', type=int, default=128) - parser.add_argument('--num_heads', type=int, default=8) - parser.add_argument('--head_size', type=int, default=64) - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument( - '--engine_dir', - type=Path, - default='outputs', - help='The directory where serialized engine files locate.') - parser.add_argument( - '--benchmark', - action='store_true', - help='Do performance benchmark compared to triton baseline.') - args = parser.parse_args() - - logger.set_level(args.log_level) - logger.info('Parameters'.center(40, '=')) - for k, v in vars(args).items(): - logger.info(f' - {k.ljust(15, ".")}: {v}') - logger.info(''.center(40, '=')) - - assert args.engine_dir.exists(), \ - f"Engine file {str(args.engine_dir)} doesn't exists." - - logger.info('Inference using the built TensorRT engine.') - run(args.engine_dir, - args.batch_size, - args.seq_len, - args.num_heads, - args.head_size, - do_benchmark=args.benchmark) - logger.info('Done.') diff --git a/examples/openai_triton/manual_plugin/tritonPlugins.cpp b/examples/openai_triton/manual_plugin/tritonPlugins.cpp deleted file mode 100644 index 27b1ece08448..000000000000 --- a/examples/openai_triton/manual_plugin/tritonPlugins.cpp +++ /dev/null @@ -1,133 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2022 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. - */ -#include "NvInferRuntime.h" -#include "TritonFlashAttentionPlugin.h" - -#include -#include -#include -#include -#include -#include -#include - -namespace -{ - -// This singleton ensures that each plugin is only registered once for a given -// namespace and type, and attempts of duplicate registration are ignored. -class TritonPluginCreatorRegistry -{ -public: - static TritonPluginCreatorRegistry& getInstance() - { - static TritonPluginCreatorRegistry instance; - return instance; - } - - template - void addPluginCreator(void* logger, char const* libNamespace) - { - // Make accesses to the plugin creator registry thread safe - std::lock_guard lock(mRegistryLock); - - std::string errorMsg; - std::string verboseMsg; - - std::unique_ptr pluginCreator{new CreatorType{}}; - pluginCreator->setPluginNamespace(libNamespace); - - nvinfer1::ILogger* trtLogger = static_cast(logger); - std::string pluginType = std::string{pluginCreator->getPluginNamespace()} - + "::" + std::string{pluginCreator->getPluginName()} + " version " - + std::string{pluginCreator->getPluginVersion()}; - - if (mRegistryList.find(pluginType) == mRegistryList.end()) - { - bool status = getPluginRegistry()->registerCreator(*pluginCreator, libNamespace); - if (status) - { - mRegistry.push(std::move(pluginCreator)); - mRegistryList.insert(pluginType); - verboseMsg = "Registered plugin creator - " + pluginType; - } - else - { - errorMsg = "Could not register plugin creator - " + pluginType; - } - } - else - { - verboseMsg = "Plugin creator already registered - " + pluginType; - } - - if (trtLogger) - { - if (!errorMsg.empty()) - { - trtLogger->log(nvinfer1::ILogger::Severity::kERROR, errorMsg.c_str()); - } - - if (!verboseMsg.empty()) - { - trtLogger->log(nvinfer1::ILogger::Severity::kVERBOSE, verboseMsg.c_str()); - } - } - } - - ~TritonPluginCreatorRegistry() - { - std::lock_guard lock(mRegistryLock); - - // Release pluginCreators in LIFO order of registration. - while (!mRegistry.empty()) - { - mRegistry.pop(); - } - mRegistryList.clear(); - } - -private: - TritonPluginCreatorRegistry() {} - - std::mutex mRegistryLock; - std::stack> mRegistry; - std::unordered_set mRegistryList; - -public: - TritonPluginCreatorRegistry(TritonPluginCreatorRegistry const&) = delete; - void operator=(TritonPluginCreatorRegistry const&) = delete; -}; - -template -void initializeTritonPlugin(void* logger, char const* libNamespace) -{ - TritonPluginCreatorRegistry::getInstance().addPluginCreator(logger, libNamespace); -} - -} // namespace - -// New Plugin APIs - -extern "C" -{ - bool initOpenAiTritonPlugins(void* logger, char const* libNamespace) - { - initializeTritonPlugin(logger, libNamespace); - return true; - } -} // extern "C" diff --git a/examples/openai_triton/plugin_autogen/README.md b/examples/openai_triton/plugin_autogen/README.md deleted file mode 100644 index 0c046330fa0b..000000000000 --- a/examples/openai_triton/plugin_autogen/README.md +++ /dev/null @@ -1,97 +0,0 @@ -# Integrating Triton Kernel with TensorRT Plugin Generator - -In the previous [OpenAI Triton Plugin in TensorRT-LLM](../../openai_triton/README.md) tutorial, it is demonstrated how to integrate a Triton kernel by manually writing a TensorRT plugin in C++ as well as a Python wrapper. In the latest TensorRT-LLM, we now have an end-to-end tool called PluginGen that simplifies this process. All you need to do is providing a plugin configuration. - -In this example, we will introduce the usage of the PluginGen tool and demonstrate the integration of the [Fused Attention](../openai_triton/fmha_triton.py) kernel. - - -To use the feature, you need a Triton version posterior to the [d0c35b3](https://github.com/openai/triton/commit/d0c35b3b7d6badf0c0d56a821dddab7ace73b4de) commit -and this example has been tested on the [d4644d6](https://github.com/openai/triton/tree/d4644d6cb3ae674e1f15932cac1f28104795744f) commit. - -## Introduction to the PluginGen Toolkit - -The PluginGen script can be found at `tensorrt_llm/tools/triton_integration/plugin_gen.py`. Its usage is as follows: - -```sh -usage: plugin_gen.py [-h] --workspace WORKSPACE --kernel_config KERNEL_CONFIG [--tensorrt_llm_include_path TENSORRT_LLM_INCLUDE_PATH] -``` - -There are three command-line arguments: - -1. `workspace`: This is the root directory to hold the temporary generation files. PluginGen should not alter anything outside of the workspace, -2. `kernel_config`: This is a Python file that holds a variable called `KERNELS` of type `List[KernelMetaData]`. PluginGen can process one or more kernels at a time, -3. `tensorrt_llm_include_path`: This is the path to the TensorRT LLM include directory. It is used to include the TensorRT LLM header files in the generated plugin. - -You can refer to [./kernel_config.py](./kernel_config.py) for an example of `KernelMetaData` for the Fused Attention kernel. It contains several fields: - -1. `ios` (short for "input and outputs"): This holds all the metadata of the inputs and outputs of the Triton kernel, including the data type, shape, and the name of the tensor. There are several kinds of arguments: - - `InputArg`: A common variable for this kernel. - - `OutputArg`: An output of the kernel. - - `ParamArg`: A special input that is a constant; it will be mapped to a PluginField in the generated plugin. - - `DimSizeArg`: A special input that is an expression of the input tensors' shape size; it requires an inference rule to compute the value. -2. `shape_infer_rules`: This field contains two types of rules: - a) Rules for deducing the shape of the output tensors from the input tensors. The syntax is like `input0[dim_names], input1[dim_names] -> output0[dim_names]`. - b) Rules for inferring `DimSizeArg`. The syntax is like `input0[dim_names]: some_dim_expression -> arg_name`. - -The user should provide the kernel configurations as well as the Triton kernel script, and the PluginGen toolkit will handle the following steps: - -1. Trigger the Triton AOT tool to obtain the necessary C files. -2. Generate the C++ code for a TensorRT plugin. -3. Generate the CMAKE code for compiling all the C/C++ files. -4. Perform the compilation and generate `libtriton_plugins.so`. -5. Generate a `functional.py` containing a Python wrapper for this plugin. - -After the generation, you should have `libtriton_plugins.so` and `functional.py` in the workspace. You can use them to integrate the Triton kernel by simply using the corresponding Python methods in the generated `functional.py` during the model-building stage, just like other layers located in the TensorRT LLM built-in `functional.py`. - -## End-to-End Example for FHMA Kernel Integration - -In this section, we will demonstrate the integration of the Fused Attention kernel. The steps are as follows: - -### Pre-Stage: Install Triton with a Specific Version - -In case the Triton AOT tool's update breaks compatibility, we recommend installing a specific version of Triton. The commit we tested is [d4644d6](https://github.com/openai/triton/tree/d4644d6cb3ae674e1f15932cac1f28104795744f). - -Install Triton with the following commands: - -```sh -git clone https://github.com/openai/triton -cd triton/python/ -pip install cmake && pip install . -cd - -``` - -### Step 1: Prepare the Configuration for FHMA - -To instruct the PluginGen toolkit on how to generate the plugin, please provide a Python file containing the metadata of the kernels. You can refer to [./kernel_config.py](./kernel_config.py) for an example of preparing `KernelMetaData` for the Fused Attention kernel. - -### Step 2: Run the PluginGen Tool and Generate the Plugin - -```sh -python3 {GIT_ROOT_DIR}/tensorrt_llm/tools/plugin_gen/plugin_gen.py --workspace ./tmp --kernel_config ./kernel_config.py -``` - -PluginGen will generate all the necessary files within the `./tmp` directory. The final output will be located in the `./tmp/output` directory, where you should ideally find two files: - -``` --rw-r--r-- 1 1001 1001 2163 Sep 21 17:13 functional.py --rwxr-xr-x 1 1001 1001 3748464 Sep 21 17:13 libtriton_plugins.so -``` - -### Post-Stage: Use the Plugin - -To use the plugin in a TensorRT LLM model, please refer to the generated `output/functional.py`. It should contain Python wrappers for all the plugins. To use the plugins, first import `functional.py` and then use the corresponding Python methods to build the model. - -For an example of using the Fused Attention plugin in a model, please refer to [build_engine.py](./build_engine.py) for building the TensorRT engine and [run_engine.py](./run_engine.py) for running the engine in the runtime. - -To run the example, you can use the following commands: - -```sh -# copy the triton script to the current directory -cp ../manual_plugin/fmha_triton.py . - -# build the TensorRT engine -python3 build_engine.py - -# run the engine -python3 run_engine.py -``` diff --git a/examples/openai_triton/plugin_autogen/build_engine.py b/examples/openai_triton/plugin_autogen/build_engine.py deleted file mode 100644 index 23b829f0b7e4..000000000000 --- a/examples/openai_triton/plugin_autogen/build_engine.py +++ /dev/null @@ -1,206 +0,0 @@ -import argparse -import math -# include plugins -# yapf: disable -import os -import sys -import time -from pathlib import Path -from typing import List, OrderedDict - -import tensorrt as trt - -# from plugin import LAYER_NAME, FmhaLayer, get_engine_name -import tensorrt_llm -from tensorrt_llm import Module, str_dtype_to_trt -from tensorrt_llm.builder import Builder, BuilderConfig -from tensorrt_llm.functional import Tensor -from tensorrt_llm.logger import logger -from tensorrt_llm.network import net_guard - -sys.path.append(os.environ.get('PLUGIN_GEN_WORKSPACE', './tmp')) -from functional import fused_attention_kernel # isort:skip -# yapf: enable - - -def get_engine_name(head_size: int, dtype: str) -> str: - return f'fmha_{head_size}_{dtype}.engine' - - -class FmhaLayer(Module): - - def __init__(self, num_heads: int, head_size: int, softmax_scale: float): - super().__init__() - self.num_heads = num_heads - self.head_size = head_size - self.softmax_scale = softmax_scale - self.dtype = str_dtype_to_trt('float16') - - def forward(self, Q: Tensor, K: Tensor, V: Tensor): - inputs = [Q, K, V] - Out, L, M = fused_attention_kernel(self.softmax_scale, self.num_heads, - *[p.trt_tensor for p in inputs]) - Out.mark_output('out', self.dtype) - L.mark_output('L', self.dtype) - M.mark_output('M', self.dtype) - return Out, L, M - - def prepare_inputs(self, max_batch_size: int, max_len: int) -> List[Tensor]: - ''' - - @brief: Prepare inputs Tensors for the model, the given sizes are used to - determine the ranges of the dimensions of when using TRT dynamic shapes. - - @return: a list contains values which can be fed into the self.forward() - ''' - - bs_range = [1, (max_batch_size + 1) // 2, max_batch_size] - max_len_range = [1, (max_len + 1) // 2, max_len] - - dynamic_shape = [-1, self.num_heads, -1, self.head_size] - Q = Tensor(name='Q', - dtype=trt.float16, - shape=dynamic_shape, - dim_range=OrderedDict([ - ('batch_size', [bs_range]), - ('num_heads', [self.num_heads]), - ('seq_len', [max_len_range]), - ('head_size', [self.head_size]), - ])) - K = Tensor(name='K', - dtype=trt.float16, - shape=dynamic_shape, - dim_range=OrderedDict([ - ('batch_size', [bs_range]), - ('num_heads', [self.num_heads]), - ('seq_len', [max_len_range]), - ('head_size', [self.head_size]), - ])) - V = Tensor(name='V', - dtype=trt.float16, - shape=dynamic_shape, - dim_range=OrderedDict([ - ('batch_size', [bs_range]), - ('num_heads', [self.num_heads]), - ('seq_len', [max_len_range]), - ('head_size', [self.head_size]), - ])) - return [Q, K, V] - - -def build_engine(builder: Builder, builder_config: BuilderConfig, - engine_name: str, args: argparse.Namespace) -> trt.IHostMemory: - ''' - @brief: Build a TensorRT engine. - @param args: The cmd line arguments. - @return: The built or refitted engine. - ''' - - # Initialize Module - softmax_scale = 1.0 / math.sqrt(args.head_size) - layer = FmhaLayer(args.num_heads, args.head_size, softmax_scale) - - # Module -> Network - network = builder.create_network() - network.trt_network.name = engine_name - network.plugin_config.to_legacy_setting() - with net_guard(network): - # Prepare - inputs = layer.prepare_inputs(args.max_batch_size, args.max_seq_len) - # Forward - logger.debug(f'model inputs: {inputs}') - layer(*inputs) - - print('dot:') - print(network.to_dot()) - - layer = network.get_layer_by_name(next( - network.get_layers()).name).as_layer() - print('layer', layer.plugin.plugin_type) - print('layer', layer.plugin.plugin_version) - print('layer', layer.plugin.plugin_namespace) - - # Network -> Engine - engine = builder.build_engine(network, builder_config) - config_path = Path(args.output_dir) / 'config.json' - builder.save_config(builder_config, str(config_path)) - return engine - - -def build(args): - tensorrt_llm.logger.set_level(args.log_level) - output_dir = Path(args.output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - - builder = Builder() - cache = None - builder_config = builder.create_builder_config( - name='fmha_triton', - precision=args.dtype, - timing_cache=args.timing_cache if cache is None else cache, - profiling_verbosity=args.profiling_verbosity) - - engine_name = get_engine_name(args.head_size, args.dtype) - engine = build_engine(builder, builder_config, engine_name, args) - assert engine is not None - - engine_path = output_dir / engine_name - logger.info(f'Serializing engine to {str(engine_path)}...') - tik = time.time() - with engine_path.open('wb') as f: - f.write(engine) - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Engine serialized. Total time: {t}') - - ok = builder.save_timing_cache(builder_config, - Path(args.output_dir) / "model.cache") - assert ok, "Failed to save timing cache." - - -if __name__ == '__main__': - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter) - parser.add_argument('--max_batch_size', type=int, default=4) - parser.add_argument('--max_seq_len', type=int, default=256) - parser.add_argument('--num_heads', type=int, default=8) - parser.add_argument('--head_size', type=int, default=64) - parser.add_argument('--dtype', - type=str, - default='float16', - choices=['float16', 'float32']) - parser.add_argument( - '--timing_cache', - type=str, - default='model.cache', - help='The path of to read timing cache from, will be ignored ' - 'if the file does not exist') - parser.add_argument( - '--profiling_verbosity', - type=str, - default='layer_names_only', - choices=['layer_names_only', 'detailed', 'none'], - help= - 'The profiling verbosity for the generated TRT engine. Set to detailed can inspect tactic choices and kernel parameters.' - ) - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument( - '--output_dir', - type=str, - default='outputs', - help='The path to save the serialized engine files, timing cache ' - 'file and model configs') - args = parser.parse_args() - - logger.set_level(args.log_level) - logger.info('Parameters'.center(40, '=')) - for k, v in vars(args).items(): - logger.info(f' - {k.ljust(15, ".")}: {v}') - logger.info(''.center(40, '=')) - - tik = time.time() - logger.info('Build TensorRT engine.') - build(args) - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Total time of building TRT engine: {t}') diff --git a/examples/openai_triton/plugin_autogen/kernel_config.py b/examples/openai_triton/plugin_autogen/kernel_config.py deleted file mode 100644 index 93d64c187d61..000000000000 --- a/examples/openai_triton/plugin_autogen/kernel_config.py +++ /dev/null @@ -1,55 +0,0 @@ -import os - -import torch - -from tensorrt_llm.tools.plugin_gen.core import * - -openai_triton_example_root = os.path.join( - os.path.dirname(os.path.abspath(__file__)), "..", "manual_plugin") - - -def get_fmha_kernel_meta_data(): - block_size = 128 - num_stages = 2 if torch.cuda.get_device_capability() >= (8, 0) else 1 - - return KernelMetaData( - kernel_name='fused_attention_kernel', - ios=[ - # outputs - OutputArg('Out', Type('tensor[fp16]'), hints=['16', '16']), - OutputArg('L', Type('tensor[fp32]'), hints=['16', '16']), - OutputArg('M', Type('tensor[fp32]'), hints=['16', '16']), - # inputs - InputArg('Q', Type('tensor[fp16]'), hints=['16', '16']), - InputArg('K', Type('tensor[fp16]'), hints=['16', '16']), - InputArg('V', Type('tensor[fp16]'), hints=['16', '16']), - ParamArg('sm_scale', Type('fp32')), - DimSizeArg('batch_size'), - ParamArg('num_heads', Type('i32')), - DimSizeArg('seq_len', hints=['', '16']), - # constexprs - Constexpr(block_size), - Constexpr(64), - Constexpr(block_size), - ], - shape_infer_rules=[ - # The following rules helps to deduce the shapes of the output tensors - "Q[*] -> Out[*]", - "Q[m,n,k,*] -> L[m,n,k]", - "Q[m,n,k,*] -> M[m,n,k]", - - # The following rules helps to deduce both DimSizeArgs: batch_size and seq_len - "Q[m,n,k,*] : m -> batch_size", - "Q[m,n,k,*] : k -> seq_len", - ], - version=0, - kernel_file=f'{openai_triton_example_root}/fmha_triton.py', - num_warps=4, - num_stages=num_stages, - grid_dims=(f"(seq_len + {block_size-1}) / {block_size}", - "batch_size * num_heads", "1")) - - -KERNELS = [ - get_fmha_kernel_meta_data(), -] diff --git a/examples/openai_triton/plugin_autogen/run_engine.py b/examples/openai_triton/plugin_autogen/run_engine.py deleted file mode 100644 index 9a438d1d8d87..000000000000 --- a/examples/openai_triton/plugin_autogen/run_engine.py +++ /dev/null @@ -1,170 +0,0 @@ -import argparse -import json -import math -# include plugins -# yapf: disable -import sys -from pathlib import Path - -import torch -from fmha_triton import fused_attention - -from tensorrt_llm import profiler -from tensorrt_llm._utils import (str_dtype_to_torch, str_dtype_to_trt, - trt_dtype_to_torch) -from tensorrt_llm.logger import logger -from tensorrt_llm.runtime.session import Session, TensorInfo - -# from tensorrt_llm.plugin import get_engine_name - - -sys.path.append('./tmp') -from functional import fused_attention_kernel # isort:skip -# yapf: enable - - -def get_engine_name(head_size, dtype): - return f'fmha_{head_size}_{dtype}.engine' - - -def run(engine_dir, - batch_size, - seq_len, - num_heads, - head_size, - do_benchmark=False): - # Load trt engine. - engine_dir = Path(engine_dir) - config_path = engine_dir / 'config.json' - with config_path.open('r') as f: - config = json.load(f) - dtype = config['builder_config']['precision'] - serialize_path = engine_dir / get_engine_name(head_size, dtype) - - with open(serialize_path, 'rb') as f: - session = Session.from_serialized_engine(f.read()) - - # Prepare input tensors. - torch_dtype = str_dtype_to_torch(dtype) if isinstance(dtype, str) else dtype - shape = (batch_size, num_heads, seq_len, head_size) - q = torch.normal(mean=0.1, - std=0.2, - size=shape, - dtype=torch_dtype, - device='cuda') - k = torch.normal(mean=0.4, - std=0.2, - size=shape, - dtype=torch_dtype, - device='cuda') - v = torch.normal(mean=0.3, - std=0.2, - size=shape, - dtype=torch_dtype, - device='cuda') - batch_size = q.shape[0] - seq_len = q.shape[2] - - inputs = {'Q': q, 'K': k, 'V': v} - - # Prepare output tensors. - output_info = session.infer_shapes([ - TensorInfo(name, str_dtype_to_trt(dtype), tensor.shape) - for name, tensor in inputs.items() - ]) - logger.debug(f'output info {output_info}') - outputs = { - t.name: - torch.empty(tuple(t.shape), - dtype=trt_dtype_to_torch(t.dtype), - device='cuda') - for t in output_info - } - - # Execute model inference - stream = torch.cuda.current_stream() - ok = session.run(inputs=inputs, outputs=outputs, stream=stream.cuda_stream) - assert ok, 'Engine execution failed' - - # Sanity check - stream.synchronize() - sm_scale = 1.0 / math.sqrt(head_size) - ref = fused_attention(q, k, v, sm_scale) - out = outputs["out"] - logger.debug( - f'Out: vals: {out.view(1, -1)} abs_sum: {out.float().abs().sum()}') - logger.debug( - f'Ref: vals: {ref.view(1, -1)} abs_sum: {ref.float().abs().sum()}') - torch.testing.assert_close(out, ref) - - if do_benchmark: - n_repeats = 10 - - # For fair comparison, pre-allocate buffers as trt plugin does. - shape = (q.shape[0] * q.shape[1], q.shape[2]) - L = torch.empty(shape, device=q.device, dtype=torch.float32) - m = torch.empty(shape, device=q.device, dtype=torch.float32) - o = torch.empty_like(q) - - # Triton warm-up - fused_attention(q, k, v, sm_scale, l_buf=L, m_buf=m, o_buf=o) - stream.synchronize() - for _ in range(n_repeats): - profiler.start('Triton') - fused_attention(q, k, v, sm_scale, l_buf=L, m_buf=m, o_buf=o) - stream.synchronize() - profiler.stop('Triton') - - # TRT warm-up - stream.synchronize() - ok = session.run(inputs=inputs, - outputs=outputs, - stream=stream.cuda_stream) - stream.synchronize() - for _ in range(n_repeats): - profiler.start('TRT Plugin') - ok = session.run(inputs=inputs, - outputs=outputs, - stream=stream.cuda_stream) - stream.synchronize() - profiler.stop('TRT Plugin') - assert ok - profiler.summary() - - -if __name__ == '__main__': - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter) - parser.add_argument('--batch_size', type=int, default=4) - parser.add_argument('--seq_len', type=int, default=128) - parser.add_argument('--num_heads', type=int, default=8) - parser.add_argument('--head_size', type=int, default=64) - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument( - '--engine_dir', - type=Path, - default='outputs', - help='The directory where serialized engine files locate.') - parser.add_argument( - '--benchmark', - action='store_true', - help='Do performance benchmark compared to triton baseline.') - args = parser.parse_args() - - logger.set_level(args.log_level) - logger.info('Parameters'.center(40, '=')) - for k, v in vars(args).items(): - logger.info(f' - {k.ljust(15, ".")}: {v}') - logger.info(''.center(40, '=')) - - assert args.engine_dir.exists(), \ - f"Engine file {str(args.engine_dir)} doesn't exists." - - logger.info('Inference using the built TensorRT engine.') - run(args.engine_dir, - args.batch_size, - args.seq_len, - args.num_heads, - args.head_size, - do_benchmark=args.benchmark) - logger.info('Done.') diff --git a/examples/python_plugin/README.md b/examples/python_plugin/README.md deleted file mode 100644 index 8079d381109a..000000000000 --- a/examples/python_plugin/README.md +++ /dev/null @@ -1,120 +0,0 @@ -# TensorRT LLM Python Plugin - -TensorRT LLM provides a Python plugin interface to integrate TensorRT LLM with pure Python. - -+ `openai_triton_plugin`: plugin package -+ `build_lookup.py`: Build a TensorRT engine with TensorRT LLM Python plugin -+ `run_lookup.py`: Run the engine and compare the result with PyTorch - -## Plugin Definition - -The following code shows how to create a look-up plugin. -We only need to do a few things to define a TensorRT LLM plugin. - -1. Inherit the `PluginBase`. -2. Register the plugin class to TensorRT LLM by using `@trtllm_plugin("your_plugin_name")`. -3. Define an `__init__` function and initialize the base class. -4. Define a shape and dtype inference function. -5. Define the compute flow. - -```python -@trtllm_plugin("TritonLookUp") -class LookUpPlugin(PluginBase): - - def __init__(self, use_torch_tensor, fp32_output): - super().__init__() - self.use_torch_tensor = use_torch_tensor - self.fp32_output = fp32_output - - def shape_dtype_inference(self, inputs: Sequence[SymTensor]) -> SymTensor: - shape = inputs[1].shape - shape[0] = inputs[0].shape[0] + inputs[1].shape[0] - inputs[1].shape[0] - return SymTensor( - inputs[1].dtype if not self.fp32_output else torch.float32, shape) - - def forward(self, inputs: Sequence[TensorWrapper], - outputs: Sequence[TensorWrapper]): - assert len(inputs) == 2 - assert inputs[0].dtype in [torch.int32 or torch.int64] - assert inputs[1].dtype in [torch.float32, torch.float16, torch.bfloat16] - assert (self.fp32_output and outputs[0].dtype - == torch.float32) or outputs[0].dtype == inputs[1].dtype - - x = inputs[0] - y = inputs[1] - z = outputs[0] - if self.use_torch_tensor: - x = convert_to_torch_tensor(x) - y = convert_to_torch_tensor(y) - z = convert_to_torch_tensor(z) - MAX_BLOCK_NUM = 65536 - MAX_BLOCK_SIZE = 512 - grid = lambda meta: (min(MAX_BLOCK_NUM, x.shape[0]) * min( - MAX_BLOCK_SIZE, y.shape[1]), ) - lookup_kernel[grid](x, y, z, y.shape[0], y.shape[1], x.shape[0]) - -``` - -## Adding a TensorRT LLM Plugin to a Network - -You only need an instance of the plugin object and then call it with `tensorrt_llm.Tensor` as input arguments. - -```python -builder = tensorrt_llm.Builder() -network = builder.create_network() -with tensorrt_llm.net_guard(network): - x = Tensor(name='x', - shape=index_shape, - dtype=tensorrt_llm.str_dtype_to_trt('int32')) - y = Tensor(name='y', - shape=(vocab_size, n_embed), - dtype=torch_dtype_to_trt(dtype)) - - def lookup(x, y): - lookup_plugin = LookUpPlugin(False) - return lookup_plugin(x, y) - - output = lookup(x, y) - output.mark_output('output', torch_dtype_to_str(dtype)) -``` - -## Plugin Code Structure - -Because TensorRT LLM performs plugin registration when importing the custom TensorRT LLM plugin, there are some code structure conventions to register the plugin at runtime. - -```text -plugin_lib -├──__init__.py -├──lookup_plugin.py -└──lookup_kernel.py -``` - -The `__init__.py` file imports all the plugins in the plugin package. -With this convention, users only need to import the plugin package to register the plugins and do not need to manually import them. - -```python -# __init__.py -from .lookup_plugin import LookUpPlugin - -__all__ = ["LookUpPlugin"] -``` - -## Deserialize an Engine with TensorRT LLM Plugin - -During deserialization, TensorRT needs to find the user-defined plugin. Thus, we need to import the plugin once to register them. If the plugin follows the code structure convention, users only need to import that package to register all the custom plugins. - -```python -from tensorrt_llm.runtime.session import Session, TensorInfo - -import openai_triton_plugin # isort: skip - -if __name__ == "__main__": - - def run_engine(dtype): - output_dir = Path('tmp') / torch_dtype_to_str(dtype) - - engine_path = output_dir / "lookup.engine" - - with engine_path.open('rb') as f: - session = Session.from_serialized_engine(f.read()) -``` diff --git a/examples/python_plugin/build_lookup.py b/examples/python_plugin/build_lookup.py deleted file mode 100644 index 48cccd58882e..000000000000 --- a/examples/python_plugin/build_lookup.py +++ /dev/null @@ -1,61 +0,0 @@ -from pathlib import Path - -import torch -from plugin_lib import LookUpPlugin - -import tensorrt_llm -from tensorrt_llm import Tensor -from tensorrt_llm._utils import torch_dtype_to_str, torch_dtype_to_trt - -if __name__ == "__main__": - - # meta data - batch_size = 10 - vocab_size = 1000 - n_embed = 1024 - - # test data - ## input index - index_shape = (batch_size, ) - index_data = torch.randint(0, vocab_size, index_shape, - dtype=torch.int32).cuda() - - def test(dtype): - builder = tensorrt_llm.Builder() - builder.strongly_typed = True - network = builder.create_network() - with tensorrt_llm.net_guard(network): - x = Tensor( - name="x", - shape=index_shape, - dtype=tensorrt_llm.str_dtype_to_trt("int32"), - ) - y = Tensor(name="y", - shape=(vocab_size, n_embed), - dtype=torch_dtype_to_trt(dtype)) - - def lookup(x, y): - lookup_plugin = LookUpPlugin(False, True) - return lookup_plugin(x, y) - - output = lookup(x, y) - - output.mark_output("output", torch_dtype_to_str(torch.float32)) - - builder_config = builder.create_builder_config("float32") - engine = builder.build_engine(network, builder_config) - assert engine is not None - - output_dir = Path("tmp") / torch_dtype_to_str(dtype) - output_dir.mkdir(parents=True, exist_ok=True) - - engine_path = output_dir / "lookup.engine" - config_path = output_dir / "config.json" - - with engine_path.open("wb") as f: - f.write(engine) - builder.save_config(builder_config, str(config_path)) - - test(torch.bfloat16) - test(torch.float16) - test(torch.float32) diff --git a/examples/python_plugin/plugin_lib/__init__.py b/examples/python_plugin/plugin_lib/__init__.py deleted file mode 100644 index f27e0ded3a1c..000000000000 --- a/examples/python_plugin/plugin_lib/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .lookup_plugin import LookUpPlugin - -__all__ = ["LookUpPlugin"] diff --git a/examples/python_plugin/plugin_lib/lookup_kernel.py b/examples/python_plugin/plugin_lib/lookup_kernel.py deleted file mode 100644 index 25cf66704b96..000000000000 --- a/examples/python_plugin/plugin_lib/lookup_kernel.py +++ /dev/null @@ -1,17 +0,0 @@ -import triton -import triton.language as tl - - -@triton.jit -def lookup_kernel(X, Y, Z, vocab_size, hidden_size, token_num): - pid = tl.program_id(axis=0) - while pid < token_num * hidden_size: - row_idx = pid // hidden_size - col_idx = pid % hidden_size - word_idx = tl.load(X + row_idx) - embedding = tl.load( - Y + word_idx * hidden_size + col_idx, - mask=word_idx < vocab_size, - ) - tl.store(Z + pid, embedding) - pid += tl.num_programs(0) diff --git a/examples/python_plugin/plugin_lib/lookup_plugin.py b/examples/python_plugin/plugin_lib/lookup_plugin.py deleted file mode 100644 index 612ee7a4db1e..000000000000 --- a/examples/python_plugin/plugin_lib/lookup_plugin.py +++ /dev/null @@ -1,60 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 1993-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. - -from typing import Sequence - -import torch - -from tensorrt_llm import PluginBase -from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor -from tensorrt_llm.python_plugin import SymTensor, trtllm_plugin - -from .lookup_kernel import lookup_kernel - - -@trtllm_plugin("TritonLookUp") -class LookUpPlugin(PluginBase): - - def __init__(self, use_torch_tensor, fp32_output): - super().__init__() - self.use_torch_tensor = use_torch_tensor - self.fp32_output = fp32_output - - def shape_dtype_inference(self, inputs: Sequence[SymTensor]) -> SymTensor: - shape = inputs[1].shape - shape[0] = inputs[0].shape[0] + inputs[1].shape[0] - inputs[1].shape[0] - return SymTensor( - inputs[1].dtype if not self.fp32_output else torch.float32, shape) - - def forward(self, inputs: Sequence[TensorWrapper], - outputs: Sequence[TensorWrapper]): - assert len(inputs) == 2 - assert inputs[0].dtype in [torch.int32 or torch.int64] - assert inputs[1].dtype in [torch.float32, torch.float16, torch.bfloat16] - assert (self.fp32_output and outputs[0].dtype - == torch.float32) or outputs[0].dtype == inputs[1].dtype - - x = inputs[0] - y = inputs[1] - z = outputs[0] - if self.use_torch_tensor: - x = convert_to_torch_tensor(x) - y = convert_to_torch_tensor(y) - z = convert_to_torch_tensor(z) - MAX_BLOCK_NUM = 65536 - MAX_BLOCK_SIZE = 512 - grid = lambda meta: (min(MAX_BLOCK_NUM, x.shape[0]) * min( - MAX_BLOCK_SIZE, y.shape[1]), ) - lookup_kernel[grid](x, y, z, y.shape[0], y.shape[1], x.shape[0]) diff --git a/examples/python_plugin/run_lookup.py b/examples/python_plugin/run_lookup.py deleted file mode 100644 index 055e86008c07..000000000000 --- a/examples/python_plugin/run_lookup.py +++ /dev/null @@ -1,65 +0,0 @@ -from pathlib import Path - -import torch - -from tensorrt_llm import logger -from tensorrt_llm._utils import (torch_dtype_to_str, torch_dtype_to_trt, - trt_dtype_to_torch) -from tensorrt_llm.runtime.session import Session, TensorInfo - -import plugin_lib # isort: skip - -if __name__ == "__main__": - - def run_engine(dtype): - output_dir = Path('tmp') / torch_dtype_to_str(dtype) - - engine_path = output_dir / "lookup.engine" - - with engine_path.open('rb') as f: - session = Session.from_serialized_engine(f.read()) - - # meta data - batch_size = 10 - vocab_size = 1000 - n_embed = 1024 - - # test data - ## input index - index_shape = (batch_size, ) - index_data = torch.randint(0, - vocab_size, - index_shape, - dtype=torch.int32).cuda() - weight_data = torch.rand(vocab_size, n_embed, dtype=dtype).cuda() - - inputs = {"x": index_data, "y": weight_data} - - output_info = session.infer_shapes([ - TensorInfo(name, torch_dtype_to_trt(tensor.dtype), tensor.shape) - for name, tensor in inputs.items() - ]) - logger.debug(f'output info {output_info}') - outputs = { - t.name: - torch.empty(tuple(t.shape), - dtype=trt_dtype_to_torch(t.dtype), - device='cuda') - for t in output_info - } - - stream = torch.cuda.Stream() - ok = session.run(inputs=inputs, - outputs=outputs, - stream=stream.cuda_stream) - assert ok, 'Engine execution failed' - - embedding = torch.nn.Embedding.from_pretrained(weight_data) - torch_out = embedding(index_data).to(torch.float32) - trt_out = outputs['output'] - - torch.testing.assert_close(trt_out, torch_out) - - run_engine(torch.bfloat16) - run_engine(torch.float16) - run_engine(torch.float32) diff --git a/examples/run.py b/examples/run.py deleted file mode 100755 index 7ce36bbe9848..000000000000 --- a/examples/run.py +++ /dev/null @@ -1,710 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. - -import argparse -import ast -import csv -import os -from pathlib import Path -from typing import List, Optional - -import numpy as np -import torch -from utils import (DEFAULT_HF_MODEL_DIRS, DEFAULT_PROMPT_TEMPLATES, - add_common_args, get_beam_width_array, load_tokenizer, - prepare_enc_dec_inputs, read_model_name, - supports_inflight_batching, throttle_generator) - -import tensorrt_llm -import tensorrt_llm.profiler -from tensorrt_llm.logger import logger -from tensorrt_llm.runtime import PYTHON_BINDINGS, ModelRunner - -if PYTHON_BINDINGS: - from tensorrt_llm.runtime import ModelRunnerCpp - -from ngram.run_dtm_ngram import run_dtm_ngram - - -def parse_arguments(args=None): - # see `add_common_args` for extended list of arguments - parser = argparse.ArgumentParser() - parser.add_argument('--max_input_length', type=int, default=923) - parser.add_argument('--max_output_len', type=int, required=True) - parser.add_argument( - '--draft_engine_dir', - type=str, - default=None, - help='Path to engine of draft model in Draft-Target-Model mode.') - parser.add_argument( - '--input_text', - type=str, - nargs='+', - default=["Born in north-east France, Soyer trained as a"]) - parser.add_argument( - '--input_file', - type=str, - help= - 'CSV or Numpy file containing tokenized input. Alternative to text input.', - default=None) - parser.add_argument('--multimodal_input_file', - type=str, - help='Path to multimodal input file.') - parser.add_argument( - '--input_token_extra_ids', - type=int, - nargs='+', - help= - 'Input token extra ids for using p-tuning and KV Cache reuse together (only available with cpp session).', - default=None) - parser.add_argument( - '--input_token_extra_ids_file', - type=str, - help= - 'CSV or Numpy file containing input token extra ids file. Alternative to text input (only available with cpp session).', - default=None) - parser.add_argument('--output_csv', - type=str, - help='CSV file where the tokenized output is stored.', - default=None) - parser.add_argument('--output_npy', - type=str, - help='Numpy file where the tokenized output is stored.', - default=None) - parser.add_argument('--output_generation_logits', - default=False, - action='store_true', - help="Enable gathering generation logits.") - parser.add_argument( - '--output_logits_npy', - type=str, - help= - 'Numpy file where the generation logits are stored. Use only when num_beams==1', - default=None) - parser.add_argument('--output_log_probs_npy', - type=str, - help='Numpy file where the log_probs are stored', - default=None) - parser.add_argument('--output_cum_log_probs_npy', - type=str, - help='Numpy file where the cum_log_probs are stored', - default=None) - parser.add_argument( - '--run_profiling', - default=False, - action='store_true', - help="Run several 10 iterations to profile the inference latencies.") - parser.add_argument( - '--fail_fast_on_attention_window_too_large', - action='store_true', - default=False, - help= - 'Exit with runtime error when attention window is too large to fit even a single sequence in the KV cache.' - ) - - parser = add_common_args(parser) - - return parser.parse_args(args=args) - - -def parse_input(tokenizer, - input_text=None, - prompt_template=None, - input_file=None, - add_special_tokens=True, - max_input_length=923, - pad_id=None, - num_prepend_vtokens=[], - model_name=None, - model_version=None): - if pad_id is None: - pad_id = tokenizer.pad_token_id - - batch_input_ids = [] - if input_file is None: - if 'whisper' in model_name.lower(): - batch_input_ids.append(tokenizer.prefix_tokens) - else: - for curr_text in input_text: - if prompt_template is not None: - curr_text = prompt_template.format(input_text=curr_text) - input_ids = tokenizer.encode( - curr_text, - add_special_tokens=add_special_tokens, - truncation=True, - max_length=max_input_length) - batch_input_ids.append(input_ids) - else: - if input_file.endswith('.csv'): - with open(input_file, 'r') as csv_file: - csv_reader = csv.reader(csv_file, delimiter=',') - for line in csv_reader: - input_ids = np.array(line, dtype='int32') - batch_input_ids.append(input_ids[-max_input_length:]) - elif input_file.endswith('.npy'): - inputs = np.load(input_file) - for row in inputs: - input_ids = row[row != pad_id] - batch_input_ids.append(input_ids[-max_input_length:]) - - elif input_file.endswith('.txt'): - with open(input_file, 'r', encoding='utf-8', - errors='replace') as txt_file: - input_text = txt_file.readlines() - batch_input_ids = tokenizer( - input_text, - add_special_tokens=add_special_tokens, - truncation=True, - max_length=max_input_length)["input_ids"] - else: - print('Input file format not supported.') - raise SystemExit - - if num_prepend_vtokens: - assert len(num_prepend_vtokens) == len(batch_input_ids) - base_vocab_size = tokenizer.vocab_size - for i, length in enumerate(num_prepend_vtokens): - batch_input_ids[i] = list( - range(base_vocab_size, - base_vocab_size + length)) + batch_input_ids[i] - - if input_file is None and 'GLM' in model_name and model_version == 'glm': - for ids in batch_input_ids: - ids.append(tokenizer.sop_token_id) - - batch_input_ids = [ - torch.tensor(x, dtype=torch.int32) for x in batch_input_ids - ] - - logger.debug(f"Input token ids (batch_size = {len(batch_input_ids)}):") - for i, input_ids in enumerate(batch_input_ids): - logger.debug(f"Request {i}: {input_ids.tolist()}") - - return batch_input_ids - - -def parse_input_token_extra_ids(prompt_table_path, kv_cache_enable_block_reuse, - input_token_extra_ids, - input_token_extra_ids_file, max_input_length): - batch_extra_ids = None - if prompt_table_path and kv_cache_enable_block_reuse: - assert input_token_extra_ids or input_token_extra_ids_file, \ - "Input token extra ids must be provided when p-tuning and KV Cache reuse are both enabled" - batch_extra_ids = [] - if input_token_extra_ids_file: - if input_token_extra_ids_file.endswith('.csv'): - with open(input_token_extra_ids_file, 'r') as csv_file: - csv_reader = csv.reader(csv_file, delimiter=',') - for line in csv_reader: - extra_ids = [int(num) for num in line] - batch_extra_ids.append(extra_ids[-max_input_length:]) - elif input_token_extra_ids_file.endswith('.npy'): - inputs = np.load(input_token_extra_ids_file) - for extra_ids in inputs: - batch_extra_ids.append(extra_ids[-max_input_length:]) - else: - print('Input file format not supported.') - raise SystemExit - else: - batch_extra_ids.append(input_token_extra_ids) - return batch_extra_ids - - -def print_output(tokenizer, - output_ids: torch.Tensor, - input_lengths: List[int], - sequence_lengths: torch.Tensor, - output_csv: Optional[str] = None, - output_npy: Optional[str] = None, - context_logits: Optional[torch.Tensor] = None, - generation_logits: Optional[torch.Tensor] = None, - cum_log_probs: Optional[torch.Tensor] = None, - log_probs: Optional[torch.Tensor] = None, - output_logits_npy: Optional[str] = None, - output_cum_log_probs_npy: Optional[str] = None, - output_log_probs_npy: Optional[str] = None): - num_output_sents, num_beams, _ = output_ids.size() - batch_size = len(input_lengths) - num_return_sequences = num_output_sents // batch_size - - if output_csv is None and output_npy is None and tokenizer is not None: - for i in range(batch_size * num_return_sequences): - batch_idx = i // num_return_sequences - seq_idx = i % num_return_sequences - inputs = output_ids[i][0][:input_lengths[batch_idx]].tolist() - input_text = tokenizer.decode(inputs) - if seq_idx == 0: - print(f'Input [Text {batch_idx}]: \"{input_text}\"') - - for beam in range(num_beams): - output_begin = input_lengths[batch_idx] - output_end = sequence_lengths[i][beam] - outputs = output_ids[i][beam][output_begin:output_end].tolist() - output_text = tokenizer.decode(outputs) - index_str = (f'Text {batch_idx} Seq {seq_idx} Beam {beam}' - if num_return_sequences > 1 else - f'Text {batch_idx} Beam {beam}') - print(f'Output [{index_str}]: \"{output_text}\"') - logger.debug(str(outputs)) - - output_ids = output_ids.reshape((-1, output_ids.size(2))) - - if output_csv is not None: - output_file = Path(output_csv) - output_file.parent.mkdir(exist_ok=True, parents=True) - outputs = output_ids.tolist() - with open(output_file, 'w') as csv_file: - writer = csv.writer(csv_file, delimiter=',') - writer.writerows(outputs) - - if output_npy is not None: - output_file = Path(output_npy) - output_file.parent.mkdir(exist_ok=True, parents=True) - outputs = np.array(output_ids.cpu().contiguous(), dtype='int32') - np.save(output_file, outputs) - - # Save context logits - if context_logits is not None and output_logits_npy is not None: - context_logits = torch.cat(context_logits, axis=0) - vocab_size_padded = context_logits.shape[-1] - context_logits = context_logits.reshape([1, -1, vocab_size_padded]) - - output_context_logits_npy = output_logits_npy.split( - '.npy')[0] + "_context" - output_context_logits_file = Path(output_context_logits_npy) - context_outputs = np.array( - context_logits.squeeze(0).cpu().contiguous(), - dtype='float32') # [promptLengthSum, vocabSize] - np.save(output_context_logits_file, context_outputs) - - # Save generation logits - if generation_logits is not None and output_logits_npy is not None and num_beams == 1: - output_generation_logits_npy = output_logits_npy.split( - '.npy')[0] + "_generation" - output_generation_logits_file = Path(output_generation_logits_npy) - generation_outputs = np.array(generation_logits.cpu().contiguous(), - dtype='float32') - np.save(output_generation_logits_file, generation_outputs) - - # Save cum log probs - if cum_log_probs is not None and output_cum_log_probs_npy is not None: - cum_log_probs_file = Path(output_cum_log_probs_npy) - cum_log_probs_outputs = np.array(cum_log_probs.cpu().contiguous(), - dtype='float32') - np.save(cum_log_probs_file, cum_log_probs_outputs) - - # Save cum log probs - if log_probs is not None and output_log_probs_npy is not None: - log_probs_file = Path(output_log_probs_npy) - log_probs_outputs = np.array(log_probs.cpu().contiguous(), - dtype='float32') - np.save(log_probs_file, log_probs_outputs) - - -def main(args): - runtime_rank = tensorrt_llm.mpi_rank() - logger.set_level(args.log_level) - - # different handling if encoder-decoder models - is_enc_dec = {'encoder', 'decoder'}.issubset({ - name - for name in os.listdir(args.engine_dir) - if os.path.isdir(os.path.join(args.engine_dir, name)) - }) - if is_enc_dec: - logger.warning( - "This path is an encoder-decoder model. Using different handling.") - assert not args.use_py_session, "Encoder-decoder models don't have a unified python runtime, please use its own examples/models/core/enc_dec/run.py instead." - - model_name, model_version = read_model_name( - args.engine_dir if not is_enc_dec else os.path. - join(args.engine_dir, 'encoder')) - - if args.tokenizer_dir is None and model_name in DEFAULT_HF_MODEL_DIRS: - logger.warning( - "tokenizer_dir is not specified. Try to infer from model_name, but this may be incorrect." - ) - args.tokenizer_dir = DEFAULT_HF_MODEL_DIRS[model_name] - - tokenizer, pad_id, end_id = load_tokenizer( - tokenizer_dir=args.tokenizer_dir, - vocab_file=args.vocab_file, - model_name=model_name, - model_version=model_version, - tokenizer_type=args.tokenizer_type, - ) - - if args.end_id: - end_id = args.end_id - - prompt_template = None - if args.use_prompt_template and model_name in DEFAULT_PROMPT_TEMPLATES: - prompt_template = DEFAULT_PROMPT_TEMPLATES[model_name] - - batch_input_ids = parse_input(tokenizer=tokenizer, - input_text=args.input_text, - prompt_template=prompt_template, - input_file=args.input_file, - add_special_tokens=args.add_special_tokens, - max_input_length=args.max_input_length, - pad_id=pad_id, - num_prepend_vtokens=args.num_prepend_vtokens, - model_name=model_name, - model_version=model_version) - - stop_words_list = None - if args.stop_words: - stop_words_list = tensorrt_llm.runtime.decode_words_list( - args.stop_words, tokenizer) - if model_version == 'glm4': # add default stop token ids for GLM-4 - glm4_stop_ids = [[151329], [151336], [151338]] - if stop_words_list is None: - stop_words_list = [glm4_stop_ids] * len(batch_input_ids) - else: - for req_stop_words_list in stop_words_list: - req_stop_words_list.extend(glm4_stop_ids) - - bad_words_list = None - if args.bad_words: - bad_words_list = tensorrt_llm.runtime.decode_words_list( - args.bad_words, tokenizer) - - if is_enc_dec: - encoder_input_ids, encoder_input_features, encoder_output_lengths, decoder_input_ids = prepare_enc_dec_inputs( - batch_input_ids, model_name, args.engine_dir, - args.multimodal_input_file) - - input_token_extra_ids = parse_input_token_extra_ids( - args.prompt_table_path, args.kv_cache_enable_block_reuse, - args.input_token_extra_ids, args.input_token_extra_ids_file, - args.max_input_length) - - input_lengths = [x.size(0) for x in decoder_input_ids - ] if is_enc_dec else [x.size(0) for x in batch_input_ids] - - encoder_input_lengths = [ - x.size(0) for x in (encoder_input_features or encoder_input_ids) - ] if is_enc_dec else None - - if args.beam_width_array is not None: - logger.info("Enable Variable-Beam-Width-Search (VBWS)") - assert not args.use_py_session, "`--use_py_session` is not supported in VBWS." - args.beam_width_array, args.num_beams = get_beam_width_array( - args.beam_width_array) - - if not args.use_py_session and not supports_inflight_batching( - os.path.join(args.engine_dir, "decoder") if is_enc_dec else args. - engine_dir): - logger.warning( - "The given engine does not support in-flight batching, fallback to python session" - ) - args.use_py_session = True - - if not PYTHON_BINDINGS and not args.use_py_session: - logger.warning( - "Python bindings of C++ session is unavailable, fallback to Python session." - ) - args.use_py_session = True - if args.debug_mode and not args.use_py_session: - logger.warning( - "Debug mode is not supported in C++ session for now, fallback to Python session." - ) - args.use_py_session = True - if args.return_all_generated_tokens and args.use_py_session: - raise ValueError( - "Returning all the generated tokens at each step is not supported in the Python session, use C++ session instead." - ) - if (not args.return_all_generated_tokens) and args.streaming and ( - args.num_beams > 1): - logger.warning( - "Setting return_all_generated_tokens to True since streaming AND beam search are done simultaneously. " - "Returning the full beams at each streaming step is needed because beam search + streaming can change previous outputs. " - "WARNING: using this option may increase network usage significantly (quadratically w.r.t output length)." - ) - args.return_all_generated_tokens = True - - logger.info(f"Using {'Python' if args.use_py_session else 'C++'} session") - - if args.draft_target_model_config is not None or args.ngram_config is not None: - # Speculative-Decoding of Draft-Target-Model (DTM) and NGram - # If the parameters of `runner_kwargs` and `runner.generate()` in the "else" branch change, the same change should be done for `examples/ngram/run_dtm_ngram.py` - assert args.kv_cache_enable_block_reuse, "`--kv_cache_enable_block_reuse` must be specified in speculative decoding." - assert not args.use_py_session, "`--use_py_session` is not supported in Speculative decoding." - assert not is_enc_dec, "Encoder-Decoder model is not supported in Speculative decoding." - assert args.num_beams == 1, "`--num_beams>1` is not supported in Speculative decoding." - - outputs = run_dtm_ngram(batch_input_ids, args, runtime_rank, end_id, - pad_id, stop_words_list, bad_words_list, - len(tokenizer)) - if not args.streaming: # Unpack runner from the return value in No-Streaming mode - outputs, runner = list(outputs)[0] - - else: # Normal run - runner_cls = ModelRunner if args.use_py_session else ModelRunnerCpp - runner_kwargs = dict( - engine_dir=args.engine_dir, - lora_dir=args.lora_dir, - rank=runtime_rank, - debug_mode=args.debug_mode, - lora_ckpt_source=args.lora_ckpt_source, - gpu_weights_percent=args.gpu_weights_percent, - max_output_len=args.max_output_len, - enable_context_fmha_fp32_acc=args.enable_context_fmha_fp32_acc, - fail_fast_on_attention_window_too_large=args. - fail_fast_on_attention_window_too_large, - ) - if args.medusa_choices is not None: - args.medusa_choices = ast.literal_eval(args.medusa_choices) - assert args.temperature == 1.0, "Medusa should use temperature == 1.0" - assert args.num_beams == 1, "Medusa should use num_beams == 1" - runner_kwargs.update(medusa_choices=args.medusa_choices) - if args.eagle_choices is not None or args.eagle_posterior_threshold is not None or args.eagle_use_dynamic_tree: - assert args.num_beams == 1, "Eagle should use num_beams == 1" - assert not args.use_py_session, "Eagle does not support py session" - if args.eagle_choices is not None and not args.eagle_use_dynamic_tree: - args.eagle_choices = ast.literal_eval(args.eagle_choices) - runner_kwargs.update(eagle_choices=args.eagle_choices) - if args.eagle_posterior_threshold is not None: - runner_kwargs.update( - eagle_posterior_threshold=args.eagle_posterior_threshold) - if args.eagle_use_dynamic_tree: - runner_kwargs.update( - eagle_use_dynamic_tree=args.eagle_use_dynamic_tree) - assert args.eagle_dynamic_tree_max_top_k is not None and args.eagle_dynamic_tree_max_top_k > 0 - runner_kwargs.update(eagle_dynamic_tree_max_top_k=args. - eagle_dynamic_tree_max_top_k) - if args.lookahead_config is not None: - args.lookahead_config = ast.literal_eval(args.lookahead_config) - assert len( - args.lookahead_config - ) == 3, "Lookahead needs [max_window_size, max_ngram_size, max_verification_set_size]" - runner_kwargs.update(lookahead_config=args.lookahead_config) - if not args.use_py_session: - runner_kwargs.update( - is_enc_dec=is_enc_dec, - max_batch_size=len(batch_input_ids), - max_input_len=max( - encoder_input_lengths if is_enc_dec else input_lengths), - max_beam_width=args.num_beams, - max_attention_window_size=args.max_attention_window_size, - sink_token_length=args.sink_token_length, - max_tokens_in_paged_kv_cache=args.max_tokens_in_paged_kv_cache, - kv_cache_enable_block_reuse=args.kv_cache_enable_block_reuse, - kv_cache_free_gpu_memory_fraction=args. - kv_cache_free_gpu_memory_fraction, - cross_kv_cache_fraction=args.cross_kv_cache_fraction - if is_enc_dec else None, - enable_chunked_context=args.enable_chunked_context, - multi_block_mode=args.multi_block_mode, - cuda_graph_mode=args.cuda_graph_mode, - gather_generation_logits=args.output_generation_logits, - use_variable_beam_width_search=(args.beam_width_array - is not None), - ) - runner = runner_cls.from_dir(**runner_kwargs) - - with torch.no_grad(): - outputs = runner.generate( - batch_input_ids=decoder_input_ids - if is_enc_dec else batch_input_ids, - encoder_input_ids=encoder_input_ids if is_enc_dec else None, - encoder_input_features=encoder_input_features - if is_enc_dec else None, - encoder_output_lengths=encoder_output_lengths - if is_enc_dec else None, - max_new_tokens=args.max_output_len, - max_attention_window_size=args.max_attention_window_size, - sink_token_length=args.sink_token_length, - end_id=end_id, - pad_id=pad_id, - temperature=args.temperature, - top_k=args.top_k, - top_p=args.top_p, - num_beams=args.num_beams, - num_return_sequences=args.num_return_sequences, - length_penalty=args.length_penalty, - early_stopping=args.early_stopping, - beam_width_array=args.beam_width_array, - repetition_penalty=args.repetition_penalty, - presence_penalty=args.presence_penalty, - frequency_penalty=args.frequency_penalty, - prompt_ignore_length=args.prompt_ignore_length, - min_p=args.min_p, - stop_words_list=stop_words_list, - bad_words_list=bad_words_list, - output_cum_log_probs=(args.output_cum_log_probs_npy != None), - output_log_probs=(args.output_log_probs_npy != None), - random_seed=args.random_seed, - lora_uids=args.lora_task_uids, - prompt_table=args.prompt_table_path, - prompt_tasks=args.prompt_tasks, - streaming=args.streaming, - output_sequence_lengths=True, - output_generation_logits=args.output_generation_logits, - no_repeat_ngram_size=args.no_repeat_ngram_size, - return_dict=True, - medusa_choices=args.medusa_choices, - eagle_choices=args.eagle_choices, - return_all_generated_tokens=args.return_all_generated_tokens, - input_token_extra_ids=input_token_extra_ids, - fail_fast_on_attention_window_too_large=args. - fail_fast_on_attention_window_too_large, - language_adapter_uids=args.language_task_uids) - torch.cuda.synchronize() - - # Receive output, print to screen or save to file - if args.streaming: - for curr_outputs in throttle_generator(outputs, - args.streaming_interval): - if runtime_rank == 0: - output_ids = curr_outputs['output_ids'] - sequence_lengths = curr_outputs['sequence_lengths'] - cum_log_probs = None - log_probs = None - if args.output_cum_log_probs_npy is not None: - cum_log_probs = curr_outputs['cum_log_probs'] - if args.output_log_probs_npy is not None: - log_probs = curr_outputs['log_probs'] - print_output( - tokenizer, - output_ids, - input_lengths, - sequence_lengths, - output_csv=args.output_csv, - output_npy=args.output_npy, - cum_log_probs=cum_log_probs, - log_probs=log_probs, - output_cum_log_probs_npy=args.output_cum_log_probs_npy, - output_log_probs_npy=args.output_log_probs_npy) - else: - if runtime_rank == 0: - output_ids = outputs['output_ids'] - sequence_lengths = outputs['sequence_lengths'] - context_logits = None - generation_logits = None - cum_log_probs = None - log_probs = None - if runner.gather_context_logits: - context_logits = outputs['context_logits'] - if runner.gather_generation_logits or args.output_generation_logits: - generation_logits = outputs['generation_logits'] - if args.output_cum_log_probs_npy is not None: - cum_log_probs = outputs['cum_log_probs'] - if args.output_log_probs_npy is not None: - log_probs = outputs['log_probs'] - print_output(tokenizer, - output_ids, - input_lengths, - sequence_lengths, - output_csv=args.output_csv, - output_npy=args.output_npy, - context_logits=context_logits, - generation_logits=generation_logits, - output_logits_npy=args.output_logits_npy, - cum_log_probs=cum_log_probs, - log_probs=log_probs, - output_cum_log_probs_npy=args.output_cum_log_probs_npy, - output_log_probs_npy=args.output_log_probs_npy) - - # Profiling - if args.run_profiling: - ite = 10 - # warmup - for _ in range(ite): - with torch.no_grad(): - outputs = runner.generate( - batch_input_ids, - max_new_tokens=args.max_output_len, - max_attention_window_size=args.max_attention_window_size, - end_id=end_id, - pad_id=pad_id, - temperature=args.temperature, - top_k=args.top_k, - top_p=args.top_p, - num_beams=args.num_beams, - length_penalty=args.length_penalty, - early_stopping=args.early_stopping, - beam_width_array=args.beam_width_array, - repetition_penalty=args.repetition_penalty, - presence_penalty=args.presence_penalty, - frequency_penalty=args.frequency_penalty, - prompt_ignore_length=args.prompt_ignore_length, - min_p=args.min_p, - stop_words_list=stop_words_list, - bad_words_list=bad_words_list, - output_cum_log_probs=(args.output_cum_log_probs_npy - is not None), - output_log_probs=(args.output_log_probs_npy is not None), - random_seed=args.random_seed, - lora_uids=args.lora_task_uids, - lookahead_config=args.lookahead_config, - prompt_table=args.prompt_table_path, - prompt_tasks=args.prompt_tasks, - streaming=args.streaming, - output_sequence_lengths=True, - return_dict=True, - return_all_generated_tokens=args. - return_all_generated_tokens, - input_token_extra_ids=input_token_extra_ids) - torch.cuda.synchronize() - - tensorrt_llm.profiler.start("tmp") - for _ in range(ite): - with torch.no_grad(): - outputs = runner.generate( - batch_input_ids, - max_new_tokens=args.max_output_len, - max_attention_window_size=args.max_attention_window_size, - end_id=end_id, - pad_id=pad_id, - temperature=args.temperature, - top_k=args.top_k, - top_p=args.top_p, - num_beams=args.num_beams, - length_penalty=args.length_penalty, - early_stopping=args.early_stopping, - beam_width_array=args.beam_width_array, - repetition_penalty=args.repetition_penalty, - presence_penalty=args.presence_penalty, - frequency_penalty=args.frequency_penalty, - prompt_ignore_length=args.prompt_ignore_length, - stop_words_list=stop_words_list, - bad_words_list=bad_words_list, - output_cum_log_probs=(args.output_cum_log_probs_npy - != None), - output_log_probs=(args.output_log_probs_npy != None), - random_seed=args.random_seed, - lora_uids=args.lora_task_uids, - prompt_table=args.prompt_table_path, - prompt_tasks=args.prompt_tasks, - streaming=args.streaming, - output_sequence_lengths=True, - return_dict=True, - return_all_generated_tokens=args. - return_all_generated_tokens, - input_token_extra_ids=input_token_extra_ids, - fail_fast_on_attention_window_too_large=args. - fail_fast_on_attention_window_too_large) - torch.cuda.synchronize() - tensorrt_llm.profiler.stop("tmp") - - print( - f"batch_size: {len(batch_input_ids)}, avg latency of {ite} iterations: : {tensorrt_llm.profiler.elapsed_time_in_sec('tmp') / ite} sec" - ) - - -if __name__ == '__main__': - args = parse_arguments() - main(args) diff --git a/examples/sample_weight_stripping/README.md b/examples/sample_weight_stripping/README.md deleted file mode 100644 index cb3c04404902..000000000000 --- a/examples/sample_weight_stripping/README.md +++ /dev/null @@ -1,275 +0,0 @@ -# Sample Weight-Stripping - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -## Table Of Contents - -- [Overview](#overview) - * [Build Weights Stripped Engine](#build-weights-stripped-engine) - * [Engine Refitter](#engine-refitter) -- [Prerequisites](#prerequisites) -- [Weight-Stripping Workflow Example](#weight-stripping-workflow-example) - * [GPT-J](#gpt-j) - * [Llama-7b INT4](#llama-7b-int4) - * [Llama-7b FP16 + WoQ INT8](#llama-7b-fp16-woq-int8) - * [Llama2-70b FP8 with TP=2](#llama2-70b-fp8-with-tp2) -- [Engine Plan File Size Results](#engine-plan-file-size-results) -- [Prototype](#prototype) - * [Checkpoint Pruner](#checkpoint-pruner) - * [Pruning a TensorRT LLM Checkpoint](#pruning-a-tensorrt-llm-checkpoint) - -## Overview - -This workflow introduces a new script `trtllm-refit`. `trtllm-refit` allows you to refit the generated engine with weights from any TensorRT LLM checkpoint matching the same architecture, so long as you build the engine as refittable or stripped. - -### Build Weights Stripped Engine -TensorRT can generate refittable engines with the same performance as the non-refittable ones when TensorRT builder optimize under the assumption that the engine will be refitted with weights identical to those provide at build time. Those refittable weights can be stripped to reduce the engine plan file size, with the option to subsequently supply them via the refit interface. - -New option `--strip_plan` is introduced in `trtllm-build` - -```bash -trtllm-build --strip_plan --checkpoint_dir ${CHECKPOINT_DIR} --output_dir ${ENGINE_DIR} ... -``` - -### Engine Refitter -The refitter allows you to refit an engine with weights in a TensorRT LLM checkpoint. It does this by doing a textual match between engine and checkpoint weight names. In order for the refitter to work, the engine must be built with refitting enabled. This can be accomplished by passing `--strip_plan` to `trtllm-build`. - -After building a stripped engine via `trtllm-build`, run - -```bash -trtllm-refit --checkpoint_dir ${CHECKPOINT_DIR} --engine_dir ${ENGINE_DIR} -``` - - -## Prerequisites - -Install [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM/blob/main/README.md) either through [pip](https://github.com/NVIDIA/TensorRT-LLM/blob/main/README.md#installation) or [from the source](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/installation/build-from-source.md). - -## Weight-Stripping Workflow Example - -### GPT-J - -1. Download the weights. -```bash -# 1. Weights & config -git clone https://huggingface.co/EleutherAI/gpt-j-6b -pushd gpt-j-6b && \ - rm -f pytorch_model.bin && \ - wget https://huggingface.co/EleutherAI/gpt-j-6b/resolve/main/pytorch_model.bin && \ -popd - -# 2. Vocab and merge table -wget https://huggingface.co/EleutherAI/gpt-j-6b/resolve/main/vocab.json -wget https://huggingface.co/EleutherAI/gpt-j-6b/resolve/main/merges.txt -``` - -2. Convert the Hugging Face checkpoint into TensorRT LLM format. -Run below command lines in [`examples/models/contrib/gpt`](../gptj) directory. -```bash -# Build a float16 checkpoint using HF weights. -python convert_checkpoint.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --output_dir ./trt_ckpt/gptj_fp16_tp1/ - -# Build an int8 weight-only checkpoint using HF weights. -python convert_checkpoint.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int8 \ - --output_dir ./trt_ckpt/gptj_int8_tp1/ - -``` - -3. Build the weights stripped engine. -```bash -# Build with --strip_plan. Requires TRT>=10.0.0 -trtllm-build --checkpoint_dir ./trt_ckpt/gptj_fp16_tp1/ \ - --output_dir ./trt_engines/gptj_fp16_tp1/ \ - --gemm_plugin float16 \ - --max_batch_size=32 \ - --max_input_len=1919 \ - --max_seq_len=2047 \ - --strip_plan -``` - -4. Refit the engine. The refit engine lives at `${ENGINE_DIR}.refit`. -```bash -# --checkpoint_dir points to the path of the weights you want refit, in this case the original weights. -trtllm-refit --checkpoint_dir ./trt_ckpt/gptj_fp16_tp1/ --engine_dir ./trt_engines/gptj_fp16_tp1/ --output_dir ./trt_engines/gptj_fp16_tp1.refit/ -``` - -5. Verify the engine. -```bash -# Run the summarization task. -python3 ../summarize.py --engine_dir ./trt_engines/gptj_fp16_tp1.refit \ - --hf_model_dir ./gpt-j-6b \ - --batch_size 1 \ - --test_trt_llm \ - --tensorrt_llm_rouge1_threshold 14 \ - --data_type fp16 \ - --check_accuracy -``` - -### Llama-7b INT4 - -1. Download the llama-7b-hf checkpoint and saved in /llm-models/llama-models/llama-7b-hf/. - -2. Calibrate the checkpoint and convert into TensorRT LLM format. -Run below command lines in [`examples/models/core/llama`](../models/core/llama) directory. -```bash -# Calibrate INT4 using AMMO. -python ../quantization/quantize.py --model_dir /llm-models/llama-models/llama-7b-hf/ \ - --dtype float16 \ - --qformat int4_awq \ - --awq_block_size 128 \ - --output_dir ./quantized_int4-awq \ - --calib_size 32 -``` - -3. Build the weights stripped engine. -```bash -# Build with --strip_plan. Requires TRT>=10.0.0 -trtllm-build --checkpoint_dir ./quantized_int4-awq \ - --strip_plan \ - --gemm_plugin float16 \ - --output_dir trt_int4_AWQ -``` - -4. Refit the engine. -```bash -trtllm-refit --checkpoint_dir ./quantized_int4-awq \ - --engine_dir trt_int4_AWQ \ - --output_dir trt_int4_AWQ_full_from_wtless -``` - -5. Verify the engine. -```bash -python3 ../summarize.py --engine_dir trt_int4_AWQ_full_from_wtless \ - --hf_model_dir /llm-models/llama-models/llama-7b-hf/ \ - --batch_size 1 \ - --test_trt_llm \ - --check_accuracy -``` - -### Llama-7b FP16 + WoQ INT8 - -1. Download the llama-7b-hf checkpoint and saved in /llm-models/llama-models/llama-7b-hf/. - -2. Convert the checkpoint into TensorRT LLM format. -Run below command lines in [`examples/models/core/llama`](../models/core/llama) directory. -```bash -python3 convert_checkpoint.py --model_dir /llm-models/llama-models/llama-7b-hf/ \ - --output_dir ./llama-7b-hf-fp16-woq \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int8 -``` - -3. Build the weights stripped engine. -```bash -# Build with --strip_plan. Requires TRT>=10.0.0 -trtllm-build --checkpoint_dir ./llama-7b-hf-fp16-woq \ - --output_dir ./engines/llama-7b-hf-fp16-woq-1gpu-wtless \ - --strip_plan \ - --gemm_plugin float16 -``` - -4. Refit the engine. -```bash -trtllm-refit --checkpoint_dir ./llama-7b-hf-fp16-woq \ - --engine_dir ./engines/llama-7b-hf-fp16-woq-1gpu-wtless \ - --output_dir ./engines/llama-7b-hf-fp16-woq-1gpu-wtless-to-full -``` - -5. Verify the engine. -```bash -python3 ../summarize.py --engine_dir ./engines/llama-7b-hf-fp16-woq-1gpu-wtless-to-full \ - --hf_model_dir /llm-models/llama-models/llama-7b-hf/ \ - --batch_size 1 \ - --test_trt_llm \ - --check_accuracy -``` - - -### Llama2-70b FP8 with TP=2 - -1. Download the llama-v2-70b-hf checkpoint and saved in /llm-models/llama-models-v2/llama-v2-70b-hf/. - -2. Calibrate the checkpoint and convert into TensorRT LLM format. -Run below command lines in [`examples/models/core/llama`](../models/core/llama) directory. -```bash -# Calibrate FP8 using AMMO. -python ../quantization/quantize.py --model_dir /llm-models/llama-models-v2/llama-v2-70b-hf/ \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir ./llama2-70b-hf-fp8-tp2 \ - --calib_size 512 \ - --tp_size 2 -``` - -3. Build the weights stripped engine. -```bash -trtllm-build --checkpoint_dir ./llama2-70b-hf-fp8-tp2 \ - --output_dir engines/llama2-70b-hf-fp8-tp2 \ - --gemm_plugin float16 \ - --workers 2 -``` - -4. Refit the engine. -```bash -trtllm-refit --checkpoint_dir ./llama2-70b-hf-fp8-tp2 \ - --engine_dir engines/llama2-70b-hf-fp8-tp2 \ - --output_dir engines/llama2-70b-hf-fp8-tp2.refit -``` - -5. Verify the engine. -```bash -python3 ../summarize.py --engine_dir engines/llama2-70b-hf-fp8-tp2.refit \ - --hf_model_dir /llm-models/llama-models-v2/llama-v2-70b-hf/ \ - --batch_size 1 \ - --test_trt_llm \ - --check_accuracy -``` - - -## Engine Plan File Size Results - -| **Model** | **Full Engine Plan Size** | **Weight-Stripped Engine Plan Size** | -|:---------:|:----------:|:----:| -|llama-7b INT4 | 3.7GB | 5.3MB | -|llama-7b FP16 + WoQ INT8 | 6.54GB | 28.69MB | -|llama2-70b FP8 + TP=2 | 64.78GB | 60.61MB | - -## Prototype -### Checkpoint Pruner -The checkpoint pruner allows you to strip `Conv` and `Gemm` weights out of a TensorRT LLM [checkpoint](https://nvidia.github.io/TensorRT-LLM/0.21.0/architecture/checkpoint.html). Since these make up the vast majority of weights, the pruner will decrease the size of your checkpoint up to 99%. - -When building an engine with a pruned checkpoint, TensorRT LLM fills in the missing weights with random ones. These weights should later be [refit](#engine-refitter) with the original weights to preserve the intended behavior. - -Building an engine from a pruned checkpoint will also allow the engine to be [refit](#engine-refitter). - -#### Pruning a TensorRT LLM Checkpoint - -1. Install [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM/blob/main/README.md) either through [pip](https://github.com/NVIDIA/TensorRT-LLM/blob/main/README.md#installation) or [from the source](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/installation/build-from-source.md). -2. Download a model of your choice and convert it to a TensorRT LLM checkpoint ([llama instructions](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/models/core/llama/README.md#usage)). -3. (Optional) Run the `trtllm-prune` command. -```bash -# Prunes the TRT-LLM checkpoint at ${CHECKPOINT_DIR}, and stores it in the directory ${CHECKPOINT_DIR}.pruned -trtllm-prune --checkpoint_dir ${CHECKPOINT_DIR} -``` - -The pruned checkpoint lives at `${CHECKPOINT_DIR}.pruned` by default, however, this can be overridden by issuing the `--out_dir` flag. - -4. Build the stripped engine. - -```bash -# From pruned checkpoint. -trtllm-build --checkpoint_dir ${CHECKPOINT_DIR}.pruned \ - --output_dir ${ENGINE_OUT_DIR} \ - ${EXTRA_ARGS} -``` diff --git a/examples/summarize.py b/examples/summarize.py deleted file mode 100644 index 1f6f8979bb7b..000000000000 --- a/examples/summarize.py +++ /dev/null @@ -1,944 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. - -import argparse -import ast -import os -from pathlib import Path - -import evaluate -import numpy as np -import torch -from datasets import load_dataset -from transformers import (AutoModel, AutoModelForCausalLM, - AutoModelForSeq2SeqLM, GenerationConfig) -from utils import (DEFAULT_HF_MODEL_DIRS, add_common_args, get_beam_width_array, - load_tokenizer, read_model_name, supports_inflight_batching) - -import tensorrt_llm -import tensorrt_llm.profiler as profiler -from tensorrt_llm._utils import mpi_broadcast, str_dtype_to_torch -from tensorrt_llm.builder import EngineConfig -from tensorrt_llm.functional import RopeEmbeddingUtils, RotaryScalingType -from tensorrt_llm.layers import MropeParams -from tensorrt_llm.logger import logger -from tensorrt_llm.models.qwen.utils import make_context -from tensorrt_llm.runtime import PYTHON_BINDINGS, ModelRunner -from tensorrt_llm.tools.ppl import ppl - -if PYTHON_BINDINGS: - from tensorrt_llm.runtime import ModelRunnerCpp - -from ngram.run_dtm_ngram import run_dtm_ngram - - -def ensemble_mrope_params(batch_input_ids, max_position_embeddings, - rotary_embedding_dim, theta): - mrope_params = MropeParams() - batch_size = len(batch_input_ids) - - _, rotary_cos_sin = RopeEmbeddingUtils.create_sinusoidal_positions_for_attention_plugin( - num_pos=max_position_embeddings, - dim=rotary_embedding_dim, - theta=1000000.0, - scale_type=RotaryScalingType.mrope, - ) - rotary_cos_sin = torch.tensor(rotary_cos_sin).to(batch_input_ids[0].device) - rotary_cos_sin = rotary_cos_sin.reshape(max_position_embeddings, - int(rotary_embedding_dim / 2), 2) - - cos_ori = rotary_cos_sin[:, :, 0] - sin_ori = rotary_cos_sin[:, :, 1] - - mrope_position_ids_padding = torch.zeros( - (batch_size, max_position_embeddings), dtype=torch.int32) - for i in range(batch_size): - seq_len = batch_input_ids[i].shape[-1] - mrope_position_ids_padding[i, :seq_len] = torch.arange( - seq_len, device=batch_input_ids[i].device) - - cos = cos_ori[mrope_position_ids_padding].unsqueeze(-1) - sin = sin_ori[mrope_position_ids_padding].unsqueeze(-1) - - mrope_params.mrope_rotary_cos_sin = torch.concatenate( - (cos, sin), axis=-1).reshape(batch_size, -1) - mrope_params.mrope_position_deltas = torch.zeros( - [batch_size, 1], device=batch_input_ids[0].device) - - return mrope_params - - -def main(args): - is_integration_test = os.getenv('INTEGRATION_TEST', '0') == '1' - if is_integration_test: - logger.info( - "Running in integration test mode - will only run one batch and skip accuracy checks" - ) - logger.info( - "Setting max_ite=1 and check_accuracy=False for integration test") - args.max_ite = 1 - args.check_accuracy = False - - runtime_rank = tensorrt_llm.mpi_rank() - logger.set_level(args.log_level) - - test_hf = args.test_hf and runtime_rank == 0 # only run hf on rank 0 - test_trt_llm = args.test_trt_llm - model_name, model_version = read_model_name( - args.engine_dir if not test_hf else args.hf_model_dir, test_hf) - if args.hf_model_dir is None: - logger.warning( - "hf_model_dir is not specified. Try to infer from model_name, but this may be incorrect." - ) - if model_name in DEFAULT_HF_MODEL_DIRS: - args.hf_model_dir = DEFAULT_HF_MODEL_DIRS[model_name] - else: - args.hf_model_dir = None - if args.tokenizer_dir is None: - args.tokenizer_dir = args.hf_model_dir - - profiler.start('load tokenizer') - tokenizer, pad_id, end_id = load_tokenizer( - tokenizer_dir=args.tokenizer_dir, - vocab_file=args.vocab_file, - model_name=model_name, - model_version=model_version, - tokenizer_type=args.tokenizer_type, - ) - profiler.stop('load tokenizer') - logger.info( - f'Load tokenizer takes: {profiler.elapsed_time_in_sec("load tokenizer")} sec' - ) - - if args.eval_task == 'code_completion': - dataset_name = "openai_humaneval" - dataset_revision = None - dataset_input_key = 'prompt' - dataset_output_key = 'canonical_solution' - dataset_split = 'test' - elif args.eval_task == 'summarize': - dataset_name = "ccdv/cnn_dailymail" - dataset_revision = "3.0.0" - dataset_input_key = 'article' - dataset_output_key = 'highlights' - dataset_split = 'test' - elif args.eval_task == 'summarize_long': - dataset_name = "tau/zero_scrolls" - dataset_revision = 'squality' - dataset_input_key = 'input' - dataset_output_key = 'output' - dataset_split = 'validation' # only this split contains reference strings - elif args.eval_task == "eval_context_ppl": - dataset_name = "SlimPajama-6B" - dataset_revision = None - dataset_input_key = 'text' - dataset_output_key = 'text' - dataset_split = 'test' - args.output_len = 1 # Only want to compute the ppl of context - args.eval_ppl = True - logger.warning( - f"Run task '{args.eval_task}', setting 'output_len' to 1, and enable 'eval_ppl'." - ) - if args.dataset_dir is not None and isinstance(args.dataset_dir, str): - args.dataset_dir = args.dataset_dir.rstrip('/') - if args.dataset_dir.endswith(dataset_name): - dataset_name = args.dataset_dir - else: - dataset_name = f"{args.dataset_dir}/{dataset_name}" - dataset = load_dataset(dataset_name, - dataset_revision, - cache_dir=args.dataset_cache_dir, - split=dataset_split, - trust_remote_code=True) - dataset = dataset.shuffle(args.random_seed) - - max_batch_size = args.batch_size - - # runtime parameters - top_k = args.top_k - top_p = args.top_p - output_len = args.output_len - test_token_num = args.max_input_length - max_attention_window_size = args.max_attention_window_size - sink_token_length = args.sink_token_length - - if args.end_id: - end_id = args.end_id - - stop_words_list = None - if args.stop_words: - stop_words_list = tensorrt_llm.runtime.decode_words_list( - args.stop_words, tokenizer) - if model_version == 'glm4': # add default stop token ids for GLM-4 - glm4_stop_ids = [[151329], [151336], [151338]] - if stop_words_list is None: - stop_words_list = [glm4_stop_ids] * args.batch_size - else: - for req_stop_words_list in stop_words_list: - req_stop_words_list.extend(glm4_stop_ids) - - bad_words_list = None - if args.bad_words: - bad_words_list = tensorrt_llm.runtime.decode_words_list( - args.bad_words, tokenizer) - - if args.beam_width_array is not None: - logger.info("Use Variable-Beam-Width-Search") - args.beam_width_array, args.num_beams = get_beam_width_array( - args.beam_width_array) - - num_beams = args.num_beams - num_return_sequences = args.num_return_sequences - num_sequences = args.num_return_sequences or num_beams - assert num_beams == 1 or num_sequences <= num_beams - - temperature = args.temperature - length_penalty = args.length_penalty - early_stopping = args.early_stopping - beam_width_array = args.beam_width_array - repetition_penalty = args.repetition_penalty - presence_penalty = args.presence_penalty - frequency_penalty = args.frequency_penalty - prompt_ignore_length = args.prompt_ignore_length - random_seed = args.random_seed - torch.manual_seed(random_seed) - - output_dir = Path(args.output_dir) if args.output_dir else None - if output_dir is not None: - output_dir.mkdir(exist_ok=True, parents=True) - if test_trt_llm: - with (output_dir / 'trtllm.out').open('w') as f: - f.write(f'Engine path: {args.engine_dir}\n') - f.write(f'Tokenizer path: {args.tokenizer_dir}\n') - if test_hf: - with (output_dir / 'hf.out').open('w') as f: - f.write(f'Model path: {args.hf_model_dir}\n') - f.write(f'Tokenizer path: {args.tokenizer_dir}\n') - - rouge_dir = args.rouge_dir if args.rouge_dir and os.path.exists( - args.rouge_dir) else "rouge" - metric_tensorrt_llm = [ - evaluate.load(rouge_dir) for _ in range(num_sequences) - ] - metric_hf = [evaluate.load(rouge_dir) for _ in range(num_sequences)] - for i in range(num_sequences): - metric_tensorrt_llm[i].seed = 0 - metric_hf[i].seed = 0 - ppls_trt_llm = [[] for _ in range(num_sequences)] - ppls_hf = [[] for _ in range(num_sequences)] - - def _prepare_inputs(batch_input_texts, - eval_task='summarize', - add_special_tokens=True, - min_input_length=0): - batch_size = len(batch_input_texts) - append_str = ' TL;DR: ' if eval_task == 'summarize' else '' - batch_input_ids = [] - for i in range(batch_size): - curr_text = batch_input_texts[i] + append_str - curr_text = curr_text.strip().replace(" n't", "n't") - - # TODO: The below lines are used to be compatible with the original code; may need fix - if 'GLM' in model_name and model_version in ('chatglm2', - 'chatglm3'): - input_ids = tokenizer.encode(curr_text, - return_tensors='pt').squeeze(0) - input_ids = input_ids[:test_token_num] - elif 'qwen' in model_name.lower() and model_version == 'qwen': - # use make_content to generate prompt - system_prompt = "You are a useful assistant, please directly output the corresponding summary according to the article entered by the user." - _, input_id_list = make_context( - tokenizer=tokenizer, - query=curr_text, - history=[], - system=system_prompt, - max_input_length=test_token_num, - ) - input_ids = torch.tensor(input_id_list) - else: - if 'qwen' in model_name.lower() and 'qwen2' in model_version: - messages = [{ - "role": - "system", - "content": - "You are a helpful assistant, please summarize the article entered by the user with one or two sentences." - }, { - "role": "user", - "content": curr_text - }] - curr_text = tokenizer.apply_chat_template( - messages, tokenize=False, add_generation_prompt=True) - input_ids = tokenizer.encode( - curr_text, - return_tensors='pt', - add_special_tokens=add_special_tokens, - truncation=True, - max_length=test_token_num).squeeze(0) - - if input_ids.numel() > min_input_length: - batch_input_ids.append(input_ids) - return batch_input_ids - - def eval_trt_llm(datapoint, - eval_task='summarize', - eval_ppl=False, - add_special_tokens=True, - min_input_length=0, - runner=None): - batch_size = len(datapoint[dataset_input_key]) - batch_input_ids = _prepare_inputs(datapoint[dataset_input_key], - eval_task=eval_task, - add_special_tokens=add_special_tokens, - min_input_length=min_input_length) - # Generate mrope params for qwen model - engine_config = EngineConfig.from_json_file( - f"{args.engine_dir}/config.json") - pretrain_config = engine_config.pretrained_config - mrope_params = None - if 'qwen' in model_name.lower(): - mrope_params = ensemble_mrope_params( - batch_input_ids, - max_position_embeddings=pretrain_config.max_position_embeddings, - rotary_embedding_dim=pretrain_config.rotary_embedding_dim, - theta=pretrain_config.rotary_base, - ) - - if batch_size == 0 or len(batch_input_ids) == 0: - return [], [], [], {} - input_lengths = [x.size(0) for x in batch_input_ids] - - if args.ngram_config is not None: - # Speculative decoding of NGram - outputs = run_dtm_ngram(batch_input_ids, - args, - runtime_rank, - end_id, - pad_id, - stop_words_list, - bad_words_list, - tokenizer.vocab_size, - target_runner=runner) - if not args.streaming: # Unpack runner from the return value in No-Streaming mode - outputs, runner = list(outputs)[0] - else: # Normal run - with torch.no_grad(): - outputs = runner.generate( - batch_input_ids, - max_new_tokens=output_len, - max_attention_window_size=max_attention_window_size, - sink_token_length=sink_token_length, - end_id=end_id, - pad_id=pad_id, - temperature=temperature, - top_k=top_k, - top_p=top_p, - stop_words_list=stop_words_list, - bad_words_list=bad_words_list, - num_beams=num_beams, - num_return_sequences=num_return_sequences, - length_penalty=length_penalty, - early_stopping=early_stopping, - beam_width_array=beam_width_array, - repetition_penalty=repetition_penalty, - presence_penalty=presence_penalty, - frequency_penalty=frequency_penalty, - prompt_ignore_length=prompt_ignore_length, - lora_uids=args.lora_task_uids, - lookahead_config=args.lookahead_config, - output_sequence_lengths=True, - output_generation_logits=eval_ppl, - return_dict=True, - random_seed=random_seed, - medusa_choices=args.medusa_choices, - eagle_choices=args.eagle_choices, - mrope_params=mrope_params) - torch.cuda.synchronize() - - # Extract a list of tensors of shape beam_width x output_ids. - if runtime_rank == 0: - output_ids = outputs['output_ids'] - output_beams_list = [ - tokenizer.batch_decode(beam_tokens[:, input_lengths[i]:], - skip_special_tokens=True) - for i, beam_tokens in enumerate(output_ids) - ] - output_ids_list = [ - beam_tokens[:, input_lengths[i]:] - for i, beam_tokens in enumerate(output_ids) - ] - - ppls = [[] for _ in range(batch_size)] - lengths_info = { - 'input_lengths': input_lengths, - 'seq_lengths': outputs["sequence_lengths"].cpu().tolist(), - } - if eval_ppl: - seq_lengths = outputs['sequence_lengths'] - context_logits = outputs['context_logits'] - # Remove the first generation logits which are same to last - # context logits. - generation_logits = outputs['generation_logits'][:, :, 1:] - for batch_idx in range(batch_size): - # [batch, beam, step] - for beam_idx in range(num_sequences): - curr_len = seq_lengths[batch_idx, beam_idx] - curr_ctx_len = input_lengths[batch_idx] - curr_gen_len = curr_len - curr_ctx_len - - curr_ids = output_ids[batch_idx, beam_idx, 1:curr_len] - curr_logits = torch.cat([ - context_logits[batch_idx], - generation_logits[batch_idx, - beam_idx, :curr_gen_len - 1] - ], - dim=0) - curr_ppl = ppl(curr_logits, curr_ids) - logger.debug(f"TensorRT LLM PPL: {curr_ppl:.3f} | " - f"Generation length: {curr_gen_len}") - ppls[batch_idx].append(curr_ppl) - return output_beams_list, output_ids_list, ppls, lengths_info - return [], [], [], {} - - def eval_hf(datapoint, - eval_task='summarize', - eval_ppl=False, - add_special_tokens=True, - min_input_length=0): - batch_size = len(datapoint[dataset_input_key]) - if batch_size > 1: - logger.warning( - f"HF does not support batch_size > 1 to verify correctness due to padding. Current batch size is {batch_size}" - ) - batch_input_ids = _prepare_inputs(datapoint[dataset_input_key], - eval_task=eval_task, - add_special_tokens=add_special_tokens, - min_input_length=min_input_length) - batch_size = len(batch_input_ids) - if batch_size == 0: - return [], [], [], [[] for _ in range(batch_size)] - input_lengths = [x.size(0) for x in batch_input_ids] - # Left padding for HF - max_length = max(input_lengths) - paddings = [ - torch.ones(max_length - l, dtype=torch.int32) * pad_id - for l in input_lengths - ] - batch_input_ids = [ - torch.cat([pad, x]) for x, pad in zip(batch_input_ids, paddings) - ] - batch_input_ids = torch.stack(batch_input_ids) - batch_input_ids = batch_input_ids.cuda() - - # specialization for HF - if early_stopping in [0, 1]: - local_early_stopping = bool(early_stopping) - else: - local_early_stopping = "never" - - with torch.no_grad(): - hf_config = {} - if num_beams == 1: - hf_config.update({ - "top_k": top_k, - "top_p": top_p, - "do_sample": True, - }) - else: - hf_config.update({ - "num_beams": num_beams, - "early_stopping": local_early_stopping, - }) - - outputs = model.generate(batch_input_ids, - max_new_tokens=output_len, - num_return_sequences=num_sequences, - temperature=temperature, - eos_token_id=end_id, - pad_token_id=pad_id, - length_penalty=length_penalty, - output_scores=True, - return_dict_in_generate=True, - **hf_config) - if eval_ppl and batch_size == 1: - # model.generate cannot return context logits? - # Will cause additional latency - context_outputs = model(batch_input_ids) - - output_ids = outputs['sequences'] - tokens_list = output_ids[:, max_length:].tolist() - output_ids = output_ids.reshape([batch_size, num_sequences, -1]) - output_lines_list = [ - tokenizer.batch_decode(output_ids[:, i, max_length:], - skip_special_tokens=True) - for i in range(num_sequences) - ] - - ppls = [[] for _ in range(batch_size)] - if eval_ppl and batch_size == 1: - # Only for batch size of 1 - seq_lens = (output_ids - != end_id).logical_and(output_ids != pad_id).sum(dim=-1) - context_logits = context_outputs['logits'] - # Remove the first generation logits which are same to last context logits - generation_logits = outputs['scores'][1:] - # When output_len is 1, generation_logits would be () and lead to error if we do torch.stack - if len(generation_logits) == 0: - generation_logits = torch.empty( - [context_logits.shape[0], 0, context_logits.shape[-1]], - device=context_logits.device) - else: - generation_logits = torch.stack(generation_logits, dim=1) - _, max_gen_len, voc_size = generation_logits.size() - generation_logits = generation_logits.view(batch_size, num_beams, - max_gen_len, voc_size) - for batch_idx in range(batch_size): - for beam_idx in range(num_sequences): - curr_len = seq_lens[batch_idx, beam_idx] - curr_ctx_len = input_lengths[batch_idx] - curr_gen_len = curr_len - curr_ctx_len - - curr_ids = output_ids[batch_idx, beam_idx, 1:curr_len] - curr_logits = torch.cat([ - context_logits[batch_idx], - generation_logits[batch_idx, - beam_idx, :curr_gen_len - 1] - ], - dim=0) - curr_ppl = ppl(curr_logits, curr_ids) - logger.debug( - f"HF PPL: {curr_ppl:.3f} | Generation length: {curr_gen_len}" - ) - ppls[batch_idx].append(curr_ppl) - - return output_lines_list, tokens_list, ppls - - if test_trt_llm: - if not supports_inflight_batching(args.engine_dir): - logger.warning( - "The given engine does not support in-flight batching, fallback to python session" - ) - args.use_py_session = True - - if not PYTHON_BINDINGS and not args.use_py_session: - logger.warning( - "Python bindings of C++ session is unavailable, fallback to Python session." - ) - args.use_py_session = True - if args.return_all_generated_tokens: - raise ValueError( - "Returning all the generated tokens at each step is not supported in summarize.py" - ) - - logger.info( - f"Using {'Python' if args.use_py_session else 'C++'} session") - - runner_cls = ModelRunner if args.use_py_session else ModelRunnerCpp - runner_kwargs = dict( - engine_dir=args.engine_dir, - rank=runtime_rank, - debug_mode=args.debug_mode, - gpu_weights_percent=args.gpu_weights_percent, - enable_context_fmha_fp32_acc=args.enable_context_fmha_fp32_acc, - ) - if not args.use_py_session: - runner_kwargs.update( - lora_dir=args.lora_dir, - lora_ckpt_source=args.lora_ckpt_source, - max_batch_size=max_batch_size, - max_input_len=test_token_num, - max_output_len=output_len, - max_beam_width=num_beams, - max_attention_window_size=max_attention_window_size, - sink_token_length=sink_token_length, - max_tokens_in_paged_kv_cache=args.max_tokens_in_paged_kv_cache, - kv_cache_enable_block_reuse=args.kv_cache_enable_block_reuse, - kv_cache_free_gpu_memory_fraction=args. - kv_cache_free_gpu_memory_fraction, - enable_chunked_context=args.enable_chunked_context, - multi_block_mode=args.multi_block_mode, - cuda_graph_mode=args.cuda_graph_mode, - gather_generation_logits=args.eval_ppl, - use_gpu_direct_storage=args.use_gpu_direct_storage, - ) - - if args.medusa_choices is not None: - args.medusa_choices = ast.literal_eval(args.medusa_choices) - assert args.temperature == 1.0, "Medusa should use temperature == 1.0" - assert args.num_beams == 1, "Medusa should use num_beams == 1" - runner_kwargs.update(medusa_choices=args.medusa_choices) - if args.eagle_choices is not None or args.eagle_posterior_threshold is not None or args.eagle_use_dynamic_tree: - assert args.num_beams == 1, "Eagle should use num_beams == 1" - if args.eagle_choices is not None and not args.eagle_use_dynamic_tree: - args.eagle_choices = ast.literal_eval(args.eagle_choices) - runner_kwargs.update(eagle_choices=args.eagle_choices) - if args.eagle_posterior_threshold is not None: - runner_kwargs.update( - eagle_posterior_threshold=args.eagle_posterior_threshold) - if args.eagle_use_dynamic_tree: - runner_kwargs.update( - eagle_use_dynamic_tree=args.eagle_use_dynamic_tree) - assert args.eagle_dynamic_tree_max_top_k is not None and args.eagle_dynamic_tree_max_top_k > 0 - runner_kwargs.update(eagle_dynamic_tree_max_top_k=args. - eagle_dynamic_tree_max_top_k) - if args.lookahead_config is not None: - args.lookahead_config = ast.literal_eval(args.lookahead_config) - assert len( - args.lookahead_config - ) == 3, "Lookahead needs [max_window_size, max_ngram_size, max_verification_set_size]" - runner_kwargs.update(lookahead_config=args.lookahead_config) - if args.ngram_config is not None: - assert args.kv_cache_enable_block_reuse, "`--kv_cache_enable_block_reuse` must be specified in speculative decoding." - assert not args.use_py_session, "`--use_py_session` is not supported in Speculative decoding." - assert args.num_beams == 1, "`--num_beams>1` is not supported in Speculative decoding." - max_draft_len, _, target_device_list = ast.literal_eval( - args.ngram_config) - args.max_output_len = output_len # Specialization for NGram - runner_kwargs.update(is_orchestrator_mode=True, - device_ids=target_device_list, - max_input_len=test_token_num + max_draft_len + - output_len) - - runner = runner_cls.from_dir(**runner_kwargs) - assert not (args.eval_ppl and not runner.gather_context_logits), \ - "PPL evaluation requires engine built with gather_context_logits enabled" - - datapoint = dataset[0:1] - output, *_ = eval_trt_llm(datapoint, - eval_task=args.eval_task, - eval_ppl=args.eval_ppl, - add_special_tokens=args.add_special_tokens, - min_input_length=args.min_input_length, - runner=runner) - if runtime_rank == 0 and args.eval_task != "eval_context_ppl": - logger.info( - "---------------------------------------------------------") - logger.info("TensorRT LLM Generated: ") - logger.info(f" Input: {datapoint[dataset_input_key]}") - logger.info(f"\n Reference: {datapoint[dataset_output_key]}") - logger.info(f"\n Output: {output}") - logger.info( - "---------------------------------------------------------") - - ite_count = 0 - data_point_idx = 0 - total_output_token_count_trt_llm = 0 # only valid for runtime_rank == 0 - while (data_point_idx < len(dataset)) and (ite_count < args.max_ite): - if runtime_rank == 0: - logger.debug( - f"run data_point {data_point_idx} ~ {data_point_idx + max_batch_size}" - ) - datapoint = dataset[data_point_idx:(data_point_idx + - max_batch_size)] - - profiler.start('tensorrt_llm') - output_tensorrt_llm, output_ids_trt_llm, curr_ppls_trt_llm, lengths_info = eval_trt_llm( - datapoint, - eval_task=args.eval_task, - eval_ppl=args.eval_ppl, - add_special_tokens=args.add_special_tokens, - min_input_length=args.min_input_length, - runner=runner) - profiler.stop('tensorrt_llm') - - empty_batch = runtime_rank == 0 and len(output_tensorrt_llm) == 0 - empty_batch = mpi_broadcast(empty_batch, 0) - if empty_batch: - # No valid samples in the current batch, skip this iteration - data_point_idx += max_batch_size - continue - - if runtime_rank == 0: - input_lengths = lengths_info['input_lengths'] - seq_lengths = lengths_info['seq_lengths'] - output_token_count_trt_llm = sum( - beam_len - input_lengths[batch_idx] - for batch_idx, beam_lens in enumerate(seq_lengths) - for beam_len in beam_lens) - total_output_token_count_trt_llm += output_token_count_trt_llm - for batch_idx, output_beams in enumerate(output_tensorrt_llm): - reference = datapoint[dataset_output_key][batch_idx] - for beam_idx, output_beam in enumerate(output_beams): - metric_tensorrt_llm[beam_idx].add_batch( - predictions=[output_beam], references=[reference]) - if args.eval_ppl: - ppls_trt_llm[beam_idx].append( - curr_ppls_trt_llm[batch_idx][beam_idx]) - if output_dir is not None: - for i in range(len(output_tensorrt_llm[0])): - for beam_idx in range(num_sequences): - with (output_dir / 'trtllm.out').open('a') as f: - f.write( - f'[{data_point_idx + i}] [Beam {beam_idx}] {output_tensorrt_llm[beam_idx][i]}\n' - ) - - logger.debug('-' * 100) - logger.debug(f"Input: {datapoint[dataset_input_key]}") - logger.debug(f'TensorRT LLM Output: {output_tensorrt_llm}') - logger.debug(f"Reference: {datapoint[dataset_output_key]}") - - data_point_idx += max_batch_size - ite_count += 1 - del runner - - if test_hf and runtime_rank == 0: - profiler.start('load HF model') - dtype_alias_mapping = { - 'fp32': 'float32', - 'fp16': 'float16', - 'bf16': 'bfloat16' - } - args.hf_data_type = dtype_alias_mapping.get(args.hf_data_type, - args.hf_data_type) - if 'GLM' in model_name and model_version == 'glm': - auto_model_cls = AutoModelForSeq2SeqLM - elif 'GLM' in model_name and model_version == 'chatglm': - auto_model_cls = AutoModel - else: - auto_model_cls = AutoModelForCausalLM - # TODO: args.hf_device_map_auto is not being correctly set - # remove in future version - if model_name == 'DeepseekV2ForCausalLM': - args.hf_device_map_auto = True - model = auto_model_cls.from_pretrained( - args.hf_model_dir, - trust_remote_code=True, - dtype=str_dtype_to_torch(args.hf_data_type), - device_map='auto' if args.hf_device_map_auto else None) - try: - model.to_bettertransformer() - except Exception as e: - logger.warning( - f'Fail to call model.to_bettertransformer(), exception:\n{str(e)}' - ) - if not args.hf_device_map_auto: - model.cuda() - if model_name == 'qwen': - model.generation_config = GenerationConfig.from_pretrained( - args.hf_model_dir, trust_remote_code=True) - profiler.stop('load HF model') - logger.info( - f'Load HF model takes: {profiler.elapsed_time_in_sec("load HF model")} sec' - ) - - datapoint = dataset[0:1] - output, *_ = eval_hf(datapoint, - eval_task=args.eval_task, - eval_ppl=args.eval_ppl, - add_special_tokens=args.add_special_tokens, - min_input_length=args.min_input_length) - if runtime_rank == 0 and args.eval_task != "eval_context_ppl": - logger.info( - "---------------------------------------------------------") - logger.info("HF Generated: ") - logger.info(f" Input: {datapoint[dataset_input_key]}") - logger.info(f"\n Reference: {datapoint[dataset_output_key]}") - logger.info(f"\n Output: {output}") - logger.info( - "---------------------------------------------------------") - - ite_count = 0 - data_point_idx = 0 - total_output_token_count_hf = 0 # only valid for runtime_rank == 0 - while (data_point_idx < len(dataset)) and (ite_count < args.max_ite): - if runtime_rank == 0: - logger.debug( - f"run data_point {data_point_idx} ~ {data_point_idx + max_batch_size}" - ) - datapoint = dataset[data_point_idx:(data_point_idx + - max_batch_size)] - - profiler.start('hf') - output_hf, token_list, curr_ppls_hf = eval_hf( - datapoint, - eval_task=args.eval_task, - eval_ppl=args.eval_ppl, - add_special_tokens=args.add_special_tokens, - min_input_length=args.min_input_length) - profiler.stop('hf') - - # HF model runs on rank 0 only - empty_batch = len(output_hf) == 0 - if empty_batch: - # No valid samples in the current batch, skip this iteration - data_point_idx += max_batch_size - continue - - if runtime_rank == 0: - seq_lengths = [len(tokens) for tokens in token_list] - total_output_token_count_hf += sum(seq_lengths) - for beam_idx in range(num_sequences): - for batch_idx in range(len(output_hf[beam_idx])): - metric_hf[beam_idx].add_batch( - predictions=[output_hf[beam_idx][batch_idx]], - references=[ - datapoint[dataset_output_key][batch_idx] - ]) - if args.eval_ppl and args.batch_size == 1: - ppls_hf[beam_idx].append( - curr_ppls_hf[batch_idx][beam_idx]) - if output_dir is not None: - for i in range(len(output_hf[0])): - for beam_idx in range(num_sequences): - with (output_dir / 'hf.out').open('a') as f: - f.write( - f'[{data_point_idx + i}] [Beam {beam_idx}] {output_hf[beam_idx][i]}\n' - ) - - logger.debug('-' * 100) - logger.debug(f"Input: {datapoint[dataset_input_key]}") - logger.debug(f'HF Output: {output_hf}') - logger.debug(f"Reference: {datapoint[dataset_output_key]}") - - data_point_idx += max_batch_size - ite_count += 1 - del model - - if runtime_rank == 0 and args.max_ite > 0: - if test_trt_llm: - np.random.seed(0) # rouge score use sampling to compute the score - logger.info( - f'TensorRT LLM (total latency: {profiler.elapsed_time_in_sec("tensorrt_llm")} sec)' - ) - - logger.info( - f'TensorRT LLM (total output tokens: {total_output_token_count_trt_llm})' - ) - logger.info( - f'TensorRT LLM (tokens per second: {total_output_token_count_trt_llm / profiler.elapsed_time_in_sec("tensorrt_llm")})' - ) - for beam_idx in range(num_sequences): - logger.info(f"TensorRT LLM beam {beam_idx} result") - if args.eval_task != "eval_context_ppl": - if args.estimate_accuracy_std_dev: - computed_metrics_tensorrt_llm = metric_tensorrt_llm[ - beam_idx].compute(use_aggregator=False) - computed_std_dev_tensorrt_llm = { - key: np.std(scores) - for key, scores in - computed_metrics_tensorrt_llm.items() - } - computed_metrics_tensorrt_llm = { - key: np.mean(scores) - for key, scores in - computed_metrics_tensorrt_llm.items() - } - for key in computed_metrics_tensorrt_llm.keys(): - logger.info( - f" {key}: {computed_metrics_tensorrt_llm[key]*100} ({computed_std_dev_tensorrt_llm[key]*100})" - ) - else: - computed_metrics_tensorrt_llm = metric_tensorrt_llm[ - beam_idx].compute() - for key in computed_metrics_tensorrt_llm.keys(): - logger.info( - f" {key}: {computed_metrics_tensorrt_llm[key]*100}" - ) - if args.check_accuracy and beam_idx == 0: - rouge1 = computed_metrics_tensorrt_llm['rouge1'] * 100 - assert rouge1 > args.tensorrt_llm_rouge1_threshold, f"[FAILED] rouge1 ({rouge1}) is smaller than threshold ({args.tensorrt_llm_rouge1_threshold})." - if args.eval_ppl: - logger.info( - f" Per-token perplexity: {np.mean(ppls_trt_llm[beam_idx])}" - ) - if args.check_accuracy and beam_idx == 0: - avg_ppl = np.mean(ppls_trt_llm[beam_idx]) - assert avg_ppl < args.tensorrt_llm_ppl_threshold, f"[FAILED] average PPL ({avg_ppl}) is larger than threshold ({args.tensorrt_llm_ppl_threshold})." - if test_hf: - np.random.seed(0) # rouge score use sampling to compute the score - logger.info( - f'Hugging Face (total latency: {profiler.elapsed_time_in_sec("hf")} sec)' - ) - logger.info( - f'Hugging Face (total output tokens: {total_output_token_count_hf})' - ) - logger.info( - f'Hugging Face (tokens per second: {total_output_token_count_hf / profiler.elapsed_time_in_sec("hf")})' - ) - - for beam_idx in range(num_sequences): - logger.info(f"HF beam {beam_idx} result") - computed_metrics_hf = metric_hf[beam_idx].compute() - if args.eval_task != "eval_context_ppl": - for key in computed_metrics_hf.keys(): - logger.info(f' {key}: {computed_metrics_hf[key]*100}') - if args.eval_ppl and args.batch_size == 1: - logger.info( - f" Per-token perplexity: {np.mean(ppls_hf[beam_idx])}") - - -if __name__ == '__main__': - # see `add_common_args` for extended list of arguments - parser = argparse.ArgumentParser() - parser.add_argument('--test_hf', action='store_true') - parser.add_argument('--test_trt_llm', action='store_true') - parser.add_argument('--eval_task', - type=str, - default='summarize', - choices=[ - 'summarize', 'summarize_long', 'code_completion', - 'eval_context_ppl' - ]) - parser.add_argument('--check_accuracy', action='store_true') - parser.add_argument('--estimate_accuracy_std_dev', action='store_true') - parser.add_argument('--tensorrt_llm_rouge1_threshold', - type=float, - default=15.0) - parser.add_argument('--eval_ppl', action='store_true') - parser.add_argument('--tensorrt_llm_ppl_threshold', - type=float, - default=15.0) - parser.add_argument( - '--dataset_dir', - type=str, - default=None, - help="The local directory of the dataset for evaluation; " - "will download the dataset from huggingface hub if not specified.") - parser.add_argument( - '--dataset_cache_dir', - type=str, - default=None, - help="The local cache directory for dataset; " - "will use `~/.cache/huggingface/datasets` if not specified.") - parser.add_argument('--batch_size', type=int, default=1) - parser.add_argument('--max_ite', type=int, default=20) - parser.add_argument('--output_len', type=int, default=100) - parser.add_argument('--max_input_length', type=int, default=923) - parser.add_argument( - '--min_input_length', - type=int, - default=0, - help='skip the sentences which are shorter than min_input_length.') - parser.add_argument( - '--output_dir', - type=str, - default=None, - help="Directory where to save output sentences. 'trtllm.out' for " - "TensorRT LLM outputs, and 'hf.out' for HF outputs. If None, do not " - "save outputs.") - parser.add_argument( - '--rouge_dir', - default=None, - type=str, - help= - "evaluate.load('rouge') will attempt to pull rouge package from HF. Use cached rouge can avoid network outage of host or HF." - ) - parser.add_argument("--use_gpu_direct_storage", - default=False, - action="store_true", - help="Use GPUDirect Storage (GDS) to load the engine") - parser = add_common_args(parser) - args = parser.parse_args() - - main(args) diff --git a/examples/utils.py b/examples/utils.py deleted file mode 100644 index c75cd255850f..000000000000 --- a/examples/utils.py +++ /dev/null @@ -1,658 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -import ast -import json -import os -import subprocess -import sys -from argparse import BooleanOptionalAction -from functools import partial -from pathlib import Path -from typing import List, Optional - -import torch -from transformers import AutoTokenizer, LlamaTokenizer - -from tensorrt_llm._utils import supports_inflight_batching # noqa -from tensorrt_llm._utils import (mpi_barrier, mpi_rank, mpi_world_size, - str_dtype_to_torch) -from tensorrt_llm.builder import get_engine_version - - -class SentencePieceTokenizer: - """Minimal SentencePiece-backed tokenizer with a transformers-like API. - - transformers v5 replaced the pure-Python SentencePiece backend of - ``T5Tokenizer`` / ``LlamaTokenizer`` with the Rust ``tokenizers`` backend, - so instantiating them with a raw SentencePiece ``.model`` vocab file no - longer reads the underlying vocabulary. This wrapper preserves the old - behavior for NEMO-style checkpoints (e.g. gpt-next) by delegating to - ``sentencepiece.SentencePieceProcessor`` directly. - """ - - def __init__(self, - vocab_file: str, - padding_side: str = 'left', - truncation_side: str = 'left'): - import sentencepiece as spm - sp = spm.SentencePieceProcessor() - sp.Load(vocab_file) - self.sp_model = sp - self.padding_side = padding_side - self.truncation_side = truncation_side - self.vocab_size = sp.GetPieceSize() - - def _opt(i: int) -> int | None: - return i if i >= 0 else None - - self.pad_token_id = _opt(sp.pad_id()) - self.eos_token_id = _opt(sp.eos_id()) - self.bos_token_id = _opt(sp.bos_id()) - self.unk_token_id = _opt(sp.unk_id()) - - def encode(self, - text: str, - return_tensors: Optional[str] = None, - add_special_tokens: bool = True, - truncation: bool = False, - max_length: Optional[int] = None, - **kwargs): - ids = self.sp_model.EncodeAsIds(text) - if add_special_tokens and self.bos_token_id is not None: - ids = [self.bos_token_id] + ids - if truncation and max_length is not None and len(ids) > max_length: - ids = ids[ - -max_length:] if self.truncation_side == 'left' else ids[: - max_length] - if return_tensors == 'pt': - return torch.tensor([ids], dtype=torch.long) - return ids - - def decode(self, ids, skip_special_tokens: bool = False, **kwargs) -> str: - if isinstance(ids, torch.Tensor): - ids = ids.tolist() - if skip_special_tokens: - special = {self.pad_token_id, self.eos_token_id, self.bos_token_id} - ids = [t for t in ids if t not in special] - return self.sp_model.DecodeIds(list(ids)) - - def batch_decode(self, - sequences, - skip_special_tokens: bool = False, - **kwargs) -> List[str]: - if isinstance(sequences, torch.Tensor): - sequences = sequences.tolist() - return [ - self.decode(seq, skip_special_tokens=skip_special_tokens) - for seq in sequences - ] - - -DEFAULT_HF_MODEL_DIRS = { - 'BaichuanForCausalLM': 'baichuan-inc/Baichuan-13B-Chat', - 'BaiChuanForCausalLM': 'baichuan-inc/Baichuan-13B-Chat', - 'BloomForCausalLM': 'bigscience/bloom-560m', - 'GLMModel': 'THUDM/glm-10b', - 'ChatGLMModel': 'THUDM/chatglm3-6b', - 'ChatGLMForCausalLM': 'THUDM/chatglm3-6b', - 'RWForCausalLM': 'tiiuae/falcon-rw-1b', - 'FalconForCausalLM': 'tiiuae/falcon-rw-1b', - 'GPT2LMHeadModel': 'gpt2', - 'GPT2LMHeadCustomModel': 'gpt2', - 'Starcoder2ForCausalLM': 'bigcode/starcoder2-3b', - 'GPTForCausalLM': 'gpt2', - 'GPTJForCausalLM': 'EleutherAI/gpt-j-6b', - 'GPTNeoXForCausalLM': 'EleutherAI/gpt-neox-20b', - 'InternLMForCausalLM': 'internlm/internlm-chat-7b', - 'InternLM2ForCausalLM': 'internlm/internlm2-chat-7b', - 'LlamaForCausalLM': 'meta-llama/Llama-2-7b-hf', - 'MPTForCausalLM': 'mosaicml/mpt-7b', - 'PhiForCausalLM': 'microsoft/phi-2', - 'OPTForCausalLM': 'facebook/opt-350m', - 'QWenLMHeadModel': 'Qwen/Qwen-7B', - 'QWenForCausalLM': 'Qwen/Qwen-7B', - 'Qwen2ForCausalLM': 'Qwen/Qwen1.5-7B', - 'Qwen2MoeForCausalLM': 'Qwen/Qwen1.5-MoE-A2.7B', - 'RecurrentGemmaForCausalLM': 'google/recurrentgemma-2b', -} - -INTERNLM_META_INSTRUCTION = """You are an AI assistant whose name is InternLM (书生·浦语). -- InternLM (书生·浦语) is a conversational language model that is developed by Shanghai AI Laboratory (上海人工智能实验室). It is designed to be helpful, honest, and harmless. -- InternLM (书生·浦语) can understand and communicate fluently in the language chosen by the user such as English and 中文. -""" - -QWEN_PROMPT_TEMPLATE = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n{input_text}<|im_end|>\n<|im_start|>assistant\n" - -DEFAULT_PROMPT_TEMPLATES = { - 'InternLMForCausalLM': "<|User|>:{input_text}\n<|Bot|>:", - 'InternLM2ForCausalLM': "<|im_start|>system\n" + INTERNLM_META_INSTRUCTION + - "<|im_end|>\n<|im_start|>user\n{input_text}<|im_end|>\n<|im_start|>assistant\n", - 'QWenLMHeadModel': QWEN_PROMPT_TEMPLATE, - 'QWenForCausalLM': QWEN_PROMPT_TEMPLATE, - 'Qwen2ForCausalLM': QWEN_PROMPT_TEMPLATE, - 'Qwen2MoeForCausalLM': QWEN_PROMPT_TEMPLATE, -} - - -def read_decoder_start_token_id(engine_dir): - with open(Path(engine_dir) / "config.json", 'r') as f: - config = json.load(f) - return config['pretrained_config']['decoder_start_token_id'] - - -def read_is_enc_dec(engine_dir: str, is_hf: bool = False): - if is_hf: - with open(Path(engine_dir) / "config.json", 'r') as f: - config = json.load(f) - is_enc_dec = config.get('is_encoder_decoder', False) - else: - is_enc_dec = {'encoder', 'decoder'}.issubset({ - name - for name in os.listdir(engine_dir) - if os.path.isdir(os.path.join(engine_dir, name)) - }) - return is_enc_dec - - -def read_model_name(engine_dir: str, is_hf: bool = False): - with open(Path(engine_dir) / "config.json", 'r') as f: - config = json.load(f) - - if is_hf: - model_arch = config['architectures'][0] - model_version = config.get('model_type', None) - return model_arch, model_version - - engine_version = get_engine_version(engine_dir) - if engine_version is None: - return config['builder_config']['name'], None - - model_arch = config['pretrained_config']['architecture'] - model_version = None - if 'GLM' in model_arch: - model_version = config['pretrained_config']['chatglm_version'] - if 'qwen' in model_arch.lower(): - model_version = config['pretrained_config']['qwen_type'] - return model_arch, model_version - - -def throttle_generator(generator, stream_interval): - for i, out in enumerate(generator): - if not i % stream_interval: - yield out - - if i % stream_interval: - yield out - - -# Load tokenizer impl, it will be called in external wrapper to avoid loading tokenizer bug under MPI env. -def _load_tokenizer(tokenizer_dir: Optional[str] = None, - vocab_file: Optional[str] = None, - model_name: str = 'GPTForCausalLM', - model_version: Optional[str] = None, - tokenizer_type: Optional[str] = None): - if vocab_file is None: - if 'whisper' in model_name.lower(): - tokenizer = AutoTokenizer.from_pretrained( - tokenizer_dir or 'openai/whisper-large-v3', - language='english', - task='transcribe', - predict_timestamps=False, - ) - elif tokenizer_type == 'language_adapter': - tokenizer = None - else: - use_fast = True - if tokenizer_type is not None and tokenizer_type == "llama": - use_fast = False - # Should set both padding_side and truncation_side to be 'left' - tokenizer = AutoTokenizer.from_pretrained( - tokenizer_dir, - legacy=False, - padding_side='left', - truncation_side='left', - trust_remote_code=True, - tokenizer_type=tokenizer_type, - use_fast=use_fast) - elif model_name == 'GemmaForCausalLM' or model_name == 'RecurrentGemmaForCausalLM': - from transformers import GemmaTokenizer - - # Initialize tokenizer from vocab file. - tokenizer = GemmaTokenizer(vocab_file=vocab_file, - padding_side='left', - truncation_side='left', - legacy=False) - elif model_name == 'Grok1ModelForCausalLM': - tokenizer = LlamaTokenizer(vocab_file=vocab_file, - padding_side='left', - truncation_side='left', - legacy=False, - use_fast=False) - else: - # For gpt-next, directly load from the SentencePiece ``tokenizer.model`` - # file. transformers v5 removed the pure-Python SentencePiece backend, - # so ``T5Tokenizer(vocab_file=...)`` no longer reads the vocabulary and - # reports vocab_size=104 with all tokens decoding to . - tokenizer = SentencePieceTokenizer(vocab_file=vocab_file, - padding_side='left', - truncation_side='left') - if 'qwen' in model_name.lower() and model_version == 'qwen': - with open(Path(tokenizer_dir) / "generation_config.json") as f: - gen_config = json.load(f) - pad_id = gen_config['pad_token_id'] - end_id = gen_config['eos_token_id'] - elif 'GLM' in model_name and model_version == 'glm': - pad_id = tokenizer.pad_token_id - end_id = tokenizer.eop_token_id - elif tokenizer_type == 'language_adapter': - pad_id = 0 - end_id = 2 - else: - if tokenizer.pad_token_id is None: - tokenizer.pad_token_id = tokenizer.eos_token_id - pad_id = tokenizer.pad_token_id - end_id = tokenizer.eos_token_id - - return tokenizer, pad_id, end_id - - -def load_tokenizer(tokenizer_dir: Optional[str] = None, - vocab_file: Optional[str] = None, - model_name: str = 'GPTForCausalLM', - model_version: Optional[str] = None, - tokenizer_type: Optional[str] = None): - func = partial(_load_tokenizer, tokenizer_dir, vocab_file, model_name, - model_version, tokenizer_type) - if mpi_world_size() > 1: - # Under MPI env, load tokenizer will result in multiple processes to download the same file to the same folder. - # This will result some random bug. Force loading on rank0 to warmup the tokenizer to avoid this issue. - if mpi_rank() == 0: - func() - mpi_barrier() - return func() - - -def prepare_enc_dec_inputs(batch_input_ids: List[torch.Tensor], model_name: str, - engine_dir: str, - multimodal_input_file: Optional[str]): - encoder_input_features = None - encoder_input_ids = None - if 'whisper' in model_name.lower(): - # cannot directly import whisper due to name collision - sys.path.append(f"{os.path.dirname(__file__)}/models/core/whisper") - from whisper_utils import log_mel_spectrogram - - config_path = os.path.join(engine_dir, 'encoder', 'config.json') - with open(config_path, 'r') as f: - config = json.load(f) - n_mels = config['pretrained_config']['n_mels'] - dtype = config['pretrained_config']['dtype'] - - # download mel filters file - subprocess.run([ - "wget", "-nc", f"--directory-prefix={engine_dir}", - "https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz" - ], - check=True) - - mel, total_duration = log_mel_spectrogram(multimodal_input_file, - n_mels, - return_duration=True, - mel_filters_dir=engine_dir) - mel = mel.type(str_dtype_to_torch(dtype)) # [featureDim, seqLen] - decoder_input_ids = batch_input_ids - encoder_input_features = [torch.einsum('DL->LD', mel)] - encoder_output_lengths = [encoder_input_features[0].shape[0] // 2] - else: - encoder_input_ids = batch_input_ids - decoder_start_token_id = read_decoder_start_token_id( - os.path.join(engine_dir, "decoder")) - decoder_input_ids = [ - torch.tensor([decoder_start_token_id], dtype=torch.int32) - for _ in batch_input_ids - ] - encoder_output_lengths = None - return encoder_input_ids, encoder_input_features, encoder_output_lengths, decoder_input_ids - - -def get_beam_width_array(bwa: str = None): - bwa = ast.literal_eval(bwa) # Short for "beam_width_array" - if isinstance(bwa, str): - bwa = ast.literal_eval(bwa) # parse again for string - - def parse_one_bwa(row): - assert isinstance(row, list), f"Beam width array must be a list." - assert len( - row - ) <= 8, "Length of beam width array must not be greater than 8 now." - assert all([isinstance(beam, int) for beam in row - ]), "Numbers in beam width array must be integer." - bwa_tensor = torch.zeros([8], dtype=torch.int32) - for j in range(len(row)): - bwa_tensor[j] = row[j] - bwa_tensor[len(row):] = row[-1] - return bwa_tensor, max(row) - - if isinstance(bwa, list): # Only one BWA - bwa_tensor, max_beam_width = parse_one_bwa(bwa) - elif isinstance(bwa, tuple): # BWA for respective requests - bwa_tensor_list = [] - max_beam_width = 0 - for row in bwa: - bwa_tensor, beam_width = parse_one_bwa(row) - bwa_tensor_list.append(bwa_tensor) - max_beam_width = max(max_beam_width, beam_width) - bwa_tensor = torch.stack(bwa_tensor_list, dim=0) - else: - raise ValueError(f"Invalid beam width array: {bwa}") - - return bwa_tensor.tolist(), max_beam_width - - -def add_common_args(parser): - # sampling arguments - parser.add_argument('--num_beams', - type=int, - help="Use beam search if num_beams > 1", - default=1) - parser.add_argument('--num_return_sequences', - type=int, - help="Number of sequences to generate for each input.", - default=None) - parser.add_argument('--temperature', type=float, default=1.0) - parser.add_argument('--top_k', type=int, default=1) - parser.add_argument('--top_p', type=float, default=0.0) - parser.add_argument('--length_penalty', type=float, default=1.0) - parser.add_argument('--repetition_penalty', type=float, default=1.0) - parser.add_argument('--presence_penalty', type=float, default=0.0) - parser.add_argument('--frequency_penalty', type=float, default=0.0) - parser.add_argument('--prompt_ignore_length', type=int, default=0) - parser.add_argument('--min_p', type=float, default=0.0) - parser.add_argument('--beam_search_diversity_rate', type=float, default=0.0) - parser.add_argument('--random_seed', type=int, default=0) - parser.add_argument('--early_stopping', - type=int, - help='Use early stopping if num_beams > 1, ' - '1 for early-stopping, 0 for non-early-stopping' - 'other values for stopping by length', - default=1) - parser.add_argument( - '--beam_width_array', - type=str, - default=None, - help= - 'Beam width array for each step. E.g.: --beam_width_array="[2,4,6,8]"', - ) - parser.add_argument( - '--end_id', - default=None, - type=int, - help="Override tokenizer end_id to stop on given end_id token.") - parser.add_argument( - '--stop_words', - default=None, - type=str, - nargs="+", - action='append', - help= - 'Set stop words for a batch. Successive invocations of --stop_words set stop words for other batches.' - ' E.g.: --stop_words " London" " chef" --stop_words "eventually became" "was not"', - ) - parser.add_argument( - '--bad_words', - default=None, - type=str, - nargs="+", - action='append', - help= - 'Set bad words for a batch. Successive invocations of --bad_words set bad words for other batches.' - ' E.g.: --bad_words " London" " chef" --bad_words "eventually became" "was not"', - ) - parser.add_argument('--no_repeat_ngram_size', type=int, default=None) - - # common runtime arguments - parser.add_argument('--sink_token_length', - type=int, - default=None, - help='The sink token length.') - parser.add_argument( - '--max_attention_window_size', - type=int, - default=None, - nargs="+", - help= - 'The attention window size that controls the sliding window attention kv cache behavior' - ) - parser.add_argument( - '--multi_block_mode', - type=lambda s: s.lower() in - ("yes", "true", "t", "1" - ), # custom boolean function to convert input string to boolean - default=True, - help= - "Distribute the work across multiple CUDA thread-blocks on the GPU for masked MHA kernel." - ) - parser.add_argument('--enable_context_fmha_fp32_acc', - action='store_true', - help="Enable FMHA runner FP32 accumulation.") - parser.add_argument('--cuda_graph_mode', - action='store_true', - help="Enable cuda graphs in the inference.") - parser.add_argument( - '--log_level', - type=str, - choices=['verbose', 'info', 'warning', 'error', 'internal_error'], - default='info') - parser.add_argument( - '--no_prompt_template', - dest='use_prompt_template', - default=True, - action='store_false', - help= - "Whether or not to use default prompt template to wrap the input text.") - parser.add_argument('--use_py_session', - default=False, - action='store_true', - help="Whether or not to use Python runtime session") - parser.add_argument('--debug_mode', - default=False, - action='store_true', - help="Whether or not to turn on the debug mode") - parser.add_argument('--streaming', default=False, action='store_true') - parser.add_argument('--streaming_interval', - type=int, - help="How often to return tokens when streaming.", - default=5) - parser.add_argument( - '--prompt_table_path', - type=str, - help="Path to .npy file, exported by nemo_prompt_convert.py") - parser.add_argument( - '--prompt_tasks', - help="Comma-separated list of tasks for prompt tuning, e.g., 0,3,1,0") - parser.add_argument('--lora_dir', - type=str, - default=None, - nargs="+", - help="The directory of LoRA weights") - parser.add_argument('--lora_ckpt_source', - type=str, - default="hf", - choices=["hf", "nemo"], - help="The source of lora checkpoint.") - parser.add_argument( - '--lora_task_uids', - type=str, - default=None, - nargs="+", - help="The list of LoRA task uids; use -1 to disable the LoRA module") - parser.add_argument( - '--num_prepend_vtokens', - nargs="+", - type=int, - help="Number of (default) virtual tokens to prepend to each sentence." - " For example, '--num_prepend_vtokens=10' will prepend the tokens" - " [vocab_size, vocab_size + 1, ..., vocab_size + 9] to the sentence.") - parser.add_argument( - '--draft_target_model_config', - type=str, - default=None, - help= - "Configuration of Draft-Target-Model decoding, see `examples/draft_target_model/README.md` for more information." - " E.g.: [4, [0], [1], False] for [draft_len, draft_model_device_list, target_model_device_list, use_logits]." - ) - parser.add_argument( - '--ngram_config', - type=str, - default=None, - help= - "Configuration of NGram decoding, see `examples/ngram/README.md` for more information." - " E.g.: [10,2,[0]] for [max_draft_len, max_matching_ngram_size, device_list].", - ) - parser.add_argument( - '--medusa_choices', - type=str, - default=None, - help="Configuration of Medusa decoding." - " E.g.: [[0, 0, 0, 0], [0, 1, 0], [1, 0], [1, 1]] for 9 medusa tokens." - ) - parser.add_argument( - '--eagle_choices', - type=str, - default=None, - help="Configuration of Eagle-1 decoding." - " E.g.: [[0, 0, 0, 0], [0, 1, 0], [1, 0], [1, 1]] for 9 draft tokens." - ) - parser.add_argument( - '--eagle_posterior_threshold', - type=float, - default=None, - help="Minimum token probability threshold for typical acceptance. " - "Enables typical acceptance in Eagle. " - "Corresponds to epsilon in https://arxiv.org/pdf/2401.10774.") - parser.add_argument('--eagle_use_dynamic_tree', - action='store_true', - help="Whether to use Ealge-2") - parser.add_argument( - '--eagle_dynamic_tree_max_top_k', - default=None, - type=int, - help= - "The maximum number of draft tokens to expand for each node in Eagle-2." - ) - parser.add_argument( - '--lookahead_config', - type=str, - default=None, - help="Configuration of executor and request lookahead decoding." - " E.g.: [5, 6, 7] for [max_window_size, max_ngram_size, max_verification_set_size]." - ) - # model arguments - parser.add_argument('--engine_dir', type=str, default='engine_outputs') - parser.add_argument( - '--tokenizer_type', - help= - 'Specify that argument when providing a .model file as the tokenizer_dir. ' - 'It allows AutoTokenizer to instantiate the correct tokenizer type.') - parser.add_argument('--vocab_file', - help="Used for sentencepiece tokenizers") - parser.add_argument('--no_add_special_tokens', - dest='add_special_tokens', - default=True, - action='store_false', - help="Whether or not to add special tokens") - parser.add_argument('--hf_model_dir', '--model_dir', type=str, default=None) - parser.add_argument( - '--tokenizer_dir', - default=None, - help='tokenizer path; defaults to hf_model_dir if left unspecified') - - # memory argument - parser.add_argument( - '--gpu_weights_percent', - default=1, - type=float, - help= - 'Specify the percentage of weights that reside on GPU instead of CPU and streaming load during runtime.', - ) - parser.add_argument( - '--max_tokens_in_paged_kv_cache', - default=None, - type=int, - help= - 'Specify the maximum number of tokens in a kv cache page (only available with cpp session).', - ) - parser.add_argument( - '--kv_cache_enable_block_reuse', - default=True, - action=BooleanOptionalAction, - help= - 'Enables block reuse in kv cache (only available with cpp session).', - ) - parser.add_argument( - '--kv_cache_free_gpu_memory_fraction', - default=0.9, - type=float, - help='Specify the free gpu memory fraction.', - ) - parser.add_argument( - '--cross_kv_cache_fraction', - default=0.5, - type=float, - help= - 'Specify the kv cache fraction reserved for cross attention. Only applicable for encoder-decoder models. By default 0.5 for self and 0.5 for cross.', - ) - parser.add_argument( - '--enable_chunked_context', - action='store_true', - help='Enables chunked context (only available with cpp session).', - ) - - # hf model argument (if use hf model) - parser.add_argument( - '--hf_data_type', - '--data_type', - type=str, - choices=['fp32', 'fp16', 'bf16', 'float32', 'float16', 'bfloat16'], - default='fp16', - help="The data type for hf model.") - parser.add_argument( - '--hf_device_map_auto', - action='store_true', - help="Use device map 'auto' to load a pretrained HF model. This may " - "help to test a large model that cannot fit into a singlue GPU.") - - parser.add_argument( - "--return_all_generated_tokens", - default=False, - action="store_true", - help="This option changes the token output only for streaming. " - "If not specified, return only generated tokens at each step. " - "If specified, return the full beams/outputs at each step. " - "It is automatically enabled for num_beams>1 (only available with cpp session). " - "WARNING: using this option may increase network usage significantly (quadratically w.r.t output length)." - ) - - parser.add_argument( - '--language_task_uids', - type=int, - nargs='+', - default=None, - help= - "language task id indicating which adapter to use in language adapter. Please include 1 locale per input text" - ) - parser.add_argument('--backend', type=str, default=None) - - return parser diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index e131587db787..607e5db43c09 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -3097,10 +3097,6 @@ def getMakoArgsFromStageName(stageName, parseSysinfo=false) { // If stageName contains "-PyTorch-", add "backend=pytorch" to makoArgs // At this point, only tests with backend=pytorch or unspecified backend will be run makoArgs += ["backend=pytorch"] - } else if (stageName.contains("-TensorRT-")) { - // If stageName contains "-TensorRT-", add "backend=tensorrt" to makoArgs - // At this point, only tests with backend=tensorrt or unspecified backend will be run - makoArgs += ["backend=tensorrt"] } else if (stageName.contains("-CPP-")) { // If stageName contains "-CPP-", add "backend=cpp" to makoArgs // At this point, only tests with backend=cpp or unspecified backend will be run @@ -3125,7 +3121,7 @@ def getMakoArgsFromStageName(stageName, parseSysinfo=false) { // At this point, only tests with backend=verl or unspecified backend will be run makoArgs += ["backend=verl"] } else { - // If stageName does not contain "-PyTorch-", "-TensorRT-", "-CPP-", "-Triton-", "-FMHA-", "-AutoDeploy-", or "-Verl-", do not add any backend + // If stageName does not contain "-PyTorch-", "-CPP-", "-Triton-", "-FMHA-", "-AutoDeploy-", or "-Verl-", do not add any backend // At this point, all tests will be run // For cases where backend is not specified in makoArgs, we will match all types of backends and tests without specified backend } @@ -4777,9 +4773,9 @@ def launchTestJobs(pipeline, testFilter) x86TestConfigs = [ "CPU-Generic-x86-1": ["cpu", "l0_cpu_x86", 1, 1], "DGX_H100-4_GPUs-CPP-1": ["dgx-h100-x4", "l0_dgx_h100", 1, 1, 4], - "A10-PyTorch-1": ["a10", "l0_a10", 1, 2], - "A10-PyTorch-2": ["a10", "l0_a10", 2, 2], - "A10-TensorRT-1": ["a10", "l0_a10", 1, 1], + "A10-PyTorch-1": ["a10", "l0_a10", 1, 3], + "A10-PyTorch-2": ["a10", "l0_a10", 2, 3], + "A10-PyTorch-3": ["a10", "l0_a10", 3, 3], "A30-PyTorch-1": ["a30", "l0_a30", 1, 2], "A30-PyTorch-2": ["a30", "l0_a30", 2, 2], "A30-CPP-1": ["a30", "l0_a30", 1, 1], @@ -4790,55 +4786,29 @@ def launchTestJobs(pipeline, testFilter) "H100_PCIe-PyTorch-Ray-1": ["h100-cr", "l0_h100", 1, 1], "H100_PCIe-AutoDeploy-1": ["h100-cr", "l0_h100", 1, 1], "H100_PCIe-CPP-1": ["h100-cr", "l0_h100", 1, 1], - "H100_PCIe-TensorRT-1": ["h100-cr", "l0_h100", 1, 1], "RTX5090-PyTorch-1": ["rtx-5090", "l0_gb202", 1, 1], - "RTX5080-TensorRT-1": ["rtx-5080", "l0_gb203", 1, 2], - "RTX5080-TensorRT-2": ["rtx-5080", "l0_gb203", 2, 2], + "RTX5080-PyTorch-1": ["rtx-5080", "l0_gb203", 1, 2], + "RTX5080-PyTorch-2": ["rtx-5080", "l0_gb203", 2, 2], // Currently post-merge test stages only run tests with "stage: post_merge" mako // in the test-db. This behavior may change in the future. - "A10-PyTorch-Post-Merge-1": ["a10", "l0_a10", 1, 1], - "A10-TensorRT-Post-Merge-1": ["a10", "l0_a10", 1, 3], - "A10-TensorRT-Post-Merge-2": ["a10", "l0_a10", 2, 3], - "A10-TensorRT-Post-Merge-3": ["a10", "l0_a10", 3, 3], + "A10-PyTorch-Post-Merge-1": ["a10", "l0_a10", 1, 4], + "A10-PyTorch-Post-Merge-2": ["a10", "l0_a10", 2, 4], + "A10-PyTorch-Post-Merge-3": ["a10", "l0_a10", 3, 4], + "A10-PyTorch-Post-Merge-4": ["a10", "l0_a10", 4, 4], "A10-FMHA-Post-Merge-1": ["a10", "l0_a10", 1, 1], - // "A30-TensorRT-Post-Merge-1": ["a30", "l0_a30", 1, 6], - // "A30-TensorRT-Post-Merge-2": ["a30", "l0_a30", 2, 6], - // "A30-TensorRT-Post-Merge-3": ["a30", "l0_a30", 3, 6], - // "A30-TensorRT-Post-Merge-4": ["a30", "l0_a30", 4, 6], - // "A30-TensorRT-Post-Merge-5": ["a30", "l0_a30", 5, 6], - // "A30-TensorRT-Post-Merge-6": ["a30", "l0_a30", 6, 6], "A30-CPP-Post-Merge-1": ["a30", "l0_a30", 1, 2], "A30-CPP-Post-Merge-2": ["a30", "l0_a30", 2, 2], // "A30-Triton-Post-Merge-1": ["a30", "l0_a30", 1, 2], // "A30-Triton-Post-Merge-2": ["a30", "l0_a30", 2, 2], - // "A100X-TensorRT-Post-Merge-1": ["a100x", "l0_a100", 1, 6], - // "A100X-TensorRT-Post-Merge-2": ["a100x", "l0_a100", 2, 6], - // "A100X-TensorRT-Post-Merge-3": ["a100x", "l0_a100", 3, 6], - // "A100X-TensorRT-Post-Merge-4": ["a100x", "l0_a100", 4, 6], - // "A100X-TensorRT-Post-Merge-5": ["a100x", "l0_a100", 5, 6], - // "A100X-TensorRT-Post-Merge-6": ["a100x", "l0_a100", 6, 6], - // "L40S-TensorRT-Post-Merge-1": ["l40s", "l0_l40s", 1, 5], - // "L40S-TensorRT-Post-Merge-2": ["l40s", "l0_l40s", 2, 5], - // "L40S-TensorRT-Post-Merge-3": ["l40s", "l0_l40s", 3, 5], - // "L40S-TensorRT-Post-Merge-4": ["l40s", "l0_l40s", 4, 5], - // "L40S-TensorRT-Post-Merge-5": ["l40s", "l0_l40s", 5, 5], + "A100X-PyTorch-Post-Merge-1": ["a100x", "l0_a100", 1, 1], + "L40S-PyTorch-Post-Merge-1": ["l40s", "l0_l40s", 1, 1], "L40S-FMHA-Post-Merge-1": ["l40s", "l0_l40s", 1, 1], "H100_PCIe-AutoDeploy-Post-Merge-1": ["h100-cr", "l0_h100", 1, 1], - // "H100_PCIe-TensorRT-Post-Merge-1": ["h100-cr", "l0_h100", 1, 5], - // "H100_PCIe-TensorRT-Post-Merge-2": ["h100-cr", "l0_h100", 2, 5], - // "H100_PCIe-TensorRT-Post-Merge-3": ["h100-cr", "l0_h100", 3, 5], - // "H100_PCIe-TensorRT-Post-Merge-4": ["h100-cr", "l0_h100", 4, 5], - // "H100_PCIe-TensorRT-Post-Merge-5": ["h100-cr", "l0_h100", 5, 5], "H100_PCIe-FMHA-Post-Merge-1": ["h100-cr", "l0_h100", 1, 1], - // "B200_PCIe-TensorRT-Post-Merge-1": ["b100-ts2", "l0_b200", 1, 2], - // "B200_PCIe-TensorRT-Post-Merge-2": ["b100-ts2", "l0_b200", 2, 2], "H100_PCIe-PyTorch-Perf-1": ["h100-cr", "l0_perf", 1, 1], "DGX_H200-8_GPUs-PyTorch-Post-Merge-1": ["dgx-h200-x8", "l0_dgx_h200", 1, 1, 8], "DGX_H200-4_GPUs-PyTorch-Post-Merge-1": ["dgx-h200-x4", "l0_dgx_h200", 1, 1, 4], "DGX_H200-8_GPUs-PyTorch-PerfSanity-Post-Merge-1": ["dgx-h200-x8", "l0_dgx_h200_perf_sanity", 1, 1, 8], - // "DGX_H200-4_GPUs-TensorRT-Post-Merge-1": ["dgx-h200-x4", "l0_dgx_h200", 1, 3, 4], - // "DGX_H200-4_GPUs-TensorRT-Post-Merge-2": ["dgx-h200-x4", "l0_dgx_h200", 2, 3, 4], - // "DGX_H200-4_GPUs-TensorRT-Post-Merge-3": ["dgx-h200-x4", "l0_dgx_h200", 3, 3, 4], // Disable RTXPro6000 stages due to nodes will be offline temporarily. // [TODO] Split tests between RTXPro6000 and RTXPro6000D and move reasonable mount of tests to pre-merge. // "RTXPro6000-PyTorch-Post-Merge-1": ["rtx-pro-6000", "l0_rtx_pro_6000", 1, 1], @@ -4981,7 +4951,7 @@ def launchTestJobs(pipeline, testFilter) // SBSA machines from the Blossom machine pool SBSATestConfigs = [ "CPU-Generic-arm-1": ["cpu", "l0_cpu_arm", 1, 1], - "GH200-TensorRT-Post-Merge-1": ["gh200", "l0_gh200", 1, 1], + "GH200-PyTorch-Post-Merge-1": ["gh200", "l0_gh200", 1, 1], // DGX Spark is also named as GB10 Grace Blackwell Superchip. "GB10-PyTorch-1": ["gb10x", "l0_gb10", 1, 1], ] @@ -5584,7 +5554,6 @@ def launchTestJobs(pipeline, testFilter) def backendMode = testFilter[(TEST_BACKEND)].collect { it.toLowerCase() } def changeMap = [ "pytorch": "-PyTorch-", - "tensorrt": "-TensorRT-", "cpp": "-CPP-", "triton": "-Triton-", "fmha": "-FMHA-", @@ -5611,9 +5580,9 @@ def launchTestJobs(pipeline, testFilter) } else { echo "ONLY_ONE_GROUP_CHANGED mode is true. The group is: ${testFilter[(ONLY_ONE_GROUP_CHANGED)]}." def excludedBackends = new HashMap() - excludedBackends["PyTorch"] = ["-CPP-", "-TensorRT-", "-FMHA-"] // Only pytorch file change also need to run triton tests - excludedBackends["Triton"] = ["-PyTorch-", "-CPP-", "-TensorRT-", "-FMHA-"] - excludedBackends["FMHA"] = ["-PyTorch-", "-CPP-", "-TensorRT-", "-Triton-"] + excludedBackends["PyTorch"] = ["-CPP-", "-FMHA-"] // Only pytorch file change also need to run triton tests + excludedBackends["Triton"] = ["-PyTorch-", "-CPP-", "-FMHA-"] + excludedBackends["FMHA"] = ["-PyTorch-", "-CPP-", "-Triton-"] def group = testFilter[(ONLY_ONE_GROUP_CHANGED)] if (excludedBackends.containsKey(group)) { parallelJobsFiltered = parallelJobsFiltered.findAll { key, value -> diff --git a/jenkins/scripts/cbts/blocks.py b/jenkins/scripts/cbts/blocks.py index 3fc34edda9af..9ce689073682 100644 --- a/jenkins/scripts/cbts/blocks.py +++ b/jenkins/scripts/cbts/blocks.py @@ -437,7 +437,6 @@ def _classify_map_var(var_name: str) -> Optional[str]: # jenkins/L0_Test.groovy (line ~2079). IMPORTANT: keep this list in sync. _BACKEND_PATTERNS = [ ("-PyTorch-", "pytorch"), - ("-TensorRT-", "tensorrt"), ("-CPP-", "cpp"), ("-Triton-", "triton"), ("-FMHA-", "fmha"), diff --git a/scripts/test_to_stage_mapping.py b/scripts/test_to_stage_mapping.py index 04626131b658..24f84c458f51 100644 --- a/scripts/test_to_stage_mapping.py +++ b/scripts/test_to_stage_mapping.py @@ -47,7 +47,7 @@ def _load_tests_file(path: str) -> List[str]: # Regex to parse Jenkins stage configurations from Groovy files -# Matches patterns like: "Stage-Name": ["platform", "yaml_file", split_id, split_count, gpu_count] +# Matches patterns like: "Stage-Name": ["platform", "yaml_file", split_id, split_count, gpu_count, node_count, runWithSbatch] # # Pattern breakdown: # "(?P[^"]+)" - Captures stage name in quotes (group 'stage') @@ -56,10 +56,12 @@ def _load_tests_file(path: str) -> List[str]: # "[^"]+" - Matches platform string in quotes (ignored) # ,\s* - Matches comma with optional whitespace # "(?P[^"]+)" - Captures yaml filename in quotes (group 'yml') -# (?:,\s*\d+)* - Matches zero or more comma-separated numbers (split_id, split_count, gpu_count) +# (?:,\s*(?:\d+|true|false))* - Matches zero or more comma-separated numbers or +# booleans (split_id, split_count, gpu_count, node_count, runWithSbatch) # \s*\] - Matches closing bracket with optional whitespace _STAGE_RE = re.compile( - r'"(?P[^"]+)"\s*:\s*\["[^"]+",\s*"(?P[^"]+)"(?:,\s*\d+)*\s*\]') + r'"(?P[^"]+)"\s*:\s*\["[^"]+",\s*"(?P[^"]+)"(?:,\s*(?:\d+|true|false))*\s*\]' +) def _extract_terms(entry): @@ -85,6 +87,8 @@ def _parse_stage_mapping(path): yaml_to_stages = defaultdict(list) with open(path, 'r') as f: for line in f: + if line.lstrip().startswith('//'): + continue m = _STAGE_RE.search(line) if m: stage = m.group('stage') diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index 67f1c9d3e950..71574f3e1aa5 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -31,7 +31,6 @@ from ctypes import byref from enum import EnumMeta from functools import lru_cache, partial, wraps -from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Sequence, TypeVar, Union import numpy as np @@ -64,7 +63,7 @@ has_nvml = False # isort: on -from tensorrt_llm.bindings import DataType, GptJsonConfig, LayerType +from tensorrt_llm.bindings import DataType, LayerType from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE from tensorrt_llm.logger import logger @@ -775,13 +774,6 @@ def __contains__(cls, item): return True -def supports_inflight_batching(engine_dir): - config_path = Path(engine_dir) / "config.json" - json_config = GptJsonConfig.parse_file(config_path) - model_config = json_config.model_config - return model_config.supports_inflight_batching - - class QuantModeWrapper: def __init__(self, objs): diff --git a/tensorrt_llm/models/automodel.py b/tensorrt_llm/models/automodel.py index 463ae334cad6..6e7ae4041c93 100644 --- a/tensorrt_llm/models/automodel.py +++ b/tensorrt_llm/models/automodel.py @@ -1,7 +1,6 @@ from pathlib import Path from typing import Optional, Union -from ..bindings.executor import DecodingMode from ..mapping import Mapping from . import MODEL_MAP from .modeling_utils import QuantConfig @@ -53,8 +52,7 @@ class AutoModelForCausalLM: @staticmethod def get_trtllm_model_class(hf_model_or_dir: Union[str, Path], - trust_remote_code: bool = False, - decoding_mode: DecodingMode = None): + trust_remote_code: bool = False): import transformers hf_model_or_dir = Path(hf_model_or_dir) if not isinstance( @@ -65,15 +63,8 @@ def get_trtllm_model_class(hf_model_or_dir: Union[str, Path], hf_config = transformers.AutoConfig.from_pretrained( hf_model_or_dir, trust_remote_code=trust_remote_code) - if decoding_mode is not None: - if decoding_mode.isMedusa(): - hf_arch = 'MedusaForCausalLM' - elif decoding_mode.isEagle(): - hf_arch = 'EagleForCausalLM' - else: - raise NotImplementedError(f"Unknown speculative decoding mode.") - elif hasattr(hf_config, - 'architectures') and hf_config.architectures is not None: + if hasattr(hf_config, + 'architectures') and hf_config.architectures is not None: hf_arch = hf_config.architectures[0] elif hasattr(hf_config, 'model_type') and hf_config.model_type.find('mamba') != -1: diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 0bf4298efaad..75e1fd8753b5 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4382,6 +4382,19 @@ def test_auto_dtype(self): task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + @pytest.mark.skip( + reason="TP2 hangs in the AUTO custom allreduce on PCIe-only (all-SYS " + "topology) nodes; the LMHead AllReduce ignores allreduce_strategy and " + "always takes the AUTO path (tunable_allreduce), so the strategy knob " + "cannot work around it. Unskip once validated on an NVLink platform " + "or the lm_head strategy plumbing is fixed.") + @pytest.mark.skip_less_device(2) + def test_tp2(self): + with LLM(self.MODEL_PATH, tensor_parallel_size=2) as llm: + task = CnnDailymail(self.MODEL_NAME) + task.evaluate(llm, + extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + class TestQwen3_4B(LlmapiAccuracyTestHarness): MODEL_NAME = "Qwen3/Qwen3-4B" diff --git a/tests/integration/defs/common.py b/tests/integration/defs/common.py index 03abc6fed72b..f6ddb0d5c7fc 100644 --- a/tests/integration/defs/common.py +++ b/tests/integration/defs/common.py @@ -21,7 +21,6 @@ import tempfile import time from difflib import SequenceMatcher -from pathlib import Path from typing import Any import yaml @@ -33,7 +32,7 @@ from tensorrt_llm.lora_manager import LoraConfig from tensorrt_llm.sampling_params import SamplingParams -from .trt_test_alternative import (check_call, check_output, exists, print_info, +from .trt_test_alternative import (check_call, check_output, print_info, print_warning) @@ -121,428 +120,6 @@ def parse_mpi_cmd(cmd): return cmd -class PluginOptions: - - def __init__(self, - gpt_attention: str = None, - bert_attention: str = None, - gemm: str = None, - layernorm: str = None): - self.gpt_attention = gpt_attention - self.bert_attention = bert_attention - self.gemm = gemm - - def to_legacy_args(self): - args = [] - if self.gpt_attention is not None: - args.extend(["--use_gpt_attention_plugin", self.gpt_attention]) - if self.bert_attention is not None: - args.extend(["--use_bert_attention_plugin", self.bert_attention]) - if self.gemm is not None: - args.extend(["--use_gemm_plugin", self.gemm]) - return args - - def to_args(self): - args = [] - if self.gpt_attention is not None: - args.extend(["--gpt_attention_plugin", self.gpt_attention]) - else: - args.extend(["--gpt_attention_plugin", "disable"]) - if self.bert_attention is not None: - args.extend(["--bert_attention_plugin", self.bert_attention]) - else: - args.extend(["--bert_attention_plugin", "disable"]) - if self.gemm is not None: - args.extend(["--gemm_plugin", self.gemm]) - else: - args.extend(["--gemm_plugin", "disable"]) - return args - - -def prune_checkpoint(llm_venv, checkpoint_dir): - pruned_checkpoint_dir = checkpoint_dir + ".pruned" - prune_cmd = [ - "trtllm-prune", f"--checkpoint_dir={checkpoint_dir}", - f"--out_dir={pruned_checkpoint_dir}" - ] - - check_call(" ".join(prune_cmd), shell=True, env=llm_venv._new_env) - return pruned_checkpoint_dir - - -def refit_model(llm_venv, engine_dir, unpruned_model_dir): - refit_engine_dir = f"{engine_dir}_refit_full" - refit_cmd = [ - "trtllm-refit", f"--checkpoint_dir={unpruned_model_dir}", - f"--engine_dir {engine_dir}", f"--output_dir {refit_engine_dir}" - ] - - check_call(" ".join(refit_cmd), shell=True, env=llm_venv._new_env) - return refit_engine_dir - - -def convert_weights(llm_venv, - example_root, - cmodel_dir, - model, - model_path, - quant_ckpt_path=None, - data_type="float16", - gpus=1, - tp_size=None, - pp_size=None, - model_type=None, - use_parallel_embedding=False, - embedding_sharding_dim=0, - load_by_shard=False, - int8_kv_cache=False, - use_weight_only=False, - workers=1, - processes=None, - smoothquant=0, - per_channel=False, - per_token=False, - fp8_kv_cache=False, - enable_fp8=False, - weight_only_precision=None, - per_group=False, - batch_size=8, - multimodal=False, - ckpt_type='hf', - load_model_on_cpu=False, - **kwargs): - "Convert weights from HF transformers format to FT format" - converted_model_path = os.path.join(cmodel_dir, model, data_type) - script = "convert_checkpoint.py" - - tp_size = gpus if tp_size is None else tp_size - pp_size = gpus // tp_size if pp_size is None else pp_size - gpus = tp_size * pp_size - model_dir = f'{converted_model_path}/{gpus}-gpu' - - # TODO: add other models command - if "gpt2" in model: - script = "convert_checkpoint.py" - convert_cmd = [ - f"{example_root}/{script}", f"--output_dir={model_dir}", - f"--dtype={data_type}", f"--tp_size={tp_size}", - f"--pp_size={pp_size}" - ] - if "next" in model: - convert_cmd.extend(["--nemo_ckpt_path", model_path]) - else: - convert_cmd.extend(["--model_dir", model_path]) - if "smooth" in model: - convert_cmd.extend(["--smoothquant", "0.5"]) - if "kv" in model and "int8" in model: - convert_cmd.append("--int8_kv_cache") - - elif "t5" in model or "bart" in model or "ul2" in model or "wmt" in model or "nougat" in model or 'pix2struct' in model: - assert model_type, "Encoder-Decoder models must specify model architecture type" - script = "convert_checkpoint.py" - convert_cmd = [ - f"{example_root}/{script}", "--model_dir", model_path, - "--output_dir", converted_model_path, f"--model_type={model_type}", - f"--tp_size={tp_size}", f"--pp_size={pp_size}", - f"--dtype={data_type}" - ] - if "nougat" in model: - convert_cmd.append("--nougat") - - model_dir = converted_model_path - - elif "opt" in model and model_type == "blip2": - convert_cmd = [ - f"{example_root}/{script}", - f"--model_dir={model_path}", - f"--output_dir={model_dir}", - f"--model_type={model_type}", - f"--dtype={data_type}", - f"--tp_size={tp_size}", - f"--pp_size={pp_size}", - ] - - elif "whisper" in model_path: - script = "convert_checkpoint.py" - convert_cmd = [ - f"{example_root}/{script}", "--model_dir", model_path, - "--output_dir", converted_model_path - ] - model_dir = converted_model_path - - elif "mamba" in model: - convert_cmd = [ - f"{example_root}/{script}", "--model_dir", model_path, - "--output_dir", model_dir, f"--dtype={data_type}", - f"--tp_size={tp_size}" - ] - - elif "llama" in model or "llava" in model or "vila" in model: - convert_cmd = [ - f"{example_root}/{script}", "--output_dir", model_dir, - f"--dtype={data_type}", f"--tp_size={tp_size}", - f"--pp_size={pp_size}" - ] - - if 'meta-ckpt' in model: - convert_cmd.extend(['--meta_ckpt_dir', model_path]) - else: - convert_cmd.extend(['--model_dir', model_path]) - - if 'code_llama_1gpu' in model: - convert_cmd.extend(['--rotary_base=1000000']) - convert_cmd.extend(['--vocab_size=32016']) - elif 'code_llama' in model: - convert_cmd.extend(['--rotary_base=1000000']) - convert_cmd.extend(['--vocab_size=32000']) - if 'int4_gptq' in model: - convert_cmd.extend([ - "--use_weight_only", "--weight_only_precision=int4_gptq", - f"--quant_ckpt_path={quant_ckpt_path}", "--per_group" - ]) - if 'int8_gptq' in model: - convert_cmd.extend([ - "--use_weight_only", "--weight_only_precision=int8_gptq", - f"--quant_ckpt_path={quant_ckpt_path}", "--per_group", - "--group_size=64" - ]) - - if 'awq' in model: - convert_cmd.extend([ - "--use_weight_only", "--weight_only_precision=int4_awq", - "--group_size=128" - ]) - if 'finegrained_fp8' in model: - convert_cmd.extend(["--use_fp8"]) - - elif "draft_target_model" in model: - if "gpt" in model_path: - example_name = "gpt" - elif "llama" in model_path: - example_name = "llama" - script = f"{example_root}/../models/core/{example_name}/convert_checkpoint.py" - convert_cmd = [ - f"{script}", - "--model_dir", - model_path, - "--output_dir", - model_dir, - f"--dtype={data_type}", - ] - - elif "ngram" in model: - if "gpt" in model_path: - example_name = "gpt" - elif "llama" in model_path: - example_name = "llama" - script = f"{example_root}/../models/core/{example_name}/convert_checkpoint.py" - convert_cmd = [ - f"{script}", - "--model_dir", - model_path, - "--output_dir", - model_dir, - f"--dtype={data_type}", - ] - - elif "medusa" in model: - convert_cmd = [ - f"{example_root}/{script}", "--model_dir", model_path[0], - "--medusa_model_dir", model_path[1], "--output_dir", model_dir, - f"--dtype={data_type}", f"--tp_size={tp_size}", - f"--pp_size={pp_size}", "--num_medusa_heads=4" - ] - elif "redrafter" in model: - redrafter_num_beams = kwargs.pop("redrafter_num_beams") - redrafter_draft_len_per_beam = kwargs.pop( - "redrafter_draft_len_per_beam") - convert_cmd = [ - f"{example_root}/{script}", "--base_model_checkpoint_dir", - model_path[0], "--drafter_model_dir", model_path[1], "--output_dir", - model_dir, f"--dtype={data_type}", f"--tp_size={tp_size}", - f"--redrafter_num_beams={redrafter_num_beams}", - f"--redrafter_draft_len_per_beam={redrafter_draft_len_per_beam}" - ] - elif "eagle" in model: - if len(model_path) == 2: - # Test the checkpoint released from HF, which requires two separate weights, - # one for the base model and one for the EagleNets. - convert_cmd = [ - f"{example_root}/{script}", "--model_dir", model_path[0], - "--eagle_model_dir", model_path[1], "--output_dir", model_dir, - f"--dtype={data_type}", f"--tp_size={tp_size}", - f"--pp_size={pp_size}", "--num_eagle_layers=4", - "--max_draft_len=63", "--max_non_leaves_per_layer=10" - ] - else: - # Test the checkpoint released from ModelOpt, which only requires one weight, - # which includes both the base model and EagleNets, and is an FP8 datatype. - convert_cmd = [ - f"{example_root}/{script}", "--model_dir", model_path, - "--output_dir", model_dir, f"--dtype={data_type}", - f"--tp_size={tp_size}", f"--pp_size={pp_size}", - "--num_eagle_layers=4", "--max_draft_len=63", - "--max_non_leaves_per_layer=10" - ] - elif "recurrentgemma" in model: - convert_cmd = [ - f"{example_root}/{script}", "--model_dir", model_path, - "--output_dir", model_dir, f"--dtype={data_type}", - f"--world_size={tp_size}", f"--ckpt_type={ckpt_type}" - ] - elif "cogvlm" in model: - convert_cmd = [ - f"{example_root}/{script}", "--model_dir", model_path, - "--output_dir", model_dir, f"--dtype={data_type}", - f"--tp_size={tp_size}", f"--pp_size={pp_size}", - "--use_prompt_tuning" - ] - elif "fuyu" in model or "kosmos" in model: - gpt_variant = "kosmos-2" if "kosmos" in model else "persimmon" - convert_cmd = [ - f"{example_root}/{script}", "--model_dir", model_path, - "--output_dir", model_dir, "--dtype", data_type, "--gpt_variant", - gpt_variant - ] - elif "neva-22b" in model: - convert_cmd = [ - f"{example_root}/{script}", "--nemo_ckpt_path", model_path, - "--output_dir", model_dir, "--dtype", data_type, - "--nemo_rename_key", "model:model.language_model", - "attention.linear_qkv.layer_norm_bias:input_layernorm.bias", - "attention.linear_qkv.layer_norm_weight:input_layernorm.weight", - "mlp.linear_fc1.layer_norm_bias:post_attention_layernorm.bias", - "mlp.linear_fc1.layer_norm_weight:post_attention_layernorm.weight", - "linear_qkv:query_key_value", "linear_fc1:dense_h_to_4h", - "linear_fc2:dense_4h_to_h", "linear_proj:dense", "decoder:encoder" - ] - elif "video-neva" in model: - - nemotron_root = os.path.join(example_root, "../", "nemotron") - - if llm_venv: - # Install Python requirements for nemotron - llm_venv.run_cmd([ - "-m", "pip", "install", "-r", - os.path.join(nemotron_root, "requirements.txt") - ]) - - qformat = 'full_prec' - model_name = 'nemotron-video-neva' - converted_model_path = os.path.join(cmodel_dir, model_name, qformat) - model_dir = f'{converted_model_path}/{gpus}-gpu' - # Overwrite the model_path with the nemotron model path - model_path = os.path.join(os.path.dirname(os.path.dirname(model_path)), - 'nemotron', 'Nemotron-4-15B-SteerLM.nemo') - convert_cmd = [ - f"{example_root}/../quantization/quantize.py", - f"--nemo_ckpt_path={model_path}", - "--batch_size=64", - f"--dtype={data_type}", - f"--qformat={qformat}", - f"--output_dir={model_dir}", - ] - elif "dit-xl" in model.lower(): - convert_cmd = [ - f"{example_root}/{script}", - f"--timm_ckpt={model_path}", - f"--output_dir={model_dir}", - f"--dtype={data_type}", - f"--tp_size={tp_size}", - f"--pp_size={pp_size}", - ] - if kwargs.get("enable_fp8_linear") is not None: - convert_cmd.append("--fp8_linear") - elif "stdit" in model.lower(): - convert_cmd = [ - f"{example_root}/{script}", - f"--timm_ckpt={model_path}/model.safetensors", - f"--output_dir={model_dir}", - f"--dtype={data_type}", - f"--tp_size={tp_size}", - f"--pp_size={pp_size}", - ] - elif "bert" in model.lower(): - convert_cmd = [ - f"{example_root}/{script}", - f"--model={model}", - f"--model_dir={model_path}", - f"--output_dir={model_dir}", - f"--dtype={data_type}", - f"--tp_size={tp_size}", - ] - elif "granite" in model.lower(): - convert_cmd = [ - f"{example_root}/{script}", - f"--model_dir={model_path}", - f"--output_dir={model_dir}", - f"--dtype={data_type}", - f"--tp_size={tp_size}", - ] - elif "stable-diffusion-3.5" in model.lower(): - convert_cmd = [ - f"{example_root}/{script}", - f"--model_path={model_path}", - f"--output_dir={model_dir}", - f"--tp_size={tp_size}", - ] - else: - convert_cmd = [ - f"{example_root}/{script}", - f"--model_dir={model_path}", - f"--output_dir={model_dir}", - f"--dtype={data_type}", - f"--tp_size={tp_size}", - f"--pp_size={pp_size}", - ] - - if use_parallel_embedding: - convert_cmd.append("--use_parallel_embedding") - convert_cmd.append(f"--embedding_sharding_dim={embedding_sharding_dim}") - if load_by_shard: - convert_cmd.extend(["--load_by_shard"]) - if load_model_on_cpu: - convert_cmd.extend(["--load_model_on_cpu"]) - if workers > 1: - convert_cmd.extend([f"--workers={workers}"]) - if int8_kv_cache: - convert_cmd.append("--int8_kv_cache") - if use_weight_only: - convert_cmd.append("--use_weight_only") - if weight_only_precision: - convert_cmd.append(f"--weight_only_precision={weight_only_precision}") - if processes is not None: - convert_cmd.append(f"--processes={processes}") - if smoothquant > 0: - convert_cmd.append(f"--smoothquant={smoothquant}") - if per_channel: - convert_cmd.append("--per_channel") - if per_token: - convert_cmd.append("--per_token") - if enable_fp8: - convert_cmd.append('--enable_fp8') - if fp8_kv_cache: - convert_cmd.append('--fp8_kv_cache') - if quant_ckpt_path: - convert_cmd.append(f"--quant_ckpt_path={quant_ckpt_path}") - if per_group: - convert_cmd.append("--per_group") - timeout = kwargs.pop('timeout', None) - - for key, value in kwargs.items(): - if isinstance(value, bool): - if value: - convert_cmd.append(f"--{key}") - else: - convert_cmd.extend([f"--{key}={value}"]) - - if llm_venv: - venv_check_call(llm_venv, convert_cmd, timeout=timeout) - return model_dir - else: - return convert_cmd, model_dir - - def similarity_score(a, b): "similar compare a and b " return SequenceMatcher(None, a, b).ratio() @@ -573,76 +150,6 @@ def generate_summary_cmd(example_root, *args, **kwargs): return summary_cmd -def quantize_data(llm_venv, - example_root, - model_dir, - dtype, - quantize_dir, - qformat="full_prec", - tp_size=1, - pp_size=1, - cp_size=1, - calib_size=512, - kv_cache_dtype=None, - **kwargs): - "quanize data and return data dir" - model_name = os.path.basename(model_dir) - output_dir = os.path.join(quantize_dir, model_name, dtype, qformat, - f"tp{tp_size}pp{pp_size}") - if kv_cache_dtype: - output_dir = os.path.join(output_dir, kv_cache_dtype) - else: - output_dir = os.path.join(output_dir, "no_kv_cache") - - quantize_script = f"{example_root}/../../../quantization/quantize.py" if "core" in example_root else f"{example_root}/../quantization/quantize.py" - quantize_cmd = [ - quantize_script, - f"--model_dir={model_dir}", - f"--dtype={dtype}", - f"--qformat={qformat}", - f"--output_dir={output_dir}", - f"--tp_size={tp_size}", - f"--pp_size={pp_size}", - f"--cp_size={cp_size}", - f"--calib_size={calib_size}", - ] - - if kv_cache_dtype: - quantize_cmd.append(f"--kv_cache_dtype={kv_cache_dtype}") - timeout = kwargs.pop('timeout', None) - - for key, value in kwargs.items(): - if isinstance(value, bool): - if value: - quantize_cmd.append(f"--{key}") - else: - quantize_cmd.extend([f"--{key}", f"{value}"]) - - if llm_venv: - if not exists(output_dir): - venv_check_call(llm_venv, quantize_cmd, timeout=timeout) - return output_dir - else: - return quantize_cmd, output_dir - - -def find_tensorrt(ld_library_path): - MAX_SEARCH_HEIGHT = 10 - ld_library_path = ld_library_path.split(os.pathsep) - for trt_lib_dir in ld_library_path: - trt_lib_dir = Path(trt_lib_dir) - trt_nvinfer_lib = trt_lib_dir / "libnvinfer.so" - if trt_nvinfer_lib.exists(): - trt_root_dir = trt_lib_dir - for i in range(MAX_SEARCH_HEIGHT): - trt_root_dir = trt_root_dir.parent - trt_include_dir = trt_root_dir / "include" - trt_nvinfer_header = trt_include_dir / "NvInfer.h" - if trt_nvinfer_header.exists(): - return str(trt_include_dir), str(trt_lib_dir) - return None, None - - def get_trt_llm_lib_dir(venv): output = venv.run_raw( "import tensorrt_llm; print(f'{tensorrt_llm.__path__[0]}/libs')", diff --git a/tests/integration/defs/conftest.py b/tests/integration/defs/conftest.py index 1b08a55ca16d..1a5beb07cdf4 100644 --- a/tests/integration/defs/conftest.py +++ b/tests/integration/defs/conftest.py @@ -239,15 +239,6 @@ def bert_example_root(llm_root): return example_root -@pytest.fixture(scope="module") -def enc_dec_example_root(llm_root): - "Get encoder-decoder example root" - example_root = os.path.join(llm_root, "examples", "models", "core", - "enc_dec") - - return example_root - - @pytest.fixture(scope="module") def whisper_example_root(llm_root, llm_venv): "Get whisper example root" @@ -260,39 +251,6 @@ def whisper_example_root(llm_root, llm_venv): return example_root -@pytest.fixture(scope="module") -def opt_example_root(llm_root, llm_venv): - "Get opt example root" - - example_root = os.path.join(llm_root, "examples", "models", "contrib", - "opt") - llm_venv.run_cmd([ - "-m", "pip", "install", "-r", - os.path.join(example_root, "requirements.txt") - ]) - - return example_root - - -@pytest.fixture(scope="module") -def llama_example_root(llm_root, llm_venv): - "Get llama example root" - - example_root = os.path.join(llm_root, "examples", "models", "core", "llama") - try: - llm_venv.run_cmd([ - "-m", - "pip", - "install", - "-r", - os.path.join(example_root, "requirements.txt"), - ]) - except: - print("pip install error!") - - return example_root - - @pytest.fixture(scope="module") def llmapi_example_root(llm_root, llm_venv): "Get llm api example root" @@ -367,41 +325,6 @@ def gpt_example_root(llm_root, llm_venv): return example_root -@pytest.fixture(scope="module") -def gptj_example_root(llm_root, llm_venv): - "Get gptj example root" - example_root = os.path.join(llm_root, "examples", "models", "contrib", - "gptj") - llm_venv.run_cmd([ - "-m", "pip", "install", "-r", - os.path.join(example_root, "requirements.txt") - ]) - - return example_root - - -@pytest.fixture(scope="module") -def glm_4_9b_example_root(llm_root, llm_venv): - "Get glm-4-9b example root" - example_root = os.path.join(llm_root, "examples", "models", "core", - "glm-4-9b") - llm_venv.run_cmd([ - "-m", "pip", "install", "-r", - os.path.join(example_root, "requirements.txt") - ]) - - return example_root - - -@pytest.fixture(scope="module") -def exaone_example_root(llm_root, llm_venv): - "Get EXAONE example root" - example_root = os.path.join(llm_root, "examples", "models", "core", - "exaone") - - return example_root - - @pytest.fixture(scope="function") def llm_exaone_model_root(request) -> str: "Get EXAONE model root" @@ -418,26 +341,6 @@ def llm_exaone_model_root(request) -> str: return exaone_model_root -@pytest.fixture(scope="module") -def falcon_example_root(llm_root, llm_venv): - "Get falcon example root" - example_root = os.path.join(llm_root, "examples", "models", "contrib", - "falcon") - llm_venv.run_cmd([ - "-m", "pip", "install", "-r", - os.path.join(example_root, "requirements.txt") - ]) - - return example_root - - -@pytest.fixture(scope="session") -def plugin_gen_path(llm_root): - "Path to the plugin_gen.py script" - return os.path.join(llm_root, "tensorrt_llm", "tools", "plugin_gen", - "plugin_gen.py") - - @pytest.fixture(scope="module") def internlm2_example_root(llm_root, llm_venv): "Get internlm2 example root" @@ -451,42 +354,6 @@ def internlm2_example_root(llm_root, llm_venv): return example_root -@pytest.fixture(scope="module") -def qwen_example_root(llm_root, llm_venv): - "Get qwen example root" - example_root = os.path.join(llm_root, "examples", "models", "core", "qwen") - llm_venv.run_cmd([ - "-m", "pip", "install", "-r", - os.path.join(example_root, "requirements.txt") - ]) - - return example_root - - -@pytest.fixture(scope="module") -def draft_target_model_example_root(llm_root, llm_venv): - "Get Draft-Target-Model example root" - example_root = os.path.join(llm_root, "examples", "draft_target_model") - llm_venv.run_cmd([ - "-m", "pip", "install", "-r", - os.path.join(example_root, "requirements.txt") - ]) - - return example_root - - -@pytest.fixture(scope="module") -def ngram_example_root(llm_root, llm_venv): - "Get NGram example root" - example_root = os.path.join(llm_root, "examples", "ngram") - llm_venv.run_cmd([ - "-m", "pip", "install", "-r", - os.path.join(example_root, "requirements.txt") - ]) - - return example_root - - @pytest.fixture(scope="module") def medusa_example_root(llm_root, llm_venv): "Get medusa example root" @@ -499,93 +366,6 @@ def medusa_example_root(llm_root, llm_venv): return example_root -@pytest.fixture(scope="module") -def redrafter_example_root(llm_root, llm_venv): - "Get ReDrafter example root" - example_root = os.path.join(llm_root, "examples", "redrafter") - llm_venv.run_cmd([ - "-m", "pip", "install", "-r", - os.path.join(example_root, "requirements.txt") - ]) - - return example_root - - -@pytest.fixture(scope="module") -def eagle_example_root(llm_root, llm_venv): - "Get EAGLE example root" - example_root = os.path.join(llm_root, "examples", "eagle") - llm_venv.run_cmd([ - "-m", "pip", "install", "-r", - os.path.join(example_root, "requirements.txt") - ]) - - return example_root - - -@pytest.fixture(scope="module") -def mamba_example_root(llm_root, llm_venv): - "Get mamba example root" - example_root = os.path.join(llm_root, "examples", "models", "core", "mamba") - llm_venv.run_cmd([ - "-m", "pip", "install", "-r", - os.path.join(example_root, "requirements.txt") - ]) - - yield example_root - - llm_venv.run_cmd([ - "-m", "pip", "install", "-r", - os.path.join(llm_root, "requirements.txt") - ]) - - -@pytest.fixture(scope="module") -def nemotron_nas_example_root(llm_root, llm_venv): - example_root = os.path.join(llm_root, "examples", "models", "core", - "nemotron_nas") - - yield example_root - - -@pytest.fixture(scope="module") -def nemotron_example_root(llm_root, llm_venv): - "Get nemotron example root" - example_root = os.path.join(llm_root, "examples", "models", "core", - "nemotron") - llm_venv.run_cmd([ - "-m", "pip", "install", "-r", - os.path.join(example_root, "requirements.txt") - ]) - return example_root - - -@pytest.fixture(scope="module") -def commandr_example_root(llm_root, llm_venv): - "Get commandr example root" - example_root = os.path.join(llm_root, "examples", "models", "core", - "commandr") - llm_venv.run_cmd([ - "-m", "pip", "install", "-r", - os.path.join(example_root, "requirements.txt") - ]) - - return example_root - - -@pytest.fixture(scope="module") -def deepseek_v2_example_root(llm_root, llm_venv): - "Get deepseek v2 example root" - example_root = os.path.join(llm_root, "examples", "models", "contrib", - "deepseek_v2") - llm_venv.run_cmd([ - "-m", "pip", "install", "-r", - os.path.join(example_root, "requirements.txt") - ]) - - return example_root - - @pytest.fixture(scope="function") def deepseek_v3_model_root(request): models_root = llm_models_root() diff --git a/tests/integration/defs/examples/run_llm_fp8_quant_llama_70b.py b/tests/integration/defs/examples/run_llm_fp8_quant_llama_70b.py deleted file mode 100644 index fffe95a4a461..000000000000 --- a/tests/integration/defs/examples/run_llm_fp8_quant_llama_70b.py +++ /dev/null @@ -1,41 +0,0 @@ -import os -from pathlib import Path - -from tensorrt_llm import SamplingParams -from tensorrt_llm._tensorrt_engine import LLM -from tensorrt_llm.llmapi import QuantAlgo, QuantConfig - -prompts = [ - "Hello, my name is", - "The president of the United States is", - "The capital of France is", - "The future of AI is", -] -sampling_params = SamplingParams(temperature=0.8, top_p=0.95) - -model_path = Path( - os.environ.get("LLM_MODELS_ROOT")) / "llama-models-v2/llama-v2-70b-chat-hf" -print(f'model_path: {model_path}') -print(f'gpus: {os.environ.get("CUDA_VISIBLE_DEVICES")}') - - -def main(): - - quant_config = QuantConfig(quant_algo=QuantAlgo.FP8, - kv_cache_quant_algo=QuantAlgo.FP8) - - llm = LLM(model=str(model_path), - quant_config=quant_config, - tensor_parallel_size=2) - - outputs = llm.generate(prompts, sampling_params) - - # Print the outputs. - for output in outputs: - prompt = output.prompt - generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - - -if __name__ == '__main__': - main() diff --git a/tests/integration/defs/examples/test_bert.py b/tests/integration/defs/examples/test_bert.py deleted file mode 100644 index f0268325ea07..000000000000 --- a/tests/integration/defs/examples/test_bert.py +++ /dev/null @@ -1,127 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -"""Module test_bert test bert examples.""" -import pytest -from defs.common import convert_weights, venv_check_call, venv_mpi_check_call -from defs.conftest import get_device_count, get_sm_version -from defs.trt_test_alternative import check_call - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -# # Build parameters -@pytest.mark.parametrize( - "model, hf_bert_model_root", - [("BertModel", 'bert/bert-base-uncased'), - ("BertForQuestionAnswering", 'bert/bert-base-cased-squad2'), - ("BertForSequenceClassification", 'bert/bert-base-uncased-yelp-polarity'), - ("RobertaModel", 'bert/roberta-base'), - ("RobertaForQuestionAnswering", 'bert/roberta-base-squad2'), - ("RobertaForSequenceClassification", 'bert/twitter-roberta-base-emotion')]) -@pytest.mark.parametrize("dtype", ["float32", "float16"]) -@pytest.mark.parametrize("pp_size", [1], ids=lambda pp_size: f'pp:{pp_size}') -@pytest.mark.parametrize("tp_size", [1, 2], ids=lambda tp_size: f'tp:{tp_size}') -@pytest.mark.parametrize( - "use_attention_plugin, context_fmha_type", [(True, 'enabled'), - (True, 'enabled_with_fp32_acc'), - (True, 'disabled'), - (False, 'disabled')], - ids=[ - 'use_attention_plugin-enable_context_fmha', - 'use_attention_plugin-enable_context_fmha_fp32_acc', - 'use_attention_plugin-disable_context_fmha', - 'disable_attention_plugin-disable_context_fmha', - ]) -@pytest.mark.parametrize( - "remove_input_padding", [True, False], - ids=["enable_remove_input_padding", "disable_remove_input_padding"]) -# Run parameters -@pytest.mark.parametrize("compare_hf", [True], ids=["compare_hf"]) -def test_llm_bert_general(bert_example_root, llm_venv, model, dtype, pp_size, - tp_size, use_attention_plugin, context_fmha_type, - hf_bert_model_root, bert_model_root, compare_hf, - cmodel_dir, engine_dir, remove_input_padding): - "Run bert for float16 and float32" - world_size = tp_size * pp_size - - if get_device_count() < world_size: - pytest.skip( - f"Running world size {world_size} on a node with only {get_device_count()} devices. Skip the test..." - ) - - print("Locate model checkpoints in test storage...") - hf_model_name, model_ckpt_path = bert_model_root - - remove_padding = remove_input_padding - if not use_attention_plugin: - remove_padding = False - else: - if get_sm_version() >= 100 and get_sm_version() < 120: - pytest.skip("Attention plugin is not supported on SM100") - - # Convert checkpoints - converted_weight_dir = convert_weights(llm_venv=llm_venv, - example_root=bert_example_root, - cmodel_dir=cmodel_dir, - model=model, - model_path=model_ckpt_path, - data_type=dtype, - tp_size=tp_size) - - # Build Engine - bert_engine_dir = f"{engine_dir}/{model}/{world_size}-gpus/{dtype}/remove_padding_{remove_padding}" - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={converted_weight_dir}", - f"--output_dir={bert_engine_dir}", - "--max_batch_size=8", - ] - - if use_attention_plugin: - build_cmd.append(f"--bert_attention_plugin={dtype}") - else: - build_cmd.append(f"--bert_attention_plugin=disable") - if remove_input_padding and use_attention_plugin: - build_cmd.extend(["--remove_input_padding=enable"]) - else: - build_cmd.extend(["--remove_input_padding=disable"]) - - if context_fmha_type == 'enabled': - build_cmd.extend(["--context_fmha=enable"]) - if context_fmha_type == 'enabled_with_fp32_acc': - build_cmd.extend(["--bert_context_fmha_fp32_acc=enable"]) - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - # Run Engine - print("Run inference...") - run_cmd = [ - f"{bert_example_root}/run.py", - f"--engine_dir={bert_engine_dir}", - f"--hf_model_dir={model_ckpt_path}", - ] - if remove_input_padding and use_attention_plugin: - run_cmd.extend(["--remove_input_padding"]) - if compare_hf: - run_cmd.extend(["--run_hf_test"]) - if world_size == 1: - venv_check_call(llm_venv, run_cmd) - else: - venv_mpi_check_call( - llm_venv, ["mpirun", "-n", - str(world_size), "--allow-run-as-root"], run_cmd) diff --git a/tests/integration/defs/examples/test_bindings.py b/tests/integration/defs/examples/test_bindings.py deleted file mode 100644 index 3d5c21f82bf3..000000000000 --- a/tests/integration/defs/examples/test_bindings.py +++ /dev/null @@ -1,34 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -"""Module test_bindings test bindings examples.""" - -import os - -import pytest -from defs.conftest import get_sm_version - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -@pytest.fixture(scope="module") -def bindings_example_root(llm_root): - "Get bindings example root" - example_root = os.path.join(llm_root, "examples", "bindings", "executor") - - return example_root diff --git a/tests/integration/defs/examples/test_chatglm.py b/tests/integration/defs/examples/test_chatglm.py deleted file mode 100644 index 37ee4a1c1740..000000000000 --- a/tests/integration/defs/examples/test_chatglm.py +++ /dev/null @@ -1,83 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -"""Module test_chatglm test chatglm examples.""" -import os -import shutil - -import pytest -from defs.common import convert_weights, venv_check_call -from defs.conftest import get_sm_version, skip_post_blackwell -from defs.trt_test_alternative import check_call, exists - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -# TODO: add more test case for input_padding, paged_kv_cache, num_beams -@pytest.mark.skip_less_device_memory(24000) -@pytest.mark.parametrize("use_weight_only", - [pytest.param(True, marks=skip_post_blackwell), False], - ids=["enable_weight_only", "disable_weight_only"]) -@pytest.mark.parametrize("llm_glm_4_9b_model_root", - ["glm-4-9b", "glm-4-9b-chat"], - indirect=True) -def test_llm_glm_4_9b_single_gpu_summary(glm_4_9b_example_root, - llm_glm_4_9b_model_root, - llm_datasets_root, llm_rouge_root, - llm_venv, cmodel_dir, engine_dir, - use_weight_only): - "Build & run glm-4-9b on single gpu." - print("Converting checkpoint...") - dtype = 'float16' - model_name = os.path.basename(llm_glm_4_9b_model_root) - ckpt_dir = convert_weights(llm_venv=llm_venv, - example_root=glm_4_9b_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=llm_glm_4_9b_model_root, - data_type=dtype, - use_weight_only=use_weight_only) - - print("Building engines...") - build_cmd = [ - "trtllm-build", f"--checkpoint_dir={ckpt_dir}", - f"--output_dir={engine_dir}", f"--max_batch_size={8}", - f"--max_input_len={924}", f"--max_seq_len={1024}", - f"--gpt_attention_plugin={dtype}" - ] - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Running inference...") - - # fix HF error in glm-4-9b, hope to remove this in the future - # nvbug 5025895 - model_temp_dir = glm_4_9b_example_root + "/model_temp_dir" - if not exists(model_temp_dir): - shutil.copytree(llm_glm_4_9b_model_root, model_temp_dir) - shutil.copy(glm_4_9b_example_root + "/tokenization_chatglm.py", - model_temp_dir) - - summary_cmd = [ - f"{glm_4_9b_example_root}/../../../summarize.py", "--test_trt_llm", - "--hf_model_dir", f"{model_temp_dir}", "--data_type", "fp16", - "--check_accuracy", f"--engine_dir={engine_dir}", - f"--dataset_dir={llm_datasets_root}", f"--rouge_dir={llm_rouge_root}" - ] - - venv_check_call(llm_venv, summary_cmd) diff --git a/tests/integration/defs/examples/test_commandr.py b/tests/integration/defs/examples/test_commandr.py deleted file mode 100644 index bf6c97ec63d8..000000000000 --- a/tests/integration/defs/examples/test_commandr.py +++ /dev/null @@ -1,83 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -"""Module test_commandr test commandr examples.""" -import os - -import pytest -from defs.common import convert_weights, venv_check_call -from defs.conftest import (get_gpu_device_list, get_sm_version, - skip_post_blackwell) -from defs.trt_test_alternative import check_call - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -@pytest.mark.skip_less_device_memory(80000) -@skip_post_blackwell -@pytest.mark.parametrize("use_weight_only", [True, False], - ids=["enable_weight_only", "disable_weight_only"]) -def test_llm_commandr_v01_single_gpu_summary(commandr_example_root, - llm_commandr_v01_model_root, - llm_datasets_root, llm_rouge_root, - llm_venv, cmodel_dir, engine_dir, - use_weight_only): - "Build & run commandr_v01 on single gpu." - if "GH200" in get_gpu_device_list()[0] and not use_weight_only: - pytest.skip("OOM on GH200. https://nvbugs/5250460") - - print("Converting checkpoint...") - dtype = 'float16' - model_name = os.path.basename(llm_commandr_v01_model_root) - ckpt_dir = convert_weights(llm_venv=llm_venv, - example_root=commandr_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=llm_commandr_v01_model_root, - data_type=dtype, - use_weight_only=use_weight_only) - - print("Building engines...") - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={ckpt_dir}", - f"--output_dir={engine_dir}", - f"--max_batch_size={8}", - f"--max_input_len={924}", - f"--max_seq_len={1024}", - f"--gemm_plugin={dtype}", - f"--gpt_attention_plugin={dtype}", - ] - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - summary_cmd = [ - f"{commandr_example_root}/../../../summarize.py", - "--test_trt_llm", - "--hf_model_dir", - f"{llm_commandr_v01_model_root}", - "--data_type", - "fp16", - "--check_accuracy", - f"--engine_dir={engine_dir}", - "--tensorrt_llm_rouge1_threshold=12", - f"--dataset_dir={llm_datasets_root}", - f"--rouge_dir={llm_rouge_root}", - ] - - venv_check_call(llm_venv, summary_cmd) diff --git a/tests/integration/defs/examples/test_draft_target_model.py b/tests/integration/defs/examples/test_draft_target_model.py deleted file mode 100644 index 55fe9da35942..000000000000 --- a/tests/integration/defs/examples/test_draft_target_model.py +++ /dev/null @@ -1,152 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. - -import csv -from copy import deepcopy - -import pytest -from defs.common import convert_weights, venv_check_call -from defs.conftest import get_device_memory, get_sm_version, skip_post_blackwell -from defs.trt_test_alternative import check_call - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -# TODO: remove skip after enable Blackwell for Speculative Decoding -@skip_post_blackwell -@pytest.mark.parametrize("batch_size", [1, 2], ids=['bs1', 'bs2']) -@pytest.mark.parametrize("data_type", ['float16']) -@pytest.mark.parametrize("draft_len", [4, 8], - ids=['draft_len_4', 'draft_len_8']) -@pytest.mark.parametrize("use_logits", [False, True], - ids=['use_tokens', 'use_logits']) -@pytest.mark.parametrize("use_py_session", [False], ids=["use_cpp_session"]) -@pytest.mark.parametrize("draft_target_model_roots", ["gpt2", "llama_v2"], - indirect=True) -@pytest.mark.parametrize("streaming", [False, True], - ids=["no_streaming", "streaming"]) -def test_llm_draft_target_model_1gpu(batch_size, data_type, draft_len, - use_logits, use_py_session, - draft_target_model_roots, streaming, - draft_target_model_example_root, - llm_datasets_root, llm_rouge_root, - llm_venv, cmodel_dir, engine_dir): - if "llama" in draft_target_model_roots[1]: - if get_device_memory() < 80000: - pytest.skip("GPU memory is insufficient.") - - model_name = "draft_target_model" - - print("Build checkpoint ...") - model_dir = convert_weights(llm_venv=llm_venv, - example_root=draft_target_model_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=draft_target_model_roots[1], - data_type=data_type) - - print("Build engines ...") - draft_engine_dir = engine_dir + "-draft" - target_engine_dir = engine_dir + "-target" - baseline_engine_dir = engine_dir + "-baseline" - common_build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={model_dir}", - f"--max_batch_size={batch_size}", - f"--max_beam_width=1", - "--max_input_len=1024", - "--max_seq_len=1536", - "--use_paged_context_fmha=enable", - f"--gpt_attention_plugin={data_type}", - f"--gemm_plugin={data_type}", - "--gather_generation_logits", - ] - draft_model_build_cmd = deepcopy(common_build_cmd) - draft_model_build_cmd.extend([ - f"--output_dir={draft_engine_dir}", - ]) - target_model_build_cmd = deepcopy(common_build_cmd) - target_model_build_cmd.extend([ - f"--output_dir={target_engine_dir}", - "--speculative_decoding_mode=draft_tokens_external", - f"--max_draft_len={draft_len}", - ]) - baseline_model_build_cmd = deepcopy(common_build_cmd) - baseline_model_build_cmd.extend([ - f"--output_dir={baseline_engine_dir}", - ]) - - check_call(" ".join(draft_model_build_cmd), - shell=True, - env=llm_venv._new_env) - check_call(" ".join(target_model_build_cmd), - shell=True, - env=llm_venv._new_env) - check_call(" ".join(baseline_model_build_cmd), - shell=True, - env=llm_venv._new_env) - - print("Run inferences ...") - draft_model_config = f"[{draft_len},[0],[0],{use_logits}]" - common_run_cmd = [ - f"{draft_target_model_example_root}/../run.py", - f"--tokenizer_dir={draft_target_model_roots[1]}", - "--max_output_len=64", - "--kv_cache_enable_block_reuse", - "--kv_cache_free_gpu_memory_fraction=0.25", - ] - if streaming: - common_run_cmd.extend(["--streaming", "--streaming_interval=1"]) - if batch_size == 1: - common_run_cmd.extend(["--input_text", "'How are you?'"]) - elif batch_size == 2: - common_run_cmd.extend(["--input_text", "'Hello'", "'How are you?'"]) - else: - assert False, "Only batch_size <=2 is supported in test." - assert not use_py_session, "Only CPP session is supported in Draft-Target-Model." - - run_cmd = deepcopy(common_run_cmd) - run_cmd.extend([ - f"--engine_dir={target_engine_dir}", - f"--draft_engine_dir={draft_engine_dir}", - f"--draft_target_model_config={draft_model_config}", - f"--output_csv={engine_dir}/draft_target_output.csv", - ]) - baseline_run_cmd = deepcopy(common_run_cmd) - baseline_run_cmd.extend([ - f"--engine_dir={baseline_engine_dir}", - f"--output_csv={engine_dir}/baseline_output.csv", - ]) - - venv_check_call(llm_venv, run_cmd) - venv_check_call(llm_venv, baseline_run_cmd) - - print("Compare outputs ...") - with open(f"{engine_dir}/draft_target_output.csv") as dt_f, open( - f"{engine_dir}/baseline_output.csv") as b_f: - for bs, (dt_request, - b_request) in enumerate(zip(csv.reader(dt_f), - csv.reader(b_f))): - assert ( - len(dt_request) == len(b_request) - ), f"Output length at ({bs=}) is different ({len(dt_request)} v.s. {len(b_request)})." - for index, (dt, b) in enumerate(zip(dt_request, b_request)): - assert ( - int(dt) == int(b) - ), f"Output at ({bs=}, {index=}) is different ({dt} v.s. {b})." diff --git a/tests/integration/defs/examples/test_eagle.py b/tests/integration/defs/examples/test_eagle.py deleted file mode 100644 index 0385edaff354..000000000000 --- a/tests/integration/defs/examples/test_eagle.py +++ /dev/null @@ -1,96 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 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. - -import pytest -from defs.common import convert_weights, venv_check_call -from defs.conftest import get_sm_version, skip_post_blackwell -from defs.trt_test_alternative import check_call - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -@skip_post_blackwell -@pytest.mark.parametrize("use_dynamic_tree", [False, True], - ids=['eagle1', 'eagle2']) -@pytest.mark.parametrize("batch_size", [1, 8], ids=['bs1', 'bs8']) -@pytest.mark.parametrize("data_type", ['float16']) -@pytest.mark.parametrize("eagle_model_roots", ["EAGLE-Vicuna-7B-v1.3"], - indirect=True) -def test_llm_eagle_1gpu(batch_size, data_type, use_dynamic_tree, - eagle_model_roots, eagle_example_root, - llm_datasets_root, llm_rouge_root, llm_venv, cmodel_dir, - engine_dir): - print("Build engines...") - model_name = "eagle" - - model_dir = convert_weights(llm_venv=llm_venv, - example_root=eagle_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=eagle_model_roots, - data_type=data_type) - - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={model_dir}", - f"--output_dir={engine_dir}", - f"--gpt_attention_plugin={data_type}", - f"--gemm_plugin={data_type}", - f"--max_beam_width=1", - "--remove_input_padding=enable", - "--context_fmha=enable", - "--use_paged_context_fmha=enable", - "--max_input_len=1024", - "--max_seq_len=1536", - f"--max_batch_size={batch_size}", - "--paged_kv_cache=enable", - '--speculative_decoding_mode=eagle', - ] - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Run run...") - run_cmd = [ - f"{eagle_example_root}/../run.py", - "--max_output_len=100", - f"--tokenizer_dir={eagle_model_roots[0]}", - "--log_level=verbose", - f"--engine_dir={engine_dir}", - ] - if use_dynamic_tree: - run_cmd.extend( - [f"--eagle_dynamic_tree_max_top_k={3}", "--eagle_use_dynamic_tree"]) - - venv_check_call(llm_venv, run_cmd) - - print("Run summarize...") - summary_cmd = [ - f"{eagle_example_root}/../summarize.py", "--test_trt_llm", - "--hf_model_dir", f"{eagle_model_roots[0]}", "--tokenizer_dir", - f"{eagle_model_roots[0]}", f"--engine_dir={engine_dir}", - "--check_accuracy", "--tensorrt_llm_rouge1_threshold=24", - "--eagle_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]", - f"--max_ite=40", f"--batch_size={batch_size}", - f"--dataset_dir={llm_datasets_root}", f"--rouge_dir={llm_rouge_root}" - ] - if use_dynamic_tree: - summary_cmd.extend( - [f"--eagle_dynamic_tree_max_top_k={3}", "--eagle_use_dynamic_tree"]) - - venv_check_call(llm_venv, summary_cmd) diff --git a/tests/integration/defs/examples/test_enc_dec.py b/tests/integration/defs/examples/test_enc_dec.py deleted file mode 100644 index 10be785b220c..000000000000 --- a/tests/integration/defs/examples/test_enc_dec.py +++ /dev/null @@ -1,340 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. - -import pytest -from defs.common import (convert_weights, quantize_data, venv_check_call, - venv_mpi_check_call) -from defs.conftest import (get_device_count, get_sm_version, skip_fp8_pre_ada, - skip_post_blackwell) -from defs.trt_test_alternative import check_call - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -@pytest.mark.parametrize("use_fp8", [True, False], - ids=["enable_fp8", "disable_fp8"]) -@pytest.mark.parametrize("num_beams", [1, 2, 3], - ids=lambda num_beams: f'nb:{num_beams}') -@pytest.mark.parametrize("pp_size", [1, 2], ids=lambda pp_size: f'pp:{pp_size}') -@pytest.mark.parametrize("tp_size", [1, 2], ids=lambda tp_size: f'tp:{tp_size}') -@pytest.mark.parametrize( - "use_paged_kv_cache", [True, False], - ids=["enable_paged_kv_cache", "disable_paged_kv_cache"]) -@pytest.mark.parametrize( - "use_attention_plugin", - [pytest.param(True, marks=skip_post_blackwell), False], - ids=["enable_attention_plugin", "disable_attention_plugin"]) -@pytest.mark.parametrize("use_gemm_plugin", [True, False], - ids=["enable_gemm_plugin", "disable_gemm_plugin"]) -@pytest.mark.parametrize("data_type", ['bfloat16', 'float16', 'float32']) -@pytest.mark.parametrize("enc_dec_model_root", [ - pytest.param('t5-small', marks=skip_post_blackwell), - pytest.param('flan-t5-small', marks=skip_post_blackwell), - pytest.param('byt5-small', marks=skip_post_blackwell), 'bart-large-cnn', - pytest.param('mbart-large-50-many-to-one-mmt', marks=skip_post_blackwell), - 'wmt14' -], - indirect=True) -@pytest.mark.parametrize("compare_hf_fp32", [True, False], - ids=["compare_hf", "no_compare_hf"]) -def test_llm_enc_dec_general(llm_venv, cmodel_dir, engine_dir, data_type, - use_attention_plugin, use_gemm_plugin, - enc_dec_example_root, enc_dec_model_root, tp_size, - pp_size, num_beams, compare_hf_fp32, - use_paged_kv_cache, use_fp8, llm_datasets_root): - - world_size = tp_size * pp_size - - if get_device_count() < world_size: - pytest.skip( - f"Running world size {world_size} on a node with only {get_device_count()} devices. Skip the test..." - ) - - skip_fp8_pre_ada(use_fp8) - - print("Locate model checkpoints in test storage...") - tllm_model_name, model_ckpt_path = enc_dec_model_root - - print("Converting Encoder-Decoder model into binary format...") - # ckpt from llm_models/ --> cmodels// - model_name = tllm_model_name - model_type = None - if "t5" in model_name or "ul2" in model_name: - if data_type != "float32": - pytest.skip("transformer:issue/34264") - model_type = "t5" - elif "bart" in model_name: - model_type = "bart" - elif "wmt" in model_name: - model_type = "nmt" - - if use_fp8: - assert use_paged_kv_cache and use_attention_plugin - # a known apex huggingface bug for t5 only - # t5 only takes float32 in quantization loop - # https://github.com/huggingface/transformers/issues/34264 - converted_weight_dir = quantize_data( - llm_venv, - enc_dec_example_root, - model_dir=model_ckpt_path, - dtype=data_type, - quantize_dir=cmodel_dir, - qformat="fp8", - tp_size=tp_size, - pp_size=pp_size, - kv_cache_dtype="fp8", - batch_size=1, - calib_dataset=f"{llm_datasets_root}/cnn_dailymail") - - enc_dec_engine_dir = f"{engine_dir}/{tllm_model_name}/{world_size}-gpu/fp8" - else: - converted_weight_dir = convert_weights( - llm_venv=llm_venv, - example_root=enc_dec_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=model_ckpt_path, - data_type=data_type, - tp_size=tp_size, - pp_size=pp_size, - model_type=model_type) - - enc_dec_engine_dir = f"{engine_dir}/{tllm_model_name}/{world_size}-gpu/{data_type}" - - print("Build engines...") - - # change plugins precision to auto if testing fp8 - data_type = "auto" if use_fp8 else data_type - - for component in ["encoder", "decoder"]: - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={converted_weight_dir}/{component}", - f"--output_dir={enc_dec_engine_dir}/{component}", - f"--max_beam_width={num_beams}", - "--moe_plugin=disable", - "--max_batch_size=8", - ] - - if component == "encoder": - build_cmd.append(f"--max_input_len=512") - else: - build_cmd.append(f"--max_input_len=1") - build_cmd.append(f"--max_seq_len=201") - build_cmd.append(f"--max_encoder_input_len=512") - - if use_paged_kv_cache and component == "decoder": - # paged_kv_cache only applies to decoder component - # As for now, we only support num_beams=1 for decoder paged kv cache in python runtime - build_cmd.append(f"--paged_kv_cache=enable") - else: - build_cmd.append(f"--paged_kv_cache=disable") - - if use_gemm_plugin: - build_cmd.append(f"--gemm_plugin={data_type}") - else: - build_cmd.append(f"--gemm_plugin=disable") - - if use_attention_plugin: - # TODO: remove skip after support bert_attention_plugin on B200 - build_cmd.append(f"--bert_attention_plugin={data_type}") - build_cmd.append(f"--gpt_attention_plugin={data_type}") - build_cmd.append("--remove_input_padding=enable") - - # for non-T5 models, FP16/BF16 - if model_type == "t5" or data_type == "float32": - build_cmd.append("--context_fmha=disable") - elif use_fp8: - build_cmd.append("--use_fp8_context_fmha=enable") - else: - build_cmd.append(f"--bert_attention_plugin=disable") - build_cmd.append(f"--gpt_attention_plugin=disable") - build_cmd.append("--remove_input_padding=disable") - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Run inference...") - if use_paged_kv_cache and pp_size == 1: - # use paged engines to cover ModelRunnerCpp tests - run_cmd = [ - f"{enc_dec_example_root}/../../../run.py", - f"--engine_dir={enc_dec_engine_dir}", - f"--tokenizer_dir={model_ckpt_path}", - "--max_output_len=24", - f"--num_beams={num_beams}", - "--input_text='translate English to German: The house is wonderful.'", - ] - else: - # old Python runtime tests - run_cmd = [ - f"{enc_dec_example_root}/run.py", - f"--engine_dir={enc_dec_engine_dir}", - f"--engine_name={model_name}", - f"--model_name={model_ckpt_path}", # use ckpt path so we can use local copy rather than cloning from HF - "--max_new_tokens=24", # shorter than 3rd example input length to capture any bug - f"--num_beams={num_beams}", - ] - if compare_hf_fp32: - run_cmd.extend(["--compare_hf_fp32"]) - - if world_size == 1: - venv_check_call(llm_venv, run_cmd) - else: - venv_mpi_check_call( - llm_venv, ["mpirun", "-n", - str(world_size), "--allow-run-as-root"], run_cmd) - - -@pytest.mark.parametrize("use_fp8", [True, False], - ids=["enable_fp8", "disable_fp8"]) -@pytest.mark.parametrize("num_beams", [1, 2, 3], - ids=lambda num_beams: f'nb:{num_beams}') -@pytest.mark.parametrize("pp_size", [1, 2], ids=lambda pp_size: f'pp:{pp_size}') -@pytest.mark.parametrize("tp_size", [1, 2], ids=lambda tp_size: f'tp:{tp_size}') -@pytest.mark.parametrize("data_type", ['bfloat16', 'float16', 'float32']) -@pytest.mark.parametrize("enc_dec_model_root", ['flan-t5-small', 'flan-t5-xl'], - indirect=True) -def test_llm_enc_dec_mmlu(llm_venv, cmodel_dir, engine_dir, data_type, - enc_dec_example_root, enc_dec_model_root, tp_size, - pp_size, num_beams, mmlu_dataset_root, use_fp8, - llm_datasets_root): - - world_size = tp_size * pp_size - - if get_device_count() < world_size: - pytest.skip( - f"Running world size {world_size} on a node with only {get_device_count()} devices. Skip the test..." - ) - - skip_fp8_pre_ada(use_fp8) - - print("Locate model checkpoints in test storage...") - tllm_model_name, model_ckpt_path = enc_dec_model_root - - print("Converting Encoder-Decoder model into binary format...") - # ckpt from llm_models/ --> cmodels// - model_name = tllm_model_name - model_type = None - if "t5" in model_name or "ul2" in model_name: - model_type = "t5" - elif "bart" in model_name: - model_type = "bart" - elif "wmt" in model_name: - model_type = "nmt" - - if use_fp8: - # a known apex huggingface bug for t5 only - # t5 only takes float32 in quantization loop - # https://github.com/huggingface/transformers/issues/34264 - converted_weight_dir = quantize_data( - llm_venv, - enc_dec_example_root, - model_dir=model_ckpt_path, - dtype=data_type, - quantize_dir=cmodel_dir, - qformat="fp8", - tp_size=tp_size, - pp_size=pp_size, - kv_cache_dtype="fp8", - batch_size=1, - calib_dataset=f"{llm_datasets_root}/cnn_dailymail") - - enc_dec_engine_dir = f"{engine_dir}/{tllm_model_name}/{world_size}-gpu/fp8" - else: - converted_weight_dir = convert_weights( - llm_venv=llm_venv, - example_root=enc_dec_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=model_ckpt_path, - data_type=data_type, - tp_size=tp_size, - pp_size=pp_size, - model_type=model_type) - - enc_dec_engine_dir = f"{engine_dir}/{tllm_model_name}/{world_size}-gpu/{data_type}" - - print("Build engines...") - - max_input_len = 2048 - - # change plugins precision to auto if testing fp8 - data_type = "auto" if use_fp8 else data_type - - for component in ["encoder", "decoder"]: - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={converted_weight_dir}/{component}", - f"--output_dir={enc_dec_engine_dir}/{component}", - f"--max_beam_width={num_beams}", - "--moe_plugin=disable", - "--max_batch_size=8", - ] - - if component == "encoder": - build_cmd.append(f"--max_input_len={max_input_len}") - else: - build_cmd.append(f"--max_input_len=1") - build_cmd.append(f"--max_seq_len=201") - build_cmd.append(f"--max_encoder_input_len={max_input_len}") - - if component == "decoder": - # paged_kv_cache only applies to decoder component - # As for now, we only support num_beams=1 for decoder paged kv cache in python runtime - build_cmd.append(f"--paged_kv_cache=enable") - - build_cmd.append(f"--gemm_plugin={data_type}") - build_cmd.append(f"--bert_attention_plugin={data_type}") - build_cmd.append(f"--gpt_attention_plugin={data_type}") - build_cmd.append("--remove_input_padding=enable") - - # for non-T5 models, FP16/BF16 - if model_type == "t5" or data_type == "float32": - build_cmd.append("--context_fmha=disable") - elif use_fp8: - build_cmd.append("--use_fp8_context_fmha=enable") - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Run MMLU test") - accuracy_threshold_map = { - "flan-t5-xl": { - "float32": 0.440, # 0.444 - }, - "flan-t5-small": { - "float32": 0.280, # 0.282 - "float16": 0.280, # 0.283 - "float8": 0.280, # 0.284 - } - } - precision = "float8" if use_fp8 else data_type - accuracy_threshold = accuracy_threshold_map[tllm_model_name][precision] - - mmlu_cmd = [ - f"{enc_dec_example_root}/../../../mmlu.py", - f"--data_dir={mmlu_dataset_root}", - f"--hf_model_dir={model_ckpt_path}", - "--test_trt_llm", - f"--engine_dir={enc_dec_engine_dir}", - "--kv_cache_free_gpu_memory_fraction=0.45", - "--cross_kv_cache_fraction=0.45", - "--check_accuracy", - f"--accuracy_threshold={accuracy_threshold}", - ] - - venv_check_call(llm_venv, mmlu_cmd) diff --git a/tests/integration/defs/examples/test_flux.py b/tests/integration/defs/examples/test_flux.py deleted file mode 100644 index 9920ae1a85da..000000000000 --- a/tests/integration/defs/examples/test_flux.py +++ /dev/null @@ -1,55 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 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. -"""Integration test for build_and_run_flux.py with multiple quantization formats.""" - -import importlib.util -import os - -import pytest -import torch -from build_and_run_flux import clip_model as load_clip_model - -# Check if CLIP is available -CLIP_AVAILABLE = importlib.util.find_spec("transformers") is not None - - -class FluxTestConfig: - """Configuration for Flux integration test.""" - - MODEL_ID = os.environ.get("FLUX_MODEL_ID", "black-forest-labs/FLUX.1-dev") - PROMPT = "a photo of an astronaut riding a horse on mars" - MIN_CLIP_SIMILARITY = 0.25 - NUM_INFERENCE_STEPS = 20 - MAX_BATCH_SIZE = 1 - BACKEND = "torch-opt" - - # Checkpoint paths for different quantization formats - # These can be set via environment variables or test parameters - FP8_CHECKPOINT = os.environ.get("FLUX_FP8_CHECKPOINT") - FP4_CHECKPOINT = os.environ.get("FLUX_FP4_CHECKPOINT") - - -@pytest.fixture(scope="module") -def clip_model(): - """Pytest fixture for loading CLIP model once per test module.""" - if not CLIP_AVAILABLE: - pytest.skip("CLIP not available") - return load_clip_model() - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for Flux model") -@pytest.mark.slow # Mark as slow test -class TestFluxIntegration: - """Integration tests for Flux model with different quantization formats.""" diff --git a/tests/integration/defs/examples/test_gpt.py b/tests/integration/defs/examples/test_gpt.py index 746c77ecd187..64200b3549db 100644 --- a/tests/integration/defs/examples/test_gpt.py +++ b/tests/integration/defs/examples/test_gpt.py @@ -13,673 +13,16 @@ # See the License for the specific language governing permissions and # limitations under the License. """Module test_gpt test gpt examples.""" -import csv -import os -import re -from pathlib import Path import defs.ci_profiler import pytest -from defs.common import (convert_weights, parse_output, quantize_data, - run_and_check, similar, similarity_score, - test_multi_lora_support, venv_check_call, - venv_check_output, venv_mpi_check_output) -from defs.conftest import (get_device_memory, get_sm_version, skip_fp8_pre_ada, - skip_post_blackwell, skip_pre_ada) -from defs.trt_test_alternative import check_call +from defs.common import similar, similarity_score from tensorrt_llm import LLM from tensorrt_llm.executor.request import LoRARequest from tensorrt_llm.lora_manager import LoraConfig from tensorrt_llm.sampling_params import SamplingParams -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - -INPUT_TEXT_1 = "After Washington had returned to Williamsburg, " + \ - "Dinwiddie ordered him to lead a larger force to assist Trent in his work. " + \ - "While en route, Washington learned of Trent's retreat. " + \ - "Since Tanaghrisson had promised support to the British, " + \ - "Washington continued toward Fort Duquesne and met with the Mingo leader. " + \ - "Learning of a French scouting party in the area, Washington, " + \ - "with Tanaghrisson and his party, surprised the Canadians on May 28 " + \ - "in what became known as the Battle of Jumonville Glen. " + \ - "They killed many of the Canadians, including their commanding officer, " + \ - "Joseph Coulon de Jumonville, whose head was reportedly split open by " + \ - "Tanaghrisson with a tomahawk. The historian Fred Anderson suggests that " + \ - "Tanaghrisson was acting to gain the support of the British and regain " + \ - "authority over his own people. They had been inclined to support the French, " + \ - "with whom they had long trading relationships. One of Tanaghrisson's men told " + \ - "Contrecoeur that Jumonville had been killed by British musket fire. " + \ - "Question: Upon learning of a French scounting party in the area, " + \ - "what did Washington do? Answer:" - -INPUT_TEXT_2 = "You hold the job title in the Wizarding World of Harry Potter where you " + \ - "say random words looking for spells" - -# streaming can can skip outputs, if the next set of outputs arrive. -# this means that the is_equal flag is currently flaky: https://nvbugspro.nvidia.com/bug/4851644 -# assert is_equal - - -@pytest.mark.parametrize("use_gemm_plugin", [True, False], - ids=["enable_gemm_plugin", "disable_gemm_plugin"]) -@pytest.mark.parametrize("use_py_session", [False, True], - ids=["use_cpp_session", "use_py_session"]) -@pytest.mark.parametrize("streaming", [False, True], - ids=["non_streaming", "streaming"]) -def test_llm_gpt2_medium_1gpu(gpt_example_root, llm_venv, - llm_gpt2_medium_model_root, cmodel_dir, - engine_dir, use_gemm_plugin, use_py_session, - streaming): - "gpt2-medium build & run" - print("Converting checkpoint...") - dtype = 'float16' - ckpt_dir = convert_weights(llm_venv=llm_venv, - example_root=gpt_example_root, - cmodel_dir=cmodel_dir, - model="gpt2-medium", - model_path=llm_gpt2_medium_model_root, - data_type=dtype) - - print("Building engines...") - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={ckpt_dir}", - f"--output_dir={engine_dir}", - f"--max_batch_size={8}", - f"--max_input_len={924}", - f"--max_seq_len={1024}", - f"--gpt_attention_plugin={dtype}", - "--paged_kv_cache=enable", - "--remove_input_padding=enable", - ] - - if use_gemm_plugin: - build_cmd.extend([f"--gemm_plugin={dtype}"]) - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - run_cmd = [ - f"{gpt_example_root}/../../../run.py", "--max_output_len=8", - f"--engine_dir={engine_dir}", - f"--tokenizer_dir={llm_gpt2_medium_model_root}", - "--no_add_special_tokens" - ] - - if streaming: - run_cmd.append("--streaming") - if use_py_session: - run_cmd.append("--use_py_session") - - print("Running inference...") - output = venv_check_output(llm_venv, run_cmd) - - valid_outputs = [ - "chef before moving to London in the early", - "chef before moving to London in the late", - "chef and eventually became a chef at a", - ] - - if not streaming: - output = parse_output(output)[0] - assert any([similar(output, expect) - for expect in valid_outputs]), f"output is: {output}" - else: - # Fetch all outputs and expect a monotonically increasing similarity - similarities = [] - for suboutput in parse_output(output): - similarities.append( - max([ - similarity_score(suboutput, expect) - for expect in valid_outputs - ])) - assert ( - all(x <= y for x, y in zip(similarities, similarities[1:])) - ), f"streaming outputs must have a monotonically increasing similarity score. similarities: {similarities}" - output = parse_output(output)[-1] - assert any([similar(output, expect) - for expect in valid_outputs]), f"output is: {output}" - - -@pytest.mark.parametrize("use_py_session", [False, True], - ids=["use_cpp_session", "use_py_session"]) -@pytest.mark.parametrize("streaming", [False, True], - ids=["non_streaming", "streaming"]) -def test_llm_gpt2_medium_bad_words_1gpu(gpt_example_root, llm_venv, - llm_gpt2_medium_model_root, cmodel_dir, - engine_dir, use_py_session, streaming): - "gpt2 build & run" - - if use_py_session and streaming: - pytest.skip( - "Streaming with py session does not return complete sequence to reliably check stop words" - ) - - print("Converting checkpoint...") - dtype = 'float16' - ckpt_dir = convert_weights(llm_venv=llm_venv, - example_root=gpt_example_root, - cmodel_dir=cmodel_dir, - model="gpt2-medium", - model_path=llm_gpt2_medium_model_root, - data_type=dtype) - - print("Building engines...") - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={ckpt_dir}", - f"--output_dir={engine_dir}", - f"--max_batch_size={8}", - f"--max_input_len={924}", - f"--max_seq_len={1024}", - f"--gpt_attention_plugin={dtype}", - "--paged_kv_cache=enable", - "--remove_input_padding=enable", - ] - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - run_cmd = [ - f"{gpt_example_root}/../../../run.py", "--max_output_len=8", - f"--engine_dir={engine_dir}", - f"--tokenizer_dir={llm_gpt2_medium_model_root}", - "--no_add_special_tokens" - ] - - if streaming: - run_cmd.append("--streaming") - if use_py_session: - run_cmd.append("--use_py_session") - - valid_outputs = [ - "chef before moving to the UK in the", - "chef and eventually became a chef at a", - ] - bad_words_args = ["--bad_words", " London"] - run_and_check(llm_venv, - run_cmd + bad_words_args, - valid_outputs, - streaming=streaming) - - bad_words_args = ["--bad_words", " to London", " irrelevant words"] - run_and_check(llm_venv, - run_cmd + bad_words_args, - valid_outputs, - streaming=streaming) - - bad_words_args = ["--bad_words", " irrelevant words", " to London"] - run_and_check(llm_venv, - run_cmd + bad_words_args, - valid_outputs, - streaming=streaming) - - -@pytest.mark.parametrize("use_py_session", [False, True], - ids=["use_cpp_session", "use_py_session"]) -@pytest.mark.parametrize("streaming", [False, True], - ids=["non_streaming", "streaming"]) -def test_llm_gpt2_medium_stop_words_1gpu(gpt_example_root, llm_venv, - llm_gpt2_medium_model_root, cmodel_dir, - engine_dir, use_py_session, streaming): - "gpt2 build & run" - if use_py_session and streaming: - pytest.skip( - "Streaming with py session does not return complete sequence to reliably check stop words" - ) - - print("Converting checkpoint...") - dtype = 'float16' - ckpt_dir = convert_weights(llm_venv=llm_venv, - example_root=gpt_example_root, - cmodel_dir=cmodel_dir, - model="gpt2-medium", - model_path=llm_gpt2_medium_model_root, - data_type=dtype) - - print("Building engines...") - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={ckpt_dir}", - f"--output_dir={engine_dir}", - f"--max_batch_size={8}", - f"--max_input_len={924}", - f"--max_seq_len={1024}", - f"--gpt_attention_plugin={dtype}", - "--paged_kv_cache=enable", - "--remove_input_padding=enable", - ] - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - run_cmd = [ - f"{gpt_example_root}/../../../run.py", "--max_output_len=8", - f"--engine_dir={engine_dir}", - f"--tokenizer_dir={llm_gpt2_medium_model_root}", - "--no_add_special_tokens" - ] - - if streaming: - run_cmd.append("--streaming") - if use_py_session: - run_cmd.append("--use_py_session") - - valid_outputs = [ - "chef before moving to London", - "chef and eventually became", - ] - stop_words_args = ["--stop_words", " London", " became"] - run_and_check(llm_venv, - run_cmd + stop_words_args, - valid_outputs, - streaming=streaming) - - stop_words_args = [ - "--stop_words", " eventually became", " to London", " irrelevant output" - ] - run_and_check(llm_venv, - run_cmd + stop_words_args, - valid_outputs, - streaming=streaming) - - stop_words_args = [ - "--stop_words", " to London", " eventually became", " irrelevant output" - ] - run_and_check(llm_venv, - run_cmd + stop_words_args, - valid_outputs, - streaming=streaming) - - -# transformers compatibility issues -@pytest.mark.parametrize("tensor_parallel", [1, 2], ids=["tp1", "tp2"]) -@pytest.mark.parametrize("use_py_session", [False, True], - ids=["use_cpp_session", "use_py_session"]) -def test_llm_gpt2_next_prompt_tuning(gpt_example_root, llm_venv, - llm_gpt2_next_model_root, cmodel_dir, - engine_dir, tensor_parallel, - use_py_session): - f"gpt-next prompt tuning on {tensor_parallel} gpu(s)" - dtype = "bfloat16" - ckpt_dir = convert_weights(llm_venv=llm_venv, - example_root=gpt_example_root, - cmodel_dir=cmodel_dir, - model="gpt2-next", - model_path=llm_gpt2_next_model_root, - gpus=tensor_parallel, - tp_size=tensor_parallel, - data_type=dtype) - - print("Building engines...") - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={ckpt_dir}", - f"--output_dir={engine_dir}", - f"--max_batch_size=4", - f"--max_input_len=924", - f"--max_seq_len=1024", - f"--gpt_attention_plugin={dtype}", - "--max_prompt_embedding_table_size=200", - ] - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Converting prompt-tuning table...") - squad_table_nemo = Path(llm_gpt2_next_model_root - ).parent / "p-tuning" / "gpt2b_gpt2-squad-vt60.nemo" - squad_table = Path(gpt_example_root) / "prompt_table_squad.npy" - train900_table_nemo = Path( - llm_gpt2_next_model_root - ).parent / "p-tuning" / "gpt2b_gpt2b-train900-v2.nemo" - train900_table = Path(gpt_example_root) / "prompt_table_train900.npy" - for (in_file, out_file) in [(squad_table_nemo, squad_table), - (train900_table_nemo, train900_table)]: - table_conv_cmd = [ - f"{gpt_example_root}/nemo_prompt_convert.py", "-i", - str(in_file), "-o", - str(out_file) - ] - venv_check_call(llm_venv, table_conv_cmd) - - merged_table = Path(gpt_example_root) / "prompt_table_train900.npy" - table_merge_cmd = [ - f"{gpt_example_root}/merge_ptuning_tables.py", - str(squad_table), - str(train900_table), - str(merged_table) - ] - venv_check_call(llm_venv, table_merge_cmd) - - inference_params = { - "squad": { - "num_v_tokens": - 50, - "input": - "Context: In Hinduism the spiritual teacher is known as a guru, and, in many traditions of Hinduism - especially those common in the West - the emphasis on spiritual mentorship is extremely high, with gurus often exercising a great deal of control over the lives of their disciples.\n\nQuestion: Who do gurus control?\n\nAnswer:", - "outputs": [ - "The answer is, of course, the disciple.", - "The guru controls the disciple's life, but", - "The guru is the one who controls the disciple." - ], - }, - "train900": { - "num_v_tokens": 20, - "input": - "Context: Carlsen faced Anand in the World Chess Championship 2013, at Hyatt Regency in Chennai, India, from 9 to 22 November. Carlsen won the match 6.5–3.5 by winning games five, six and nine and drawing the remainder, becoming the new World Chess Champion.\n\nQuestion: When did Carlsen become World Chess Champion?\n\nAnswer:", - "outputs": - ["2013", "2013" + os.linesep + os.linesep + "Question: Who"], - } - } - - print("Running inference...") - - def parse_output(text: str) -> list[str]: - results = [] - while True: - match = re.search( - r"Output \[Text \d+ Beam \d+\]: \"([^\"]*)\"" + os.linesep, - text, re.MULTILINE) - if match is None: - break - _, end = match.span() - results.append(match.group(1)) - text = text[end:] - return results - - # test model without p-tuning dict - run_cmd = [ - f"{gpt_example_root}/../../../run.py", - "--no_add_special_tokens", - "--max_output_len=10", - f"--engine_dir={engine_dir}", - f"--vocab_file={ckpt_dir}/tokenizer.model", - f"--input_text={inference_params['squad']['input']}", - ] - - if use_py_session: - run_cmd.append("--use_py_session") - - output = venv_mpi_check_output( - llm_venv, ["mpirun", "-n", f"{tensor_parallel}", "--allow-run-as-root"], - run_cmd) - assert any( - similar(parse_output(output)[0][:len(ref) + 1], ref) - for ref in inference_params["squad"]["outputs"]), "incorrect output" - - # test p-tuning task separately - run_cmd = [ - f"{gpt_example_root}/../../../run.py", - "--no_add_special_tokens", - "--max_output_len=10", - f"--engine_dir={engine_dir}", - f"--vocab_file={ckpt_dir}/tokenizer.model", - f"--prompt_table={squad_table}", - f"--num_prepend_vtokens={inference_params['squad']['num_v_tokens']}", - f"--input_text={inference_params['squad']['input']}", - f"--no-kv_cache_enable_block_reuse", - ] - - if use_py_session: - run_cmd.append("--use_py_session") - - output = venv_mpi_check_output( - llm_venv, ["mpirun", "-n", f"{tensor_parallel}", "--allow-run-as-root"], - run_cmd) - assert any( - similar(parse_output(output)[0][:len(ref) + 1], ref) - for ref in inference_params["squad"]["outputs"]), "incorrect output" - - run_cmd = [ - f"{gpt_example_root}/../../../run.py", - "--no_add_special_tokens", - "--max_output_len=10", - f"--engine_dir={engine_dir}", - f"--vocab_file={ckpt_dir}/tokenizer.model", - f"--prompt_table={train900_table}", - f"--num_prepend_vtokens={inference_params['train900']['num_v_tokens']}", - f"--input_text={inference_params['train900']['input']}", - f"--no-kv_cache_enable_block_reuse", - ] - - if use_py_session: - run_cmd.append("--use_py_session") - - output = venv_mpi_check_output( - llm_venv, ["mpirun", "-n", f"{tensor_parallel}", "--allow-run-as-root"], - run_cmd) - assert any( - similar(parse_output(output)[0][:len(ref) + 1], ref) - for ref in inference_params["train900"]["outputs"]), "incorrect output" - - # test batched p-tuning tasks - run_cmd = [ - f"{gpt_example_root}/../../../run.py", - "--no_add_special_tokens", - "--max_output_len=10", - f"--engine_dir={engine_dir}", - f"--vocab_file={ckpt_dir}/tokenizer.model", - f"--prompt_table={merged_table}", - f"--num_prepend_vtokens", - str(inference_params['squad']['num_v_tokens']), - str(inference_params['train900']['num_v_tokens']), - f"--prompt_tasks=0,1", - f"--input_text", - inference_params["squad"]["input"], - inference_params['train900']['input'], - f"--no-kv_cache_enable_block_reuse", - ] - - if use_py_session: - run_cmd.append("--use_py_session") - - output = venv_mpi_check_output( - llm_venv, ["mpirun", "-n", f"{tensor_parallel}", "--allow-run-as-root"], - run_cmd) - - outputs = parse_output(output) - assert any( - similar(outputs[0][:len(ref) + 1], ref) - for ref in inference_params["squad"]["outputs"]), "incorrect output" - assert any( - similar(outputs[1][:len(ref) + 1], ref) - for ref in inference_params["train900"]["outputs"]), "incorrect output" - - # test batched and streamed p-tuning tasks - # Streaming with py session does not return complete sequence to reliably check stop words" - - if not use_py_session and tensor_parallel == 1: - run_cmd = [ - f"{gpt_example_root}/../../../run.py", - "--no_add_special_tokens", - "--max_output_len=10", - f"--engine_dir={engine_dir}", - f"--vocab_file={ckpt_dir}/tokenizer.model", - f"--prompt_table={merged_table}", - f"--num_prepend_vtokens", - str(inference_params['squad']['num_v_tokens']), - str(inference_params['train900']['num_v_tokens']), - f"--prompt_tasks=0,1", - "--streaming", - f"--input_text", - inference_params["squad"]["input"], - inference_params['train900']['input'], - f"--no-kv_cache_enable_block_reuse", - ] - - output = venv_mpi_check_output( - llm_venv, - ["mpirun", "-n", f"{tensor_parallel}", "--allow-run-as-root"], - run_cmd) - - outputs = parse_output(output) - squad_outputs = outputs[::2] - train900_outputs = outputs[1::2] - for outputs, valid_outputs in [ - (squad_outputs, inference_params["squad"]["outputs"]), - (train900_outputs, inference_params["train900"]["outputs"]) - ]: - assert any( - similar(outputs[-1][:len(ref) + 1], ref) - for ref in valid_outputs), "incorrect output" - similarities = [] - for suboutput in outputs: - similarities.append( - max([ - similarity_score(suboutput, expect) - for expect in valid_outputs - ])) - assert ( - all(x <= y for x, y in zip(similarities, similarities[1:])) - ), f"streaming outputs must have a monotonically increasing similarity score. valid_outputs: {valid_outputs}, outputs: {outputs}, similarities: {similarities}" - - -@skip_post_blackwell -@pytest.mark.skip_less_device_memory(50000) -@pytest.mark.parametrize("data_type", ['float16', 'fp8'], - ids=['base_fp16', 'base_fp8']) -@pytest.mark.parametrize("lora_data_type", ['float16'], ids=['lora_fp16']) -@pytest.mark.parametrize("llm_gpt2_starcoder_model_root", ['starcoder2'], - indirect=True) -@pytest.mark.parametrize("llm_lora_model_root", - ['peft-lora-starcoder2-15b-unity-copilot'], - indirect=True) -def test_llm_gpt_starcoder_lora_1gpu(data_type, lora_data_type, - gpt_example_root, - llm_gpt2_starcoder_model_root, - llm_datasets_root, llm_venv, cmodel_dir, - engine_dir, llm_lora_model_root, - qcache_dir): - "run starcoder2 lora test on 1gpu" - if data_type == 'fp8': - skip_fp8_pre_ada(use_fp8=True) - else: - if get_device_memory() < 80000: - pytest.skip("GPU memory is not sufficient.") - - print("Converting checkpoint...") - model_name = 'starcoder2-lora' - - if data_type == 'fp8': - model_dir = quantize_data( - llm_venv, - gpt_example_root, - model_dir=llm_gpt2_starcoder_model_root, - calib_dataset=f"{llm_datasets_root}/cnn_dailymail", - dtype="float16", - qformat="fp8", - kv_cache_dtype="fp8", - quantize_dir=qcache_dir, - calib_size=512) - else: - model_dir = convert_weights(llm_venv=llm_venv, - example_root=gpt_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=llm_gpt2_starcoder_model_root) - - print("Build engines...") - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={model_dir}", - f"--output_dir={engine_dir}", - "--lora_plugin=auto", - "--gemm_plugin=auto", - f"--lora_dir={llm_lora_model_root}", - ] - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - ref_1 = [ - 610, 1489, 100, 7670, 100, 5879, 2284, 303, 1489, 459, 8302, 10914, - 16013, 222, 222, 610, 1489, 100, 7670, 100, 5879, 100, 115, 100, 5598, - 45, 115 - ] - ref_2 = [ - 610, 1489, 100, 7670, 100, 5879, 2284, 303, 1489, 459, 8302, 10914, 678, - 222, 222, 610, 1489, 100, 7670, 100, 5879, 100, 115, 100, 5598, 45, 115 - ] - - input_text = "def print_hello_world():" - - print(f"Run inference with lora id 0...") - venv_check_call(llm_venv, [ - f"{gpt_example_root}/../../../run.py", - "--max_output_len=20", - f"--input_text={input_text}", - "--lora_task_uids=0", - f"--tokenizer_dir={llm_gpt2_starcoder_model_root}", - f"--engine_dir={engine_dir}", - f"--output_csv={llm_venv.get_working_directory()}/use_lora.csv", - "--no_add_special_tokens", - "--use_py_session", - ]) - - with open(f"{llm_venv.get_working_directory()}/use_lora.csv") as f: - predict = csv.reader(f) - predict = next(predict) - predict = [int(p) for p in predict] - assert ref_1 == predict or data_type != "float16" - - print(f"Run inference with lora id -1...") - venv_check_call(llm_venv, [ - f"{gpt_example_root}/../../../run.py", - "--max_output_len=20", - f"--input_text={input_text}", - "--lora_task_uids=-1", - f"--tokenizer_dir={llm_gpt2_starcoder_model_root}", - f"--engine_dir={engine_dir}", - f"--output_csv={llm_venv.get_working_directory()}/no_lora.csv", - "--no_add_special_tokens", - "--use_py_session", - ]) - - with open(f"{llm_venv.get_working_directory()}/no_lora.csv") as f: - predict = csv.reader(f) - predict = next(predict) - predict = [int(p) for p in predict] - assert ref_2 == predict or data_type != "float16" - - -@skip_pre_ada -@pytest.mark.parametrize("minitron_model_root", ["4b"], indirect=True) -def test_llm_minitron_fp8_with_pseudo_loras(gpt_example_root, - minitron_model_root, - llm_datasets_root, - llm_venv, - cmodel_dir, - engine_dir, - dtype='bfloat16'): - "Run Minitron model with multiple pseudo LoRAs." - - # Quantize the base model to fp8. - print("Quantizing and converting checkpoint...") - ckpt_dir = f"{cmodel_dir}/minitron/fp8/1-gpu" - - quantize_cmd = [ - f"{gpt_example_root}/../../../quantization/quantize.py", - f"--model_dir={minitron_model_root}", - f"--calib_dataset={llm_datasets_root}/cnn_dailymail", - f"--dtype={dtype}", - "--qformat=fp8", - "--kv_cache_dtype=fp8", - f"--output_dir={ckpt_dir}", - ] - venv_check_call(llm_venv, quantize_cmd) - - test_multi_lora_support( - hf_model_dir=minitron_model_root, - tllm_ckpt_dir=ckpt_dir, - engine_dir=engine_dir, - llm_venv=llm_venv, - example_root=gpt_example_root, - num_loras=2, - lora_rank=8, - target_hf_modules=["q_proj", "k_proj", "v_proj"], - target_trtllm_modules=["attn_q", "attn_k", "attn_v"], - zero_lora_weights=True, - ) - @pytest.mark.skip_less_device_memory( 20000) # Conservative 20GB requirement for GPT-OSS-20B diff --git a/tests/integration/defs/examples/test_gptj.py b/tests/integration/defs/examples/test_gptj.py deleted file mode 100644 index 51a2a1002b6c..000000000000 --- a/tests/integration/defs/examples/test_gptj.py +++ /dev/null @@ -1,27 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. - -import pytest -from defs.conftest import get_sm_version - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - -INPUT_TEXT = """ -Write a Python function `find_max(words)` to solve the following problem:\nWrite a function that accepts a list of strings.\nThe list contains different words. Return the word with maximum number\nof unique characters. If multiple strings have maximum number of unique\ncharacters, return the one which comes first in lexicographical order.\nfind_max(["name", "of", "string"]) == "string"\nfind_max(["name", "enam", "game"]) == "enam"\nfind_max(["aaaaaaa", "bb" ,"cc"]) == ""aaaaaaa" -""" diff --git a/tests/integration/defs/examples/test_granite.py b/tests/integration/defs/examples/test_granite.py deleted file mode 100644 index f789c4065653..000000000000 --- a/tests/integration/defs/examples/test_granite.py +++ /dev/null @@ -1,160 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. - -import os -import time - -import pytest -from defs.common import (convert_weights, test_multi_lora_support, - venv_mpi_check_call) -from defs.conftest import get_sm_version -from defs.trt_test_alternative import check_call - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -@pytest.fixture(scope="module", autouse=True) -def disable_unified_converter(): - os.environ['TRTLLM_DISABLE_UNIFIED_CONVERTER'] = '1' - yield - del os.environ['TRTLLM_DISABLE_UNIFIED_CONVERTER'] - - -@pytest.mark.parametrize("dtype", ["float16", "bfloat16"]) -@pytest.mark.parametrize( - "llm_granite_model_root", - ["granite-3.0-1b-a400m-instruct", "granite-3.0-2b-instruct"], - indirect=True) -def test_llm_granite(llama_example_root, llm_granite_model_root, - llm_datasets_root, llm_rouge_root, llm_venv, cmodel_dir, - engine_dir, dtype): - print("Converting checkpoint...") - model_name = os.path.basename(llm_granite_model_root) - - ckpt_dir = convert_weights( - llm_venv=llm_venv, - example_root=llama_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=llm_granite_model_root, - data_type=dtype, - ) - - print("Building engines...") - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={ckpt_dir}", - f"--output_dir={engine_dir}", - "--max_batch_size=8", - "--max_input_len=924", - "--max_seq_len=1024", - f"--gpt_attention_plugin={dtype}", - f"--gemm_plugin={dtype}", - f"--moe_plugin={dtype}", - f"--workers=1", - ] - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Run engines...") - summary_cmd = [ - f"{llama_example_root}/../../../summarize.py", - f"--engine_dir={engine_dir}", - f"--hf_model_dir={llm_granite_model_root}", - f"--dataset_dir={llm_datasets_root}", - f"--rouge_dir={llm_rouge_root}" - "--test_trt_llm", - "--check_accuracy", - "--tensorrt_llm_rouge1_threshold=25", - "--batch_size=8", - "--max_ite=40", - ] - venv_mpi_check_call(llm_venv, ["mpirun", "-n", "1", "--allow-run-as-root"], - summary_cmd) - - -@pytest.mark.parametrize( - "llm_granite_model_root", - ["granite-3.0-1b-a400m-instruct", "granite-3.0-2b-instruct"], - indirect=True) -def test_granite_bf16_lora(llama_example_root, - llm_datasets_root, - qcache_dir, - llm_rouge_root, - llm_venv, - engine_dir, - cmodel_dir, - llm_granite_model_root, - num_beams=1): - "Run Granite 3.0 models with multiple dummy LoRAs." - - # TODO: Enable fp8 quantization when ModelOpt changes for Granite are available. - start_time = time.time() - print("Converting checkpoint...") - convert_start = time.time() - model_name = os.path.basename(llm_granite_model_root) - dtype = 'bfloat16' - - ckpt_dir = convert_weights( - llm_venv=llm_venv, - example_root=llama_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=llm_granite_model_root, - data_type=dtype, - ) - convert_end = time.time() - print( - f"Convert checkpoint completed in {(convert_end - convert_start):.2f} seconds." - ) - - target_hf_modules = [ - "q_proj", - "k_proj", - "v_proj", - ] - target_trtllm_modules = [ - "attn_q", - "attn_k", - "attn_v", - ] - if model_name == "granite-3.0-1b-a400m-instruct": - target_hf_modules += ["moe_h_to_4h", "moe_4h_to_h", "moe_gate"] - target_trtllm_modules += ["moe_h_to_4h", "moe_4h_to_h", "moe_gate"] - - print("Calling test_multi_lora_support...") - test_multi_lora_start = time.time() - test_multi_lora_support( - hf_model_dir=llm_granite_model_root, - tllm_ckpt_dir=ckpt_dir, - engine_dir=engine_dir, - llm_venv=llm_venv, - example_root=llama_example_root, - num_loras=2, - lora_rank=8, - target_hf_modules=target_hf_modules, - target_trtllm_modules=target_trtllm_modules, - zero_lora_weights=True, - ) - test_multi_lora_end = time.time() - print( - f"test_multi_lora_support completed in {(test_multi_lora_end - test_multi_lora_start):.2f} seconds" - ) - - total_time = time.time() - start_time - print(f"Total function execution time: {total_time:.2f} seconds") diff --git a/tests/integration/defs/examples/test_internlm.py b/tests/integration/defs/examples/test_internlm.py deleted file mode 100644 index 144f86b32fab..000000000000 --- a/tests/integration/defs/examples/test_internlm.py +++ /dev/null @@ -1,104 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -import pytest -from defs.common import convert_weights, parse_mpi_cmd, venv_mpi_check_call -from defs.conftest import get_device_memory, get_sm_version -from defs.trt_test_alternative import check_call - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -# @pytest.mark.skip_less_device(2) -@pytest.mark.parametrize("num_beams", [1, 2, 4], - ids=lambda num_beams: f'nb:{num_beams}') -@pytest.mark.parametrize( - "use_gpt_attention_plugin", [True, False], - ids=["enable_attention_plugin", "disable_attention_plugin"]) -@pytest.mark.parametrize("use_gemm_plugin", [True, False], - ids=["enable_gemm_plugin", "disable_gemm_plugin"]) -@pytest.mark.parametrize("context_fmha_type", [ - "enable_context_fmha", "enable_context_fmha_fp32_acc", - "disable_context_fmha" -]) -@pytest.mark.parametrize("dtype", ['float16', 'bfloat16']) -def test_llm_internlm2_7b_1node_1gpu(internlm2_example_root, - llm_internlm2_7b_model_root, - llm_datasets_root, llm_rouge_root, - llm_venv, cmodel_dir, engine_dir, - use_gpt_attention_plugin, use_gemm_plugin, - context_fmha_type, dtype, num_beams): - "Build & Run internlm2-7b with 1 gpu" - if dtype == "bfloat16" and not use_gemm_plugin: - pytest.skip("Please use gemm plugin when dtype is bfloat16.") - if num_beams == 4 and get_device_memory() < 50000: - pytest.skip("device memory is insufficient.") - - model_dir = convert_weights(llm_venv=llm_venv, - example_root=f"{internlm2_example_root}", - cmodel_dir=cmodel_dir, - model="internlm2-7b", - model_path=llm_internlm2_7b_model_root, - data_type=dtype, - gpus=1, - tp_size=1) - - build_cmd = [ - "python3 -m tensorrt_llm.commands.build", - f"--checkpoint_dir={model_dir}", - f"--output_dir={engine_dir}", - f"--max_beam_width={num_beams}", - f"--max_batch_size=1", - ] - - if use_gpt_attention_plugin: - build_cmd.append("--remove_input_padding=enable") - build_cmd.append(f"--gpt_attention_plugin={dtype}") - else: - build_cmd.append("--gpt_attention_plugin=disable") - build_cmd.append("--remove_input_padding=disable") - build_cmd.append("--paged_kv_cache=disable") - - if use_gemm_plugin: - build_cmd.append(f"--gemm_plugin={dtype}") - else: - build_cmd.append("--gemm_plugin=disable") - - if context_fmha_type == "enable_context_fmha": - build_cmd.append("--context_fmha=enable") - elif context_fmha_type == "disable_context_fmha": - build_cmd.append("--context_fmha=disable") - - print("Building engines...") - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print('Run internlm2-7b...') - data_type = "fp16" if dtype == "float16" else "bf16" - summary_cmd = [ - f"{internlm2_example_root}/../../../summarize.py", "--test_trt_llm", - "--hf_model_dir", llm_internlm2_7b_model_root, "--engine_dir", - engine_dir, "--data_type", data_type, "--check_accuracy", - f"--num_beams={num_beams}", f"--dataset_dir={llm_datasets_root}", - f"--rouge_dir={llm_rouge_root}" - ] - if context_fmha_type == "enable_context_fmha_fp32_acc": - summary_cmd.append("--enable_context_fmha_fp32_acc") - - venv_mpi_check_call( - llm_venv, parse_mpi_cmd(["mpirun", "-n", "1", "--allow-run-as-root"]), - summary_cmd) diff --git a/tests/integration/defs/examples/test_llama.py b/tests/integration/defs/examples/test_llama.py deleted file mode 100644 index ebfcee69c760..000000000000 --- a/tests/integration/defs/examples/test_llama.py +++ /dev/null @@ -1,533 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -import copy -import csv -import os - -import defs.ci_profiler -import pytest -from defs.common import (convert_weights, generate_summary_cmd, quantize_data, - test_llm_torch_multi_lora_support, - test_multi_lora_support, venv_check_call, - venv_mpi_check_call) -# yapf: disable -from defs.conftest import (get_device_count, get_device_memory, - skip_fp8_pre_ada, skip_post_blackwell, skip_pre_ada) -# yapf: enable -from defs.trt_test_alternative import check_call - -# skip trt flow cases on post-Blackwell-Ultra -# if get_sm_version() >= 103: -# pytest.skip( -# "TRT workflow tests are not supported on post Blackwell-Ultra architecture", -# allow_module_level=True) - -INPUT_TEXT_1 = "After Washington had returned to Williamsburg, " + \ - "Dinwiddie ordered him to lead a larger force to assist Trent in his work. " + \ - "While en route, Washington learned of Trent's retreat. " + \ - "Since Tanaghrisson had promised support to the British, " + \ - "Washington continued toward Fort Duquesne and met with the Mingo leader. " + \ - "Learning of a French scouting party in the area, Washington, " + \ - "with Tanaghrisson and his party, surprised the Canadians on May 28 " + \ - "in what became known as the Battle of Jumonville Glen. " + \ - "They killed many of the Canadians, including their commanding officer, " + \ - "Joseph Coulon de Jumonville, whose head was reportedly split open by " + \ - "Tanaghrisson with a tomahawk. The historian Fred Anderson suggests that " + \ - "Tanaghrisson was acting to gain the support of the British and regain " + \ - "authority over his own people. They had been inclined to support the French, " + \ - "with whom they had long trading relationships. One of Tanaghrisson's men told " + \ - "Contrecoeur that Jumonville had been killed by British musket fire. " + \ - "Question: Upon learning of a French scounting party in the area, " + \ - "what did Washington do? Answer:" - -INPUT_TEXT_2 = "Born in north-east France, Soyer trained as a" - - -@pytest.mark.parametrize("num_beams", [1, 2, 4], - ids=lambda num_beams: f'nb:{num_beams}') -@pytest.mark.parametrize("run_type", ['inference', 'summarization']) -@pytest.mark.parametrize("data_type", ['bfloat16', 'float16']) -@pytest.mark.parametrize("fp8_cache", [True, False], - ids=["enable_fp8", "disable_fp8"]) -@pytest.mark.parametrize( - "llama_model_root", - ['llama-v3-8b-instruct-hf', 'llama-3.1-8b-instruct-hf-fp8'], - indirect=True) -def test_llm_llama_1gpu(run_type, data_type, fp8_cache, llama_example_root, - llama_model_root, llm_datasets_root, llm_rouge_root, - llm_venv, cmodel_dir, engine_dir, - qcache_dir_without_install_package, num_beams): - if num_beams > 2 and get_device_memory() < 80000: - pytest.skip("device memory is insufficient.") - - use_fp8 = fp8_cache if "fp8" not in llama_model_root.lower() else True - skip_fp8_pre_ada(use_fp8=use_fp8) - - model_name = os.path.basename(llama_model_root) - - if llama_model_root.endswith('Llama-3.1-8B-Instruct-FP8'): - model_dir = convert_weights(llm_venv=llm_venv, - example_root=llama_example_root, - cmodel_dir=cmodel_dir, - model="llama_v3_finegrained_fp8", - model_path=llama_model_root, - fp8_kv_cache=fp8_cache, - data_type=data_type) - elif fp8_cache: - # Quantize HF llama checkpoint into FP8 format - model_dir = quantize_data( - llm_venv, - llama_example_root, - model_dir=llama_model_root, - calib_dataset=f"{llm_datasets_root}/cnn_dailymail", - dtype=data_type, - qformat="fp8", - quantize_dir=qcache_dir_without_install_package, - calib_size=512, - kv_cache_dtype="fp8") - else: - model_dir = convert_weights( - llm_venv=llm_venv, - example_root=llama_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=llama_model_root, - data_type=data_type, - enable_fp8=fp8_cache, - fp8_kv_cache=fp8_cache, - quant_ckpt_path= - f"{qcache_dir_without_install_package}/quantized_fp8/llama_tp1_rank0.npz" - if fp8_cache else None) - - print("Build engines...") - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={model_dir}", - f"--output_dir={engine_dir}", - f"--gpt_attention_plugin={data_type}", - f"--gemm_plugin={data_type}", - "--remove_input_padding=enable", - f"--max_beam_width={num_beams}", - ] - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - if run_type == "inference": - print("Run inference...") - venv_check_call(llm_venv, [ - f"{llama_example_root}/../run.py", - "--max_output_len=50", - f"--tokenizer_dir={llama_model_root}", - f"--engine_dir={engine_dir}", - f"--num_beams={num_beams}", - ]) - elif run_type == "summarization": - print("Run summarize...") - tensorrt_llm_rouge1_threshold = { - 1: 14, - 2: 19, - 4: 19, - }[num_beams] - - summary_cmd = generate_summary_cmd( - llama_example_root, - hf_model_dir=llama_model_root, - data_type="fp16", - engine_dir=engine_dir, - tensorrt_llm_rouge1_threshold=tensorrt_llm_rouge1_threshold, - num_beams=num_beams, - dataset_dir=llm_datasets_root, - rouge_dir=llm_rouge_root) - - venv_check_call(llm_venv, summary_cmd) - - -@pytest.mark.parametrize( - "data_type", ['float16', 'fp8', 'sq_ootb', 'awq', 'int8_wo'], - ids=['base_fp16', 'base_fp8', 'base_sq_ootb', 'base_awq', 'base_int8_wo']) -@pytest.mark.parametrize("llama_model_root", ['llama-v3-8b-hf'], indirect=True) -@pytest.mark.parametrize("llm_dora_model_root", - ['commonsense-llama-v3-8b-dora-r32'], - indirect=True) -def test_llm_llama_v3_dora_1gpu(data_type, llama_example_root, llama_model_root, - llm_dora_model_root, llm_datasets_root, - llm_venv, cmodel_dir, engine_dir, - qcache_dir_without_install_package): - "run llama dora test on 1gpu" - print("Build engines...") - - model_name = 'llama_v3-dora' - if data_type == 'fp8': - skip_fp8_pre_ada(use_fp8=True) - - model_dir = quantize_data( - llm_venv, - llama_example_root, - model_dir=llama_model_root, - calib_dataset=f"{llm_datasets_root}/cnn_dailymail", - dtype="float16", - qformat="fp8", - quantize_dir=qcache_dir_without_install_package, - calib_size=512, - kv_cache_dtype="fp8") - elif data_type == 'sq_ootb': - model_dir = quantize_data( - llm_venv, - llama_example_root, - model_dir=llama_model_root, - calib_dataset=f"{llm_datasets_root}/cnn_dailymail", - dtype="float16", - qformat="int8_sq", - quantize_dir=qcache_dir_without_install_package, - calib_size=32) - elif data_type == 'awq': - model_dir = quantize_data( - llm_venv, - llama_example_root, - model_dir=llama_model_root, - calib_dataset=f"{llm_datasets_root}/cnn_dailymail", - dtype="float16", - qformat="int4_awq", - awq_block_size=128, - quantize_dir=qcache_dir_without_install_package, - calib_size=32) - elif data_type == 'int8_wo': - model_dir = convert_weights(llm_venv=llm_venv, - example_root=llama_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=llama_model_root, - use_weight_only=True, - weight_only_precision='int8') - else: - model_dir = convert_weights(llm_venv=llm_venv, - example_root=llama_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=llama_model_root) - - # normalize dora magnitude - dora_weights = f"{llm_venv.get_working_directory()}/dora_weights" - - normalize_cmd = [ - f"{llama_example_root}/../../../dora/normalize_weights.py", "-i", - llm_dora_model_root, "-b", llama_model_root, "-o", dora_weights - ] - - venv_check_call(llm_venv, normalize_cmd) - - print("Build engines...") - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={model_dir}", - f"--output_dir={engine_dir}", - "--lora_plugin=auto", - "--dora_plugin=enable", - "--remove_input_padding=enable", # otherwise no cpp runtime - "--kv_cache_type=paged", # otherwise no cpp runtime - "--gemm_plugin=auto", - f"--lora_dir={dora_weights}", - ] - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - input_tokens = [ - 128000, 39314, 374, 459, 7754, 430, 16964, 264, 3465, 11, 35526, 449, - 459, 1988, 430, 5825, 4726, 2317, 13, 9842, 264, 2077, 430, 36001, - 45695, 279, 1715, 382, 394, 17010, 30151, 512, 394, 5321, 5268, 279, - 4495, 4320, 311, 279, 3488, 25, 578, 842, 1121, 304, 279, 1920, 315, - 7397, 74767, 374, 279, 5788, 315, 13465, 323, 24463, 13, 16299, 3094, - 17738, 279, 7314, 315, 7397, 74767, 31931, 16533, 16, 25, 36424, 4907, - 374, 42101, 1555, 279, 20282, 13, 22559, 17, 25, 8828, 4907, 374, 16489, - 311, 11742, 4907, 13, 22559, 18, 25, 92479, 5237, 25734, 304, 279, - 16312, 41255, 3177, 4907, 13, 22559, 19, 25, 8219, 4238, 374, 16489, - 1139, 37833, 5237, 25734, 4286, 16533, 3645, 25, 4320, 16, 14, 9399, 17, - 14, 9399, 18, 14, 9399, 19, 271, 394, 17010, 5688, 512, 72348, 394, - 17010, 6075, 1473 - ] - - out_ref = [ - 128000, 39314, 374, 459, 7754, 430, 16964, 264, 3465, 11, 35526, 449, - 459, 1988, 430, 5825, 4726, 2317, 13, 9842, 264, 2077, 430, 36001, - 45695, 279, 1715, 382, 394, 17010, 30151, 512, 394, 5321, 5268, 279, - 4495, 4320, 311, 279, 3488, 25, 578, 842, 1121, 304, 279, 1920, 315, - 7397, 74767, 374, 279, 5788, 315, 13465, 323, 24463, 13, 16299, 3094, - 17738, 279, 7314, 315, 7397, 74767, 31931, 16533, 16, 25, 36424, 4907, - 374, 42101, 1555, 279, 20282, 13, 22559, 17, 25, 8828, 4907, 374, 16489, - 311, 11742, 4907, 13, 22559, 18, 25, 92479, 5237, 25734, 304, 279, - 16312, 41255, 3177, 4907, 13, 22559, 19, 25, 8219, 4238, 374, 16489, - 1139, 37833, 5237, 25734, 4286, 16533, 3645, 25, 4320, 16, 14, 9399, 17, - 14, 9399, 18, 14, 9399, 19, 271, 394, 17010, 5688, 512, 72348, 394, - 17010, 6075, 1473, 394, 279, 4495, 4320, 374, 4320, 18, 128001, 128001, - 128001, 128001, 128001, 128001, 128001, 128001, 128001, 128001, 128001, - 128001, 128001, 128001, 128001, 128001, 128001, 128001, 128001, 128001, - 128001, 128001, 128001, 128001, 128001 - ] - - in_csv = f"{llm_venv.get_working_directory()}/input.csv" - out_csv = f"{llm_venv.get_working_directory()}/output.csv" - with open(in_csv, "w") as f: - writer = csv.writer(f) - writer.writerow(input_tokens) - - base_run_cmd = [ - f"{llama_example_root}/../../../run.py", "--max_output_len=20", - f"--input_file={in_csv}", f"--tokenizer_dir={llama_model_root}", - f"--engine_dir={engine_dir}", "--max_output_len=32" - ] - - for use_py_session in [True, False]: - if use_py_session: - print("Run inference with Python runtime...") - else: - print("Run inference with C++ runtime...") - - print(f"Run inference with lora id 0...") - run_cmd = copy.deepcopy(base_run_cmd) - run_cmd.extend(["--lora_task_uids=0", f"--output_csv={out_csv}"]) - if use_py_session: - run_cmd.append("--use_py_session") - venv_check_call(llm_venv, run_cmd) - - with open(out_csv) as f: - predict = csv.reader(f) - predict = next(predict) - - predict = [int(p) for p in predict] - assert out_ref == predict or data_type != "float16" - - -@pytest.mark.skip_less_device(8) -@pytest.mark.skip_less_device_memory(80000) -@pytest.mark.parametrize("num_beams", [1, 4], - ids=lambda num_beams: f'nb:{num_beams}') -@pytest.mark.parametrize( - "tp_pp_size", [(8, 1), (4, 2)], - ids=lambda tp_pp_size: f'tp{tp_pp_size[0]}pp{tp_pp_size[1]}') -@pytest.mark.parametrize("test_case", ["pg64317"], indirect=True) -def test_llm_llama_long_alpaca_8gpu_summary(llama_example_root, - llm_long_alpaca_model_root, - llm_datasets_root, llm_rouge_root, - llm_venv, cmodel_dir, engine_dir, - num_beams, tp_pp_size, test_case): - "llama test for long alpaca" - tp_size, pp_size = tp_pp_size - world_size = 8 - assert tp_size * pp_size == world_size, \ - f'tp_size({tp_size}) x pp_size({pp_size}) != 8' - - model_name = 'llama_long_alpaca' - model_dir = convert_weights(llm_venv=llm_venv, - example_root=llama_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=llm_long_alpaca_model_root, - gpus=world_size, - tp_size=tp_size, - pp_size=pp_size, - data_type="bfloat16") - - print("Build engines...") - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={model_dir}", - f"--output_dir={engine_dir}", - "--gpt_attention_plugin=bfloat16", - "--remove_input_padding=enable", - "--gemm_plugin=bfloat16", - f"--max_beam_width={num_beams}", - "--max_input_len=32768", - "--max_seq_len=49152", - "--max_batch_size=1", - "--max_num_tokens=32768", - ] - print("Build engines...") - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Run summarize...") - max_output_len = test_case["max_output_len"] - run_cmd = [ - f"{llama_example_root}/../../../run.py", - f"--max_output_len={max_output_len}", - f"--input_file={test_case['input_file']}", f"--engine_dir={engine_dir}", - f"--num_beams={num_beams}", - f"--tokenizer_dir={llm_long_alpaca_model_root}", - "--max_input_length=32768" - ] - - venv_mpi_check_call( - llm_venv, ["mpirun", "-n", f"{world_size}", "--allow-run-as-root"], - run_cmd) - - summary_cmd = generate_summary_cmd(llama_example_root, - hf_model_dir=llm_long_alpaca_model_root, - max_input_length=16384, - output_len=max_output_len, - data_type="fp16", - num_beams=num_beams, - engine_dir=engine_dir, - dataset_dir=llm_datasets_root, - rouge_dir=llm_rouge_root) - - venv_mpi_check_call( - llm_venv, ["mpirun", "-n", f"{world_size}", "--allow-run-as-root"], - summary_cmd) - - -@pytest.mark.parametrize("fp8_quant", [ - 'disable_fp8', - pytest.param('enable_fp8', marks=skip_post_blackwell), - pytest.param('enable_fp8_meta_recipe', marks=skip_post_blackwell) -]) -@pytest.mark.parametrize("llama_model_root", ['llama-3.1-8b', 'llama-3.2-1b'], - indirect=True) -def test_llm_llama_v3_1_1node_single_gpu(llama_example_root, llama_model_root, - llm_venv, cmodel_dir, - llm_datasets_root, llm_rouge_root, - engine_dir, fp8_quant): - "Run llama3.1 test on 1 gpu." - data_type = "bfloat16" - model_name = os.path.basename(llama_model_root) - - use_fp8_rowwise = False - use_meta_fp8_rowwise_recipe = False - if fp8_quant == 'enable_fp8': - use_fp8_rowwise = True - elif fp8_quant == 'enable_fp8_meta_recipe': - use_fp8_rowwise = True - use_meta_fp8_rowwise_recipe = True - - print("Convert weight...") - model_dir = convert_weights( - llm_venv=llm_venv, - example_root=llama_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=llama_model_root, - data_type=data_type, - tp_size=1, - pp_size=1, - use_fp8_rowwise=use_fp8_rowwise, - use_meta_fp8_rowwise_recipe=use_meta_fp8_rowwise_recipe) - - print("Build engines...") - build_cmd = [ - "trtllm-build", f"--checkpoint_dir={model_dir}", - f"--output_dir={engine_dir}", f"--max_batch_size={8}", - f"--max_seq_len={2048}" - ] - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Run summarize...") - summary_cmd = [ - f"{llama_example_root}/../../../summarize.py", - "--test_trt_llm", - f"--hf_model_dir={llama_model_root}", - f"--engine_dir={engine_dir}", - "--check_accuracy", - f"--tensorrt_llm_rouge1_threshold={14}", - f"--dataset_dir={llm_datasets_root}", - f"--rouge_dir={llm_rouge_root}", - ] - venv_check_call(llm_venv, summary_cmd) - - -@skip_pre_ada -@pytest.mark.skip_less_device_memory(80000) -@pytest.mark.parametrize( - "llama_model_root", - ['llama-v3-8b-instruct-hf', 'llama-3.1-8b', 'llama-3.2-1b', 'llama-3.2-3b'], - indirect=True) -def test_llama_3_x_fp8_with_bf16_lora(llama_example_root, llm_datasets_root, - qcache_dir_without_install_package, - llm_venv, engine_dir, llama_model_root): - "Run Llama 3.1 and 3.2 models with multiple dummy LoRAs." - - print("Quantizing model to fp8...") - - defs.ci_profiler.start("quantize_model") - qmodel_dir = quantize_data( - llm_venv, - llama_example_root, - model_dir=llama_model_root, - calib_dataset=f"{llm_datasets_root}/cnn_dailymail", - dtype="bfloat16", - qformat="fp8", - quantize_dir=qcache_dir_without_install_package, - calib_size=32, - kv_cache_dtype="fp8") - defs.ci_profiler.stop("quantize_model") - print( - f"quantize_model: {defs.ci_profiler.elapsed_time_in_sec('quantize_model')} sec" - ) - - defs.ci_profiler.start("test_multi_lora_support") - test_multi_lora_support( - hf_model_dir=llama_model_root, - tllm_ckpt_dir=qmodel_dir, - engine_dir=engine_dir, - llm_venv=llm_venv, - example_root=llama_example_root, - num_loras=2, - lora_rank=8, - target_hf_modules=["q_proj", "k_proj", "v_proj"], - target_trtllm_modules=["attn_q", "attn_k", "attn_v"], - zero_lora_weights=True, - ) - defs.ci_profiler.stop("test_multi_lora_support") - print( - f"test_multi_lora_support: {defs.ci_profiler.elapsed_time_in_sec('test_multi_lora_support')} sec" - ) - - -@skip_pre_ada -@pytest.mark.skip_less_device_memory(80000) -@pytest.mark.parametrize("llama_model_root", [ - 'llama-v3-8b-instruct-hf', - 'llama-3.1-8b-instruct', - 'llama-3.2-1b-instruct', - 'llama-3.2-3b-instruct', - 'llama-3.3-70b-instruct', -], - indirect=True) -def test_llama_3_x_with_bf16_lora_torch(llama_example_root, llm_datasets_root, - qcache_dir_without_install_package, - llm_venv, engine_dir, llama_model_root): - """Run Llama models with multiple dummy LoRAs using LLM-API Torch backend.""" - - if "llama-3.3-70b-instruct" in llama_model_root.lower(): - tensor_parallel_size = 8 - if get_device_count() < 8: - pytest.skip( - "Skipping: llama-3.3-70b-instruct model requires 8 GPUs") - else: - tensor_parallel_size = 1 - - print("Testing with LLM-API Torch backend...") - - defs.ci_profiler.start("test_llm_torch_multi_lora_support") - - test_llm_torch_multi_lora_support( - hf_model_dir=llama_model_root, - llm_venv=llm_venv, - num_loras=2, - lora_rank=8, - target_hf_modules=["q_proj", "k_proj", "v_proj"], - zero_lora_weights=True, - tensor_parallel_size=tensor_parallel_size) - defs.ci_profiler.stop("test_llm_torch_multi_lora_support") - print( - f"test_llm_torch_multi_lora_support: {defs.ci_profiler.elapsed_time_in_sec('test_llm_torch_multi_lora_support')} sec" - ) diff --git a/tests/integration/defs/examples/test_mamba.py b/tests/integration/defs/examples/test_mamba.py deleted file mode 100644 index c771278bdbe2..000000000000 --- a/tests/integration/defs/examples/test_mamba.py +++ /dev/null @@ -1,123 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. - -import os - -import pytest -from defs.common import (convert_weights, generate_summary_cmd, venv_check_call, - venv_mpi_check_call) -from defs.conftest import get_sm_version, skip_post_blackwell -from defs.trt_test_alternative import check_call - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -@pytest.mark.parametrize("gemm_plugin", [True, False], - ids=["enable_gemm_plugin", "disable_gemm_plugin"]) -@pytest.mark.parametrize("dtype", ['bfloat16', 'float16']) -@pytest.mark.parametrize("mamba_model_root", [ - pytest.param('mamba-130m', marks=skip_post_blackwell), 'mamba-2.8b', - 'mamba-1.4b', 'mamba-790m', 'mamba-370m', 'mamba2-130m', 'mamba2-2.7b', - 'mamba2-1.3b', 'mamba2-780m', 'mamba2-370m', - pytest.param('mamba-codestral-7B-v0.1', marks=skip_post_blackwell) -], - indirect=True) -def test_llm_mamba_1gpu(mamba_example_root, mamba_model_root, - llm_gptneox_model_root, llm_mathstral_model_root, - llm_datasets_root, llm_rouge_root, llm_venv, - gemm_plugin, dtype, cmodel_dir, engine_dir): - "Build & Run mamba model with one gpu" - print("Build engines...") - - model_name = os.path.basename(mamba_model_root) - model_dir = convert_weights(llm_venv=llm_venv, - example_root=mamba_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=mamba_model_root, - data_type=dtype) - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={model_dir}", - f"--output_dir={engine_dir}", - "--paged_kv_cache=disable", - "--max_batch_size=8", - ] - if gemm_plugin: - build_cmd.append("--gemm_plugin=auto") - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print(f'Run {model_name}...') - tokenizer_dir = llm_mathstral_model_root if model_name == "mamba-codestral-7B-v0.1" else llm_gptneox_model_root - summary_cmd = generate_summary_cmd(mamba_example_root, - hf_model_dir=mamba_model_root, - tokenizer_dir=tokenizer_dir, - data_type=dtype, - engine_dir=engine_dir, - batch_size=8, - tensorrt_llm_rouge1_threshold="13.5", - dataset_dir=llm_datasets_root, - rouge_dir=llm_rouge_root) - - venv_check_call(llm_venv, summary_cmd) - - -@pytest.mark.parametrize("mamba_model_root", ['mamba-codestral-7B-v0.1'], - indirect=True) -def test_llm_mamba2_2gpu(mamba_example_root, mamba_model_root, - llm_gptneox_model_root, llm_mathstral_model_root, - llm_datasets_root, llm_rouge_root, llm_venv, - cmodel_dir, engine_dir): - "Build & Run mamba2 model with two gpus" - print("Build engines...") - - model_name = mamba_model_root.split('/')[-1] - model_dir = convert_weights(llm_venv=llm_venv, - example_root=mamba_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=mamba_model_root, - data_type='float16', - tp_size=2) - build_cmd = [ - "trtllm-build", - "--gemm_plugin=auto", - f"--checkpoint_dir={model_dir}", - f"--output_dir={engine_dir}", - "--paged_kv_cache=disable", - "--max_batch_size=8", - ] - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print(f'Run {model_name}...') - tokenizer_dir = llm_mathstral_model_root - summary_cmd = generate_summary_cmd(mamba_example_root, - hf_model_dir=mamba_model_root, - tokenizer_dir=tokenizer_dir, - data_type='float16', - engine_dir=engine_dir, - batch_size=8, - tensorrt_llm_rouge1_threshold="19.0", - dataset_dir=llm_datasets_root, - rouge_dir=llm_rouge_root) - - venv_mpi_check_call(llm_venv, ["mpirun", "-n", "2", "--allow-run-as-root"], - summary_cmd) diff --git a/tests/integration/defs/examples/test_medusa.py b/tests/integration/defs/examples/test_medusa.py deleted file mode 100644 index 49975a591569..000000000000 --- a/tests/integration/defs/examples/test_medusa.py +++ /dev/null @@ -1,156 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. - -import os - -import pytest -from defs.common import convert_weights, venv_check_call -from defs.conftest import get_sm_version, skip_post_blackwell -from defs.trt_test_alternative import check_call - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -@skip_post_blackwell -@pytest.mark.parametrize("batch_size", [1, 8], ids=['bs1', 'bs8']) -@pytest.mark.parametrize("data_type", ['bfloat16']) -@pytest.mark.parametrize("num_medusa_heads", [4], ids=['4-heads']) -@pytest.mark.parametrize("medusa_model_roots", ["medusa-vicuna-7b-v1.3"], - indirect=True) -@pytest.mark.parametrize("use_py_session", [False, True], - ids=["use_cpp_session", "use_py_session"]) -def test_llm_medusa_1gpu(batch_size, data_type, medusa_model_roots, - medusa_example_root, llm_datasets_root, llm_rouge_root, - num_medusa_heads, llm_venv, cmodel_dir, engine_dir, - use_py_session): - print("Build engines...") - model_name = "medusa" - - model_dir = convert_weights(llm_venv=llm_venv, - example_root=medusa_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=medusa_model_roots, - data_type=data_type) - - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={model_dir}", - f"--output_dir={engine_dir}", - f"--gpt_attention_plugin={data_type}", - f"--gemm_plugin={data_type}", - f"--max_beam_width=1", - "--remove_input_padding=enable", - "--context_fmha=enable", - "--max_input_len=1024", - "--max_seq_len=1536", - f"--max_batch_size={batch_size}", - "--paged_kv_cache=enable", - '--speculative_decoding_mode=medusa', - ] - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Run summarize...") - - summary_cmd = [ - f"{medusa_example_root}/../summarize.py", "--test_trt_llm", - "--hf_model_dir", f"{medusa_model_roots[0]}", "--tokenizer_dir", - f"{medusa_model_roots[0]}", f"--engine_dir={engine_dir}", - "--check_accuracy", "--tensorrt_llm_rouge1_threshold=24", - "--medusa_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]", - f"--temperature=1.0", f"--max_ite=40", f"--batch_size={batch_size}", - f"--dataset_dir={llm_datasets_root}", f"--rouge_dir={llm_rouge_root}" - ] - - if use_py_session: - summary_cmd.append("--use_py_session") - - venv_check_call(llm_venv, summary_cmd) - - -@skip_post_blackwell -@pytest.mark.parametrize("batch_size", [1, 8], ids=['bs1', 'bs8']) -@pytest.mark.parametrize("data_type", ['bfloat16', 'float16']) -@pytest.mark.parametrize("num_medusa_heads", [4], ids=['4-heads']) -@pytest.mark.parametrize("medusa_model_roots", ["medusa-vicuna-7b-v1.3"], - indirect=True) -@pytest.mark.parametrize("use_py_session", [False, True], - ids=["use_cpp_session", "use_py_session"]) -@pytest.mark.parametrize("base_model_datatype", ['fp8']) -def test_llm_medusa_with_qaunt_base_model_1gpu( - batch_size, data_type, medusa_model_roots, medusa_example_root, - base_model_datatype, llm_datasets_root, llm_rouge_root, - num_medusa_heads, llm_venv, cmodel_dir, engine_dir, use_py_session): - - model_name = f"vicuna_meudsa_quant_base_mode_{base_model_datatype}" - quant_model_ckpt_output_path = os.path.join(cmodel_dir, model_name) - - print("Quant base model to FP8 and combine medusa head") - quant_cmd = [ - f"{medusa_example_root}/../quantization/quantize.py", - f"--model_dir={medusa_model_roots[0]}", f"--dtype={data_type}", - f"--qformat={base_model_datatype}", - f"--kv_cache_dtype={base_model_datatype}", - f"--output_dir={quant_model_ckpt_output_path}", "--calib_size=512", - f"--medusa_model_dir={medusa_model_roots[1]}", - f"--num_medusa_heads={num_medusa_heads}" - ] - - # https://nvbugs/4658787 - # WAR before medusa tests can work offline - env = {"HF_DATASETS_OFFLINE": "0"} - venv_check_call(llm_venv, quant_cmd, env=env) - - print("Build engines...") - - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={quant_model_ckpt_output_path}", - f"--output_dir={engine_dir}", - f"--gpt_attention_plugin={data_type}", - f"--gemm_plugin={data_type}", - f"--max_beam_width=1", - "--remove_input_padding=enable", - "--context_fmha=enable", - "--max_input_len=1024", - "--max_seq_len=1536", - f"--max_batch_size={batch_size}", - "--paged_kv_cache=enable", - '--speculative_decoding_mode=medusa', - ] - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Run summarize...") - - summary_cmd = [ - f"{medusa_example_root}/../summarize.py", "--test_trt_llm", - "--hf_model_dir", f"{medusa_model_roots[0]}", "--tokenizer_dir", - f"{medusa_model_roots[0]}", f"--engine_dir={engine_dir}", - "--check_accuracy", "--tensorrt_llm_rouge1_threshold=24", - "--medusa_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]", - f"--temperature=1.0", f"--max_ite=40", f"--batch_size={batch_size}", - f"--dataset_dir={llm_datasets_root}", f"--rouge_dir={llm_rouge_root}" - ] - - if use_py_session: - summary_cmd.append("--use_py_session") - - venv_check_call(llm_venv, summary_cmd) diff --git a/tests/integration/defs/examples/test_mistral.py b/tests/integration/defs/examples/test_mistral.py deleted file mode 100644 index b5d07fb015cb..000000000000 --- a/tests/integration/defs/examples/test_mistral.py +++ /dev/null @@ -1,219 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -"""Module test_mistral test mistral examples.""" -import multiprocessing - -import defs.ci_profiler -import psutil -import pytest -from defs.common import (convert_weights, test_llm_torch_multi_lora_support, - venv_check_call) -from defs.conftest import (get_device_count, get_sm_version, - skip_post_blackwell, skip_pre_ada) -from defs.trt_test_alternative import check_call - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -def get_optimal_jobs(): - cpu_count = multiprocessing.cpu_count() - available_memory = psutil.virtual_memory().available / (1024 * 1024 * 1024) - memory_per_job = 4 - memory_based_jobs = int(available_memory / memory_per_job) - system_load = psutil.getloadavg()[0] / cpu_count - if system_load > 0.7: - cpu_factor = 0.5 - else: - cpu_factor = 0.75 - cpu_based_jobs = max(1, int(cpu_count * cpu_factor)) - optimal_jobs = max(1, min(cpu_based_jobs, memory_based_jobs)) - return optimal_jobs - - -@skip_post_blackwell #nvbug 5298661 -@pytest.mark.parametrize( - "run_type", - ['inference', 'summarization_long', 'chunked_summarization_long']) -@pytest.mark.parametrize("max_attention_window", [4096], - ids=['max_attention_window_size_4096']) -@pytest.mark.parametrize("data_type", ['float16']) -@pytest.mark.parametrize("llm_mistral_model_root", ['mistral-7b-v0.1'], - indirect=True) -def test_llm_mistral_v1_1gpu(run_type, data_type, llama_example_root, - max_attention_window, llm_mistral_model_root, - llm_datasets_root, llm_rouge_root, llm_venv, - cmodel_dir, engine_dir): - - print("Build engines...") - if run_type == "summarization_long": - model_name = 'mistral-{}'.format(run_type) - model_dir = convert_weights(llm_venv=llm_venv, - example_root=llama_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=llm_mistral_model_root, - data_type=data_type) - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={model_dir}", - f"--output_dir={engine_dir}", - "--max_input_len", - "6400", - f"--max_batch_size={1}", - "--max_seq_len", - "6528", - f"--gpt_attention_plugin={data_type}", - f"--gemm_plugin={data_type}", - "--context_fmha=enable", - "--use_paged_context_fmha=disable", - ] - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Run long context summarize...") - # using shorter input length since A30 doesn't have enough device memory. - summary_cmd = [ - f"{llama_example_root}/summarize_long.py", - "--test_trt_llm", - "--test_hf", - "--hf_model_location", - f"{llm_mistral_model_root}", - "--data_type", - "fp16", - f"--engine_dir={engine_dir}", - f"--max_attention_window_size={max_attention_window}", - "--max_ite", - "3", - "--max_input_len", - "6400", - "--tensorrt_llm_rouge1_threshold", - "90", - "--check_accuracy", - ] - # https://nvbugs/4658787 - # WAR before summarize_long.py can work offline - env = {"HF_DATASETS_OFFLINE": "0"} - venv_check_call(llm_venv, summary_cmd, env=env) - - # multi block + sliding window attention tests. - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={model_dir}", - f"--output_dir={engine_dir}", - "--max_input_len", - "6400", - "--max_seq_len", - "6528", - f"--gpt_attention_plugin={data_type}", - f"--gemm_plugin={data_type}", - "--use_paged_context_fmha=disable", - ] - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Run long context summarize with multi_block_mode enabled...") - # using shorter input length since A30 doesn't have enough device memory. - summary_cmd = [ - f"{llama_example_root}/summarize_long.py", "--test_trt_llm", - "--test_hf", "--hf_model_location", f"{llm_mistral_model_root}", - "--data_type", "fp16", f"--engine_dir={engine_dir}", - f"--max_attention_window_size={max_attention_window}", "--max_ite", - "3", "--max_input_len", "6400", "--tensorrt_llm_rouge1_threshold", - "90", "--check_accuracy" - ] - venv_check_call(llm_venv, summary_cmd, env=env) - - elif run_type == "chunked_summarization_long": - model_name = 'mistral-{}'.format(run_type) - model_dir = convert_weights(llm_venv=llm_venv, - example_root=llama_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=llm_mistral_model_root, - data_type=data_type) - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={model_dir}", - f"--output_dir={engine_dir}", - "--max_input_len", - "6400", - "--max_num_tokens=2048", - "--use_paged_context_fmha=enable", - f"--max_batch_size={1}", - "--max_seq_len", - "6528", - f"--gpt_attention_plugin={data_type}", - f"--gemm_plugin={data_type}", - "--context_fmha=enable", - ] - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Run long context summarize...") - summary_cmd = [ - f"{llama_example_root}/../../../summarize.py", - "--eval_task=summarize_long", "--test_trt_llm", "--test_hf", - "--hf_model_dir", f"{llm_mistral_model_root}", "--data_type", - "fp16", f"--engine_dir={engine_dir}", - f"--max_attention_window_size={max_attention_window}", - "--max_input_length", "6400", "--tensorrt_llm_rouge1_threshold", - "21", "--check_accuracy", "--enable_chunked_context" - ] - # https://nvbugs/4658787 - # WAR before summarize_long.py can work offline - env = {"HF_DATASETS_OFFLINE": "0"} - venv_check_call(llm_venv, summary_cmd, env=env) - - -@skip_pre_ada -@pytest.mark.skip_less_device_memory(80000) -@pytest.mark.parametrize("llm_mistral_model_root", [ - 'mistral-7b-v0.1', - 'mistral-nemo-instruct-2407', -], - indirect=True) -def test_mistral_with_bf16_lora_torch(llama_example_root, llm_datasets_root, - qcache_dir_without_install_package, - llm_venv, engine_dir, - llm_mistral_model_root): - """Run Mistral models with multiple dummy LoRAs using LLM-API Torch backend.""" - - if "mistral-nemo-instruct-2407" in llm_mistral_model_root.lower(): - tensor_parallel_size = 2 - if get_device_count() < 2: - pytest.skip( - "Skipping: mistral-nemo-instruct-2407 model requires 2 GPUs") - else: - tensor_parallel_size = 1 - - print(f"Testing {llm_mistral_model_root} with LLM-API Torch backend...") - - defs.ci_profiler.start("test_llm_torch_multi_lora_support") - test_llm_torch_multi_lora_support( - hf_model_dir=llm_mistral_model_root, - llm_venv=llm_venv, - num_loras=2, - lora_rank=8, - target_hf_modules=["q_proj", "k_proj", "v_proj"], - zero_lora_weights=True, - tensor_parallel_size=tensor_parallel_size) - defs.ci_profiler.stop("test_llm_torch_multi_lora_support") - print( - f"test_llm_torch_multi_lora_support: {defs.ci_profiler.elapsed_time_in_sec('test_llm_torch_multi_lora_support')} sec" - ) diff --git a/tests/integration/defs/examples/test_mixtral.py b/tests/integration/defs/examples/test_mixtral.py deleted file mode 100644 index 56d762f45257..000000000000 --- a/tests/integration/defs/examples/test_mixtral.py +++ /dev/null @@ -1,119 +0,0 @@ -# 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. - -import csv -import os - -import pytest -from defs.common import convert_weights, venv_mpi_check_call -from defs.conftest import get_sm_version -from defs.trt_test_alternative import check_call - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -@pytest.mark.skip_less_device(4) -@pytest.mark.skip_less_device_memory(45000) -@pytest.mark.parametrize("llm_lora_model_root", ["chinese-mixtral-lora"], - indirect=True) -@pytest.mark.parametrize("llm_mixtral_model_root", ["Mixtral-8x7B-v0.1"], - indirect=True) -def test_llm_mixtral_moe_plugin_lora_4gpus( - llama_example_root, - llm_mixtral_model_root, - llm_venv, - cmodel_dir, - engine_dir, - llm_lora_model_root, -): - "run Mixtral MoE lora test on 4 gpu." - print("Build engines...") - dtype = 'float16' - model_name = os.path.basename(llm_mixtral_model_root) - ckpt_dir = convert_weights(llm_venv=llm_venv, - example_root=llama_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - tp_size=4, - pp_size=1, - model_path=llm_mixtral_model_root, - data_type=dtype) - - print("Build engines...") - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={ckpt_dir}", - f"--output_dir={engine_dir}", - "--lora_plugin=auto", - "--moe_plugin=auto", - f"--lora_dir={llm_lora_model_root}", - "--worker=4", - "--max_batch_size=8", - ] - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - ref_1 = [ - 1, 28705, 29242, 30731, 31182, 235, 158, 142, 234, 182, 152, 28924, - 29926, 28971, 29242, 28988 - ] - ref_2 = [ - 1, 315, 2016, 285, 4284, 526, 5680, 28723, 28705, 28740, 28723, 661 - ] - - input_text = "我爱吃蛋糕" - print("Run inference with lora id 0...") - run_cmd = [ - f"{llama_example_root}/../../../run.py", - "--max_output_len=5", - f"--input_text={input_text}", - "--lora_task_uids=0", - f"--tokenizer_dir={llm_lora_model_root}", - f"--engine_dir={engine_dir}", - f"--output_csv={llm_venv.get_working_directory()}/use_lora.csv", - "--use_py_session", - ] - venv_mpi_check_call(llm_venv, ["mpirun", "-n", "4", "--allow-run-as-root"], - run_cmd) - - with open(f"{llm_venv.get_working_directory()}/use_lora.csv") as f: - predict = csv.reader(f) - predict = next(predict) - predict = [int(p) for p in predict] - assert ref_1 == predict - - print("Run inference with lora id -1...") - input_text = "I love french quiche" - run_cmd = [ - f"{llama_example_root}/../../../run.py", - "--max_output_len=5", - f"--input_text={input_text}", - "--lora_task_uids=-1", - f"--tokenizer_dir={llm_lora_model_root}", - f"--engine_dir={engine_dir}", - f"--output_csv={llm_venv.get_working_directory()}/no_lora.csv", - "--use_py_session", - ] - venv_mpi_check_call(llm_venv, ["mpirun", "-n", "4", "--allow-run-as-root"], - run_cmd) - - with open(f"{llm_venv.get_working_directory()}/no_lora.csv") as f: - predict = csv.reader(f) - predict = next(predict) - predict = [int(p) for p in predict] - assert ref_2 == predict diff --git a/tests/integration/defs/examples/test_multimodal.py b/tests/integration/defs/examples/test_multimodal.py deleted file mode 100644 index 591205e9afd7..000000000000 --- a/tests/integration/defs/examples/test_multimodal.py +++ /dev/null @@ -1,726 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2023 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. - -import os - -import pytest -import torch -from defs.common import convert_weights, venv_check_call, venv_mpi_check_call -from defs.conftest import (get_device_memory, get_sm_version, - skip_post_blackwell, skip_pre_ada) -from defs.trt_test_alternative import check_call - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -@pytest.fixture(scope="module") -def multimodal_example_root(llm_root): - "Get multimodal example root" - example_root = os.path.join(llm_root, "examples", "models", "core", - "multimodal") - - return example_root - - -@pytest.fixture(scope="function") -def recover_transformers(llm_venv, llm_root): - "Recover transformers" - - yield - - llm_venv.run_cmd([ - "-m", "pip", "install", "-r", - os.path.join(llm_root, "requirements.txt") - ]) - - -def _call_run_cmd(llm_venv, llm_root, cmd, world_size): - if world_size == 1: - venv_check_call(llm_venv, cmd) - else: - venv_mpi_check_call( - llm_venv, ["mpirun", "-n", - str(world_size), "--allow-run-as-root"], cmd) - - -dataset_path_mapping = { - 'cnn_dailymail': 'cnn_dailymail', - 'scienceqa': 'derek-thomas___science_qa', -} - - -def _test_llm_multimodal_general(llm_venv, - llm_root, - llm_datasets_root, - cmodel_dir, - engine_dir, - batch_size, - data_type, - tp_size, - pp_size, - multimodal_example_root, - multimodal_model_root, - recover_transformers, - calibration_dataset=None, - qformat=None, - kv_cache_dtype=None, - cpp_e2e=False, - num_beams=1): - - # Empty the torch CUDA cache before each multimodal test to reduce risk of OOM errors. - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - world_size = tp_size * pp_size - print("Locate model checkpoints in test storage...") - tllm_model_name, model_ckpt_path = multimodal_model_root - - if "neva-22b" in tllm_model_name and get_device_memory() < 80000: - pytest.skip("GPU memory is insufficient.") - if "Mistral-Small" in tllm_model_name and get_device_memory() < 80000: - pytest.skip("GPU memory is insufficient.") - - print("Converting huggingface model into binary format...") - # ckpt from llm_models/ --> cmodels// - model_name = tllm_model_name - model_name = "pix2struct" if model_name == "deplot" else model_name - opt_example_root = multimodal_example_root + "/../models/contrib/opt" - enc_dec_example_root = multimodal_example_root + "/../enc_dec" - llama_example_root = multimodal_example_root + "/../llama" - cogvlm_example_root = multimodal_example_root + "/../cogvlm" - gpt_example_root = multimodal_example_root + "/../gpt" - nemotron_example_root = multimodal_example_root + "/../nemotron" - phi_example_root = multimodal_example_root + "/../phi" - mllama_example_root = multimodal_example_root + "/../mllama" - qwen_example_root = multimodal_example_root + "/../qwen" - internlm_example_root = multimodal_example_root + "/../internlm2" - - opt_model = "opt" in model_name - nougat_model = "nougat" in model_name - gpt_model = "fuyu" in model_name or "neva-22b" in model_name or "kosmos" in model_name - pix2struct_model = "pix2struct" in model_name - enc_dec_model = "t5" in model_name or nougat_model or pix2struct_model - llava_model = "llava" in model_name - llava_next_model = "llava-v1.6" in model_name - llava_next_vision_trtllm_engine_model = "vision-trtllm" in model_name and llava_next_model - llava_onevision_model = "llava-onevision" in model_name - llava_onevision_video_model = "video" in model_name and llava_onevision_model - vila_model = "VILA" in model_name - cogvlm_model = "cogvlm" in model_name - nemotron_model = "video-neva" in model_name - phi3_model = "phi-3" in model_name.lower() - phi4_model = "phi-4" in model_name.lower() - mllama_model = 'Llama-3.2' in model_name - qwen2_vl_model = 'Qwen2-VL' in model_name - internlm_model = 'internlm-xcomposer2' in model_name - mistral_model = 'Mistral-Small' in model_name - if enc_dec_model: - builder_root = enc_dec_example_root - if nougat_model: - model_type = "bart" - if pix2struct_model: - model_type = "pix2struct" - if "t5" in model_name: - model_type = "blip2" - elif gpt_model: - builder_root, model_type = gpt_example_root, "gpt" - elif llava_onevision_model: - builder_root, model_type = qwen_example_root, "qwen" - elif qwen2_vl_model: - builder_root, model_type = qwen_example_root, "qwen" - elif internlm_model: - builder_root, model_type = internlm_example_root, "internlm" - elif llava_model or vila_model: - builder_root, model_type = llama_example_root, "llama" - elif mistral_model: - builder_root, model_type = llama_example_root, "llama" - elif cogvlm_model: - builder_root, model_type = cogvlm_example_root, "cogvlm" - elif nemotron_model: - builder_root, model_type = nemotron_example_root, "nemotron" - elif phi3_model: - model_name = model_name.split('/')[-1] # Remove HF directory name - builder_root, model_type = phi_example_root, "phi-3-vision" - elif phi4_model: - builder_root, model_type = phi_example_root, "phi-4-multimodal" - elif opt_model: - builder_root, model_type = opt_example_root, "blip2" - elif mllama_model: - builder_root, model_type = mllama_example_root, "mllama" - - use_weight_only = (not enc_dec_model) and (data_type in [ - 'int4_weight_only', 'int8_weight_only' - ]) - weight_only_precision = data_type.split('_')[0] if use_weight_only else None - if use_weight_only: data_type = 'float16' - - if vila_model: - print( - "VILA model has dependencies on certain HuggingFace version. Need to pip install until this limitation is removed." - ) - check_call( - f"pip install -r {multimodal_example_root}/requirements-vila.txt", - shell=True, - env=llm_venv._new_env) - elif llava_onevision_model: - check_call( - f"pip install -r {multimodal_example_root}/requirements-llava_onevision.txt", - shell=True, - env=llm_venv._new_env) - elif qwen2_vl_model: - check_call( - f"pip install -r {multimodal_example_root}/requirements-qwen2vl.txt", - shell=True, - env=llm_venv._new_env) - elif internlm_model: - check_call( - f"pip install -r {multimodal_example_root}/requirements-internlm-xcomposer2.txt", - shell=True, - env=llm_venv._new_env) - elif mllama_model: - check_call(f"pip install -r {mllama_example_root}/requirements.txt", - shell=True, - env=llm_venv._new_env) - if qformat == 'fp8': - convert_cmd = [ - f"{multimodal_example_root}/../../../quantization/quantize.py", - f"--model_dir={model_ckpt_path}", - f"--calib_dataset={llm_datasets_root}/{dataset_path_mapping[calibration_dataset]}", - f"--dtype={data_type}", - f"--qformat={qformat}", - f"--kv_cache_dtype={kv_cache_dtype}", - f"--output_dir={cmodel_dir}", - f"--calib_size=16", - ] - venv_check_call(llm_venv, convert_cmd) - converted_weight_dir = cmodel_dir - else: - converted_weight_dir = convert_weights( - llm_venv, - builder_root, - cmodel_dir, - model_name, - model_ckpt_path, - data_type=data_type, - gpus=tp_size, - model_type=model_type, - use_weight_only=use_weight_only, - weight_only_precision=weight_only_precision, - tp_size=tp_size, - pp_size=pp_size, - batch_size=batch_size, - multimodal=True) - - print("Build LLM engines...") - model_name = model_name.split('/')[-1] # Remove HF directory name - llm_engine_dir = f"{engine_dir}/{model_name}/{world_size}-gpu" - if "opt" in model_name or llava_model or vila_model or gpt_model or nemotron_model or phi3_model or phi4_model or qwen2_vl_model or mistral_model: - max_input_len_text = 1024 - max_output_len = 200 - if llava_next_model: - multimodal_len = 4096 - elif llava_onevision_model: - multimodal_len = 7300 - elif llava_model: - multimodal_len = 576 - elif vila_model: - multimodal_len = 196 - elif phi3_model: - multimodal_len = 5120 - elif phi4_model: - multimodal_len = 5120 - elif mistral_model: - multimodal_len = 5120 - elif "fuyu" in model_name: - multimodal_len = 2640 - elif "neva-22b" in model_name: - multimodal_len = 729 - elif "video-neva" in model_name: - multimodal_len = 3072 - elif "kosmos" in model_name: - multimodal_len = 64 - elif "Qwen2-VL" in model_name: - multimodal_len = 3552 - else: - multimodal_len = 32 - max_input_len = max_input_len_text + batch_size * multimodal_len - max_seq_len = max_input_len + max_output_len - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={converted_weight_dir}", - f"--output_dir={llm_engine_dir}/llm", - f"--gpt_attention_plugin {data_type}", - f"--gemm_plugin={data_type}", - f"--max_batch_size={batch_size}", - f"--max_multimodal_len={batch_size * multimodal_len}", - f"--max_input_len={max_input_len}", - f"--max_seq_len={max_seq_len}", - f"--max_num_tokens={max_input_len}", - f"--max_beam_width={num_beams}", - ] - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - elif internlm_model: - max_input_len_text = 1536 - max_output_len = 200 - multimodal_len = 1225 - max_input_len = max_input_len_text + batch_size * multimodal_len - max_seq_len = max_input_len + max_output_len - - max_lora_rank = 256 - lora_dir = "." - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={converted_weight_dir}", - f"--output_dir={llm_engine_dir}", - f"--gpt_attention_plugin {data_type}", - f"--gemm_plugin={data_type}", - f"--lora_plugin={data_type}", - f"--lora_dir={lora_dir}", - f"--max_lora_rank={max_lora_rank}", - f"--max_batch_size={batch_size}", - f"--max_multimodal_len={batch_size * multimodal_len}", - f"--max_input_len={max_input_len}", - f"--max_seq_len={max_seq_len}", - f"--max_num_tokens={max_input_len}", - f"--max_beam_width={num_beams}", - ] - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - elif enc_dec_model: - components = ["decoder"] if nougat_model or pix2struct_model else [ - "encoder", "decoder" - ] - for component in components: - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={converted_weight_dir}/{component}", - f"--output_dir={llm_engine_dir}/{data_type}/llm/{component}", - "--paged_kv_cache=enable", - "--moe_plugin=disable", - f"--max_batch_size={batch_size}", - "--max_seq_len=412", - f"--gemm_plugin={data_type}", - f"--bert_attention_plugin={data_type}", - f"--gpt_attention_plugin={data_type}", - "--remove_input_padding=enable", - f"--max_beam_width={num_beams}", - ] - - # for non-T5 models, FP16/BF16 - if model_type == "t5" or data_type == "float32": - build_cmd.append("--context_fmha=disable") - if "t5" in model_name: - if component == "encoder": - build_cmd.append(f"--max_multimodal_len={32 * batch_size}") - build_cmd.append("--max_input_len=412") - else: - build_cmd.append("--max_encoder_input_len=412") - build_cmd.append(f"--max_input_len=1") - else: # Nougat - assert nougat_model or pix2struct_model - if component == "encoder": - build_cmd.append(f"--max_multimodal_len={588 * batch_size}") - - # only decoder for nougat - if nougat_model: - build_cmd.append( - f"--max_encoder_input_len={588 * batch_size}") - else: - build_cmd.append( - f"--max_encoder_input_len={2048 * batch_size}") - - build_cmd.append(f"--max_input_len=1") - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - elif cogvlm_model: - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={converted_weight_dir}", - f"--output_dir={llm_engine_dir}/llm", - f"--gemm_plugin={data_type}", - f"--gpt_attention_plugin={data_type}", - f"--remove_input_padding=enable", - f"--max_batch_size={batch_size}", - f"--max_input_len=2048", - f"--max_seq_len=2048", - f"--paged_kv_cache=enable", - f"--bert_attention_plugin=disable", - f"--moe_plugin=disable", - f"--max_multimodal_len=61440", - f"--max_beam_width={num_beams}", - ] - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - elif mllama_model: - # set max_encoder_input_len = 6404 for running both non-instruct model and instruct model - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={converted_weight_dir}", - f"--output_dir={llm_engine_dir}/llm", - f"--gemm_plugin={data_type}", - f"--max_num_tokens=4096", - f"--max_seq_len=2048", - f"--max_batch_size={batch_size}", - f"--max_encoder_input_len=6404", - f"--max_beam_width={num_beams}", - ] - if kv_cache_dtype == 'fp8': - build_cmd.extend([ - "--use_fp8_context_fmha=enable", - "--use_paged_context_fmha=enable", - ]) - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Build visual engines...") - vision_model_type = model_name - if 'llava' in model_name: vision_model_type = 'llava' - if 'llava-v1.6' in model_name: vision_model_type = 'llava_next' - elif llava_onevision_model: vision_model_type = 'llava_onevision' - elif 'VILA' in model_name: vision_model_type = 'vila' - elif nougat_model: vision_model_type = 'nougat' - elif pix2struct_model: vision_model_type = 'pix2struct' - elif 'cogvlm' in model_name: vision_model_type = 'cogvlm' - elif 'fuyu' in model_name: vision_model_type = 'fuyu' - elif 'neva-22b' in model_name: vision_model_type = 'neva' - elif 'video-neva' in model_name: vision_model_type = 'video-neva' - elif phi3_model: vision_model_type = "phi-3-vision" - elif phi4_model: vision_model_type = "phi-4-multimodal" - elif 'blip2' in model_name: vision_model_type = 'blip2' - elif 'Llama-3.2' in model_name: vision_model_type = 'mllama' - elif "Qwen2-VL" in model_name: vision_model_type = 'qwen2_vl' - elif 'internlm' in model_name: vision_model_type = 'internlm-xcomposer2' - elif 'Mistral-Small' in model_name: vision_model_type = 'pixtral' - - vit_batch_size = batch_size - if vision_model_type == "llava_next": - vit_batch_size = vit_batch_size * 5 - elif vision_model_type == 'llava_onevision': - vit_batch_size = vit_batch_size * 32 - - llm_engine_subdir = f"{data_type}" if enc_dec_model else "" - # Phi4MM has both vision and audio. Engine build dumps to vision and audio dirs automatically by builder. - component_dir = "vision" if vision_model_type != "phi-4-multimodal" else "" - build_cmd = [ - f"{multimodal_example_root}/build_multimodal_engine.py", - f"--output_dir={os.path.join(llm_engine_dir, llm_engine_subdir, component_dir)}", - f"--model_type={vision_model_type}", - f"--model_path={model_ckpt_path}", - f"--max_batch_size={vit_batch_size}", - ] - if vision_model_type == "vila": - vila_path = model_ckpt_path + "/../VILA" - build_cmd.extend([f"--vila_path={vila_path}"]) - if llava_next_vision_trtllm_engine_model: - script_root = f"{multimodal_example_root}/../vit" - convert_cmd = [ - f"{script_root}/convert_checkpoint.py", - f"--model_dir={model_ckpt_path}", - f"--output_dir={os.path.join(cmodel_dir, model_name, data_type, component_dir)}", - f"--dtype={data_type}", - f"--vision_tp_size={tp_size}", - ] - venv_check_call(llm_venv, convert_cmd) - - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={os.path.join(cmodel_dir, model_name, data_type, component_dir)}", - f"--output_dir={os.path.join(llm_engine_dir, llm_engine_subdir, component_dir)}", - f"--max_batch_size={vit_batch_size}", - f"--remove_input_padding disable", - f"--bert_attention_plugin disable", - ] - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - else: - venv_check_call(llm_venv, build_cmd) - - if llava_next_vision_trtllm_engine_model: - cp_cmd = [ - "cp", - f"{os.path.join(cmodel_dir, model_name, data_type, 'vision', 'image_newlines.safetensors')}", - f"{os.path.join(llm_engine_dir, llm_engine_subdir, 'vision')}", - ] - check_call(" ".join(cp_cmd), shell=True, env=llm_venv._new_env) - - print("Run inference...") - hf_model_dir = model_ckpt_path + "/../vicuna-7b-v1.5" if cogvlm_model else model_ckpt_path - hf_model_dir = converted_weight_dir if "neva" in model_name else hf_model_dir - video_path = os.path.join( - os.path.dirname(model_ckpt_path), "test_video", - "video_test.mp4") if "video-neva" in model_name else "" - run_cmd = [ - f"{multimodal_example_root}/run.py", - f"--engine_dir={llm_engine_dir}/{llm_engine_subdir}", - f"--hf_model_dir={hf_model_dir}", "--max_new_tokens=30", - f"--batch_size={batch_size}", "--check_accuracy", - "--enable_context_fmha_fp32_acc" - ] - if vision_model_type == 'phi-4-multimodal': - audio_path = f"{model_ckpt_path}/examples/what_is_shown_in_this_image.wav" - run_cmd.extend(["--audio_path", f"{audio_path}"]) - if vision_model_type in ['llava', 'vila'] and batch_size > 1: - # batch inference test - if vision_model_type == 'vila': - input_text = [ - '"\n Please elaborate what you see in the images?"' - ] * batch_size - else: - input_text = ['"\\n Which city is this? Answer:"'] * batch_size - run_cmd.append("--input_text") - run_cmd.extend(input_text) - if enc_dec_model: - run_cmd.extend(["--cross_kv_cache_fraction", "0.5"]) - if vision_model_type == "neva" and not cpp_e2e: - # randomly pick one to test the python runtime - run_cmd.extend(["--session", "python"]) - if vision_model_type == "video-neva": - run_cmd.extend(["--video_path", video_path]) - if llava_onevision_video_model: - run_cmd.extend(["--video_path", 'llava-onevision-accuracy']) - if phi3_model or phi4_model: - run_cmd.extend(["--kv_cache_free_gpu_memory_fraction", "0.4"]) - if cpp_e2e: - run_cmd.extend(["--session", "cpp"]) - if num_beams > 1: - run_cmd.extend(["--num_beams", str(num_beams)]) - - if mllama_model: - if qformat is None: - run_cmd_vision = run_cmd.copy() - run_cmd_vision.extend([ - "--cross_kv_cache_fraction=0.5", # mllama uses cross attention - "--image_path", - "https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg", - "--input_text", - "If I had to write a haiku for this one" - ]) - - print("Run mllama vision test in with example image ...") - _call_run_cmd(llm_venv, llm_root, run_cmd_vision, world_size) - - print("multimodal_example_root: ", multimodal_example_root) - print("llm_root: ", llm_root) - run_cmd_vision = run_cmd.copy() - run_cmd_vision.extend([ - "--cross_kv_cache_fraction=0.5", # mllama uses cross attention - "--image_path", - os.path.join( - llm_root, - "tests/integration/test_input_files/excel_table_test.jpg"), - "--input_text", - "What is the total income? Answer:" - ]) - - print("Run mllama vision test with random image ...") - - run_cmd_text = run_cmd.copy() - run_cmd_text.extend([ - "--cross_kv_cache_fraction=0.5", # mllama uses cross attention - "--input_text", - "The key to life is", - ]) - print("Run mllama text test...") - _call_run_cmd(llm_venv, llm_root, run_cmd_text, world_size) - else: - _call_run_cmd(llm_venv, llm_root, run_cmd, world_size) - - # Run evaluation test - if batch_size == 1 and (data_type == "float16" or qformat == 'fp8'): - print(f"prepare to run eval test") - - # for blip2-t5, ref: https://github.com/huggingface/transformers/issues/25491 - if "t5" in model_name: - check_call("pip uninstall -y apex", - shell=True, - env=llm_venv._new_env) - - # Threshold are set based on the HF correctness for 20 iterations - threshold_map = { - 'blip2-opt-2.7b': 35, - 'blip2-flan-t5-xl': 55, - 'llava-1.5-7b-hf': 65, - 'llava-v1.6-mistral-7b-hf': 65, - 'llava-onevision-qwen2-7b-ov-hf': 80, - 'VILA1.5-3b': 75, # from local TRT-LLM run - 'fuyu-8b': 70, - 'kosmos-2': 60, - 'Phi-3-vision-128k-instruct': 75, - 'Phi-3.5-vision-instruct': 85, - 'Llama-3.2-11B-Vision': 60, # The expected score is 62 - 'Llama-3.2-11B-Vision-Instruct': 75, # The expected score is 77 - 'Qwen2-VL-7B-Instruct': 80, - } - - if model_name not in threshold_map: - print(f"Skip {model_name} evaluation test.") - return - - # TODO: Delete these lines after resolving the issues - # For llava - input tokens are not parsed correctly with '\n' - # For llava_next - correctness lower than HF, and needs lower transformer version built - # For Phi-3 - correctness lower than HF - # For qwen_vl - runtime issue with eval.py -- need to unify prompt generation logics - # For internvl - not added to the test - if llava_model or llava_next_model or phi3_model or qwen2_vl_model: - return - - eval_task = "lmms-lab/ai2d" if mllama_model else "lmms-lab/VQAv2" - - eval_cmd = [ - f"{multimodal_example_root}/eval.py", - f"--model_type={vision_model_type}", - f"--engine_dir={llm_engine_dir}/{llm_engine_subdir}", - f"--hf_model_dir={hf_model_dir}", "--enable_context_fmha_fp32_acc", - "--test_trtllm", - f"--accuracy_threshold={threshold_map[model_name]}", - f"--eval_task={eval_task}" - ] - - if mllama_model: - eval_cmd.extend([ - f"--dataset_dir={llm_datasets_root}/lmms-lab___ai2d/", - "--cross_kv_cache_fraction=0.5", "--max_ite=100" - ]) - else: - eval_cmd.extend([ - f"--dataset_dir={llm_datasets_root}/lmms-lab__VQAv2_valid_2000samples/" - ]) - - if phi3_model: - eval_cmd.extend(["--kv_cache_free_gpu_memory_fraction", "0.4"]) - elif enc_dec_model: - eval_cmd.extend(["--cross_kv_cache_fraction", "0.5"]) - - print(f"Run {model_name} evaluation test...") - _call_run_cmd(llm_venv, llm_root, eval_cmd, world_size) - - -@pytest.mark.parametrize("num_beams", [1, 4], - ids=lambda num_beams: f'nb:{num_beams}') -@pytest.mark.parametrize('cpp_e2e', [False, True], - ids=lambda cpp_e2e: f'cpp_e2e:{cpp_e2e}') -@pytest.mark.parametrize("batch_size", [1, 8], - ids=lambda batch_size: f'bs:{batch_size}') -@pytest.mark.parametrize( - "data_type", - ['float16', 'bfloat16', 'int4_weight_only', 'int8_weight_only']) -@pytest.mark.parametrize("tp_size", [1, 2], ids=lambda tp_size: f'tp:{tp_size}') -@pytest.mark.parametrize("pp_size", [1, 2], ids=lambda pp_size: f'pp:{pp_size}') -@pytest.mark.parametrize("multimodal_model_root", [ - 'blip2-opt-2.7b', - 'blip2-flan-t5-xl', - 'llava-1.5-7b-hf', - 'llava-v1.6-mistral-7b-hf', - pytest.param('llava-v1.6-mistral-7b-hf-vision-trtllm', - marks=pytest.mark.skipif(get_device_memory() < 50000, - reason="Skip due to low memory")), - 'llava-onevision-qwen2-7b-ov-hf', - 'llava-onevision-qwen2-7b-ov-hf-video', - pytest.param('nougat-base', marks=skip_post_blackwell), - 'VILA1.5-3b', - 'cogvlm-chat', - 'fuyu-8b', - pytest.param('deplot', marks=skip_post_blackwell), - pytest.param('neva-22b', - marks=pytest.mark.skip(reason="RCCA https://nvbugs/5220761")), - 'kosmos-2', - pytest.param('video-neva', marks=skip_post_blackwell), - pytest.param('Phi-3-vision-128k-instruct', marks=skip_post_blackwell), - pytest.param('Phi-3.5-vision-instruct', marks=skip_post_blackwell), - pytest.param('Phi-4-multimodal-instruct', marks=skip_post_blackwell), - pytest.param('Llama-3.2-11B-Vision', marks=skip_post_blackwell), - 'Qwen2-VL-7B-Instruct', - 'internlm-xcomposer2-vl-7b', - 'Mistral-Small-3.1-24B-Instruct-2503', -], - indirect=True) -def test_llm_multimodal_general(llm_venv, llm_root, llm_datasets_root, - cmodel_dir, engine_dir, batch_size, data_type, - tp_size, pp_size, multimodal_example_root, - multimodal_model_root, recover_transformers, - cpp_e2e, num_beams): - _test_llm_multimodal_general(llm_venv, - llm_root, - llm_datasets_root, - cmodel_dir, - engine_dir, - batch_size, - data_type, - tp_size, - pp_size, - multimodal_example_root, - multimodal_model_root, - recover_transformers, - cpp_e2e=cpp_e2e, - num_beams=num_beams) - - -@skip_pre_ada -@pytest.mark.parametrize('cpp_e2e', [False, True], - ids=lambda cpp_e2e: f'cpp_e2e:{cpp_e2e}') -@pytest.mark.parametrize("batch_size", [1, 8], - ids=lambda batch_size: f'bs:{batch_size}') -@pytest.mark.parametrize("data_type", ['float16', 'bfloat16']) -@pytest.mark.parametrize("tp_size", [1, 2], ids=lambda tp_size: f'tp:{tp_size}') -@pytest.mark.parametrize("pp_size", [1, 2], ids=lambda pp_size: f'pp:{pp_size}') -@pytest.mark.parametrize("multimodal_model_root", [ - 'blip2-opt-2.7b', - 'blip2-flan-t5-xl', - 'llava-1.5-7b-hf', - 'llava-v1.6-mistral-7b-hf', - 'llava-onevision-qwen2-7b-ov-hf', - 'llava-onevision-qwen2-7b-ov-hf-video', - 'nougat-base', - 'VILA1.5-3b', - 'cogvlm-chat', - 'fuyu-8b', - 'deplot', - 'neva-22b', - 'kosmos-2', - 'video-neva', - 'Phi-3-vision-128k-instruct', - 'Phi-3.5-vision-instruct', - 'Phi-4-multimodal-instruct', - pytest.param('Llama-3.2-11B-Vision-Instruct', marks=skip_post_blackwell), - pytest.param('Llama-3.2-11B-Vision', marks=skip_post_blackwell), - 'Qwen2-VL-7B-Instruct', -], - indirect=True) -@pytest.mark.parametrize('calibration_dataset', ['scienceqa', 'cnn_dailymail']) -@pytest.mark.parametrize('qformat', ['fp8']) -@pytest.mark.parametrize('kv_cache_dtype', ['fp8']) -def test_llm_fp8_multimodal_general( - llm_venv, llm_root, llm_datasets_root, cmodel_dir, engine_dir, - batch_size, data_type, tp_size, pp_size, multimodal_example_root, - multimodal_model_root, recover_transformers, calibration_dataset, - qformat, kv_cache_dtype, cpp_e2e): - _test_llm_multimodal_general(llm_venv, - llm_root, - llm_datasets_root, - cmodel_dir, - engine_dir, - batch_size, - data_type, - tp_size, - pp_size, - multimodal_example_root, - multimodal_model_root, - recover_transformers, - calibration_dataset, - qformat, - kv_cache_dtype, - cpp_e2e=cpp_e2e) diff --git a/tests/integration/defs/examples/test_nemotron.py b/tests/integration/defs/examples/test_nemotron.py deleted file mode 100644 index 0f7a69134bd4..000000000000 --- a/tests/integration/defs/examples/test_nemotron.py +++ /dev/null @@ -1,72 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. - -import pytest -from defs.common import venv_check_call -from defs.conftest import get_sm_version -from defs.trt_test_alternative import check_call - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -@pytest.mark.skip_less_device_memory(50000) -@pytest.mark.parametrize("qformat", ["full_prec", "fp8", "int4_awq"]) -@pytest.mark.parametrize("dtype", ["float16", "bfloat16"]) -def test_llm_nemotron_3_8b_1gpu(nemotron_example_root, - llm_nemotron_3_8b_model_root, llm_datasets_root, - llm_rouge_root, llm_venv, cmodel_dir, - engine_dir, dtype, qformat): - print("Converting checkpoint...") - model_name = 'nemotron-3-8b' - ckpt_dir = f"{cmodel_dir}/{model_name}/{qformat}/1-gpu" - - quantize_cmd = [ - f"{nemotron_example_root}/../quantization/quantize.py", - f"--nemo_ckpt_path={llm_nemotron_3_8b_model_root}", - f"--calib_dataset={llm_datasets_root}/cnn_dailymail", - "--batch_size=64", - f"--dtype={dtype}", - f"--qformat={qformat}", - f"--output_dir={ckpt_dir}", - ] - venv_check_call(llm_venv, quantize_cmd) - - print("Building engines...") - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={ckpt_dir}", - f"--output_dir={engine_dir}", - "--max_batch_size=8", - "--max_input_len=924", - "--max_seq_len=1024", - f"--gpt_attention_plugin={dtype}", - f"--gemm_plugin={dtype}", - ] - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Run engines...") - summary_cmd = [ - f"{nemotron_example_root}/../summarize.py", "--test_trt_llm", - f"--engine_dir={engine_dir}", - f"--vocab_file={ckpt_dir}/tokenizer.model", "--no_add_special_tokens", - "--batch_size=8", "--max_ite=40", "--check_accuracy", - "--tensorrt_llm_rouge1_threshold=18", - f"--dataset_dir={llm_datasets_root}", f"--rouge_dir={llm_rouge_root}" - ] - venv_check_call(llm_venv, summary_cmd) diff --git a/tests/integration/defs/examples/test_nemotron_nas.py b/tests/integration/defs/examples/test_nemotron_nas.py deleted file mode 100644 index d1663eab672e..000000000000 --- a/tests/integration/defs/examples/test_nemotron_nas.py +++ /dev/null @@ -1,124 +0,0 @@ -from pathlib import Path - -import pytest -from defs.common import convert_weights, venv_check_call, venv_mpi_check_call -from defs.conftest import get_device_memory, get_sm_version -from defs.trt_test_alternative import check_call - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - -ROUGE1_ACCURACY_THRESHOLD = 20 - - -@pytest.mark.parametrize("nemotron_nas_model_root", [ - "DeciLM-7B", -], - indirect=True) -def test_nemotron_nas_summary_1gpu(nemotron_nas_example_root, llm_venv, - nemotron_nas_model_root, llm_datasets_root, - llm_rouge_root, engine_dir, cmodel_dir): - model_name = Path(nemotron_nas_model_root).name - if "51B" in model_name and get_device_memory() < 80000: - pytest.skip("device memory is insufficient.") - - print(f"Model name: {model_name}") - dtype = 'float16' - ckpt_type = "hf" - - print("Converting checkpoint...") - model_dir = convert_weights(llm_venv=llm_venv, - example_root=nemotron_nas_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=nemotron_nas_model_root, - data_type=dtype, - ckpt_type=ckpt_type, - gpus=1, - tp_size=1, - trust_remote_code=True) - - print("Building engines...") - build_cmd = [ - "trtllm-build", f"--checkpoint_dir={model_dir}", - f"--output_dir={engine_dir}", f"--max_batch_size={4}", - f"--max_input_len={2048}", "--kv_cache_type=paged", - "--remove_input_padding=enable", "--gemm_plugin=auto", - "--gpt_attention_plugin=auto" - ] - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Running inference...") - - summary_cmd = [ - f"{nemotron_nas_example_root}/../../../summarize.py", - f"--engine_dir={engine_dir}", "--test_hf", "--hf_device_map_auto", - "--batch_size=1", "--test_trt_llm", - f"--hf_model_dir={nemotron_nas_model_root}", "--check_accuracy", - f"--tensorrt_llm_rouge1_threshold={ROUGE1_ACCURACY_THRESHOLD}", - "--no_add_special_tokens", f"--dataset_dir={llm_datasets_root}", - f"--rouge_dir={llm_rouge_root}" - ] - - venv_check_call(llm_venv, summary_cmd) - - -@pytest.mark.skip_less_device(2) -@pytest.mark.parametrize("nemotron_nas_model_root", [ - "DeciLM-7B", - "Llama-3_1-Nemotron-51B-Instruct", -], - indirect=True) -def test_nemotron_nas_summary_2gpu(nemotron_nas_example_root, llm_venv, - nemotron_nas_model_root, llm_datasets_root, - llm_rouge_root, engine_dir, cmodel_dir): - model_name = Path(nemotron_nas_model_root).name - if "51B" in model_name and get_device_memory() < 80000: - pytest.skip("device memory is insufficient.") - - print(f"Model name: {model_name}") - dtype = 'float16' - ckpt_type = "hf" - - print("Converting checkpoint...") - model_dir = convert_weights(llm_venv=llm_venv, - example_root=nemotron_nas_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=nemotron_nas_model_root, - data_type=dtype, - ckpt_type=ckpt_type, - gpus=2, - tp_size=2, - trust_remote_code=True) - - print("Building engines...") - build_cmd = [ - "trtllm-build", f"--checkpoint_dir={model_dir}", - f"--output_dir={engine_dir}", f"--max_batch_size={4}", - f"--max_input_len={2048}", "--kv_cache_type=paged", - "--remove_input_padding=enable", "--gemm_plugin=auto", - "--gpt_attention_plugin=auto" - ] - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Running inference...") - - mpi_cmd = ["mpirun", "-n", "2", "--allow-run-as-root"] - - summary_cmd = [ - f"{nemotron_nas_example_root}/../../../summarize.py", - f"--engine_dir={engine_dir}", "--test_hf", "--hf_device_map_auto", - "--batch_size=1", "--test_trt_llm", - f"--hf_model_dir={nemotron_nas_model_root}", "--check_accuracy", - f"--tensorrt_llm_rouge1_threshold={ROUGE1_ACCURACY_THRESHOLD}", - "--no_add_special_tokens", f"--dataset_dir={llm_datasets_root}", - f"--rouge_dir={llm_rouge_root}" - ] - - venv_mpi_check_call(llm_venv, mpi_cmd, summary_cmd) diff --git a/tests/integration/defs/examples/test_ngram.py b/tests/integration/defs/examples/test_ngram.py deleted file mode 100644 index 2de49e8322f0..000000000000 --- a/tests/integration/defs/examples/test_ngram.py +++ /dev/null @@ -1,163 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. - -import csv -from copy import deepcopy - -import pytest -from defs.common import convert_weights, venv_check_call -from defs.conftest import get_sm_version, skip_post_blackwell -from defs.trt_test_alternative import check_call - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -# TODO: remove skip after support NGram on B200 -@skip_post_blackwell -@pytest.mark.parametrize("batch_size", [1, 2], ids=['bs1', 'bs2']) -@pytest.mark.parametrize("data_type", ['float16']) -@pytest.mark.parametrize("max_draft_len", [4, 8], - ids=['max_draft_len_4', 'max_draft_len_8']) -@pytest.mark.parametrize( - "max_matching_ngram_size", [2, 4], - ids=['max_matching_ngram_size_2', 'max_matching_ngram_size_4']) -@pytest.mark.parametrize("use_logits", [False, True], - ids=['use_tokens', 'use_logits']) # useless yet -@pytest.mark.parametrize("use_py_session", [False], ids=["use_cpp_session"]) -@pytest.mark.parametrize("ngram_root", ["gpt2"], indirect=True) -@pytest.mark.parametrize("streaming", [False, True], - ids=["no_streaming", "streaming"]) -def test_llm_ngram_1gpu(batch_size, data_type, max_draft_len, - max_matching_ngram_size, use_logits, use_py_session, - ngram_root, streaming, ngram_example_root, - llm_datasets_root, llm_rouge_root, llm_venv, cmodel_dir, - engine_dir): - model_name = "ngram" - - print("Build checkpoint ...") - model_dir = convert_weights(llm_venv=llm_venv, - example_root=ngram_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=ngram_root, - data_type=data_type) - - print("Build engines ...") - target_engine_dir = engine_dir + "-target" - baseline_engine_dir = engine_dir + "-baseline" - common_build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={model_dir}", - f"--max_batch_size={batch_size}", - f"--max_beam_width=1", - "--max_input_len=1024", - "--max_seq_len=1536", - "--use_paged_context_fmha=enable", - f"--gpt_attention_plugin={data_type}", - f"--gemm_plugin={data_type}", - ] - target_model_build_cmd = deepcopy(common_build_cmd) - target_model_build_cmd.extend([ - f"--output_dir={target_engine_dir}", - "--speculative_decoding_mode=draft_tokens_external", - f"--max_draft_len={max_draft_len+1}", - ]) - baseline_model_build_cmd = deepcopy(common_build_cmd) - baseline_model_build_cmd.extend([ - f"--output_dir={baseline_engine_dir}", - ]) - - check_call(" ".join(target_model_build_cmd), - shell=True, - env=llm_venv._new_env) - check_call(" ".join(baseline_model_build_cmd), - shell=True, - env=llm_venv._new_env) - - print("Run inferences ...") - common_run_cmd = [ - f"{ngram_example_root}/../run.py", - f"--tokenizer_dir={ngram_root}", - f"--max_output_len=64", - f"--kv_cache_enable_block_reuse", - f"--kv_cache_free_gpu_memory_fraction=0.25", - ] - if streaming: - common_run_cmd.extend(["--streaming", "--streaming_interval=1"]) - if batch_size == 1: - common_run_cmd.extend(["--input_text", "'How are you?'"]) - elif batch_size == 2: - common_run_cmd.extend(["--input_text", "'Hello'", "'How are you?'"]) - else: - assert False, "Only batch_size <=2 is supported in test." - assert not use_py_session, "Only CPP session is supported in Draft-Target-Model." - - run_cmd = deepcopy(common_run_cmd) - ngram_config = f"[{max_draft_len},{max_matching_ngram_size},[0]]" - run_cmd.extend([ - f"--engine_dir={target_engine_dir}", - f"--ngram_config={ngram_config}", - f"--output_csv={engine_dir}/ngram_output.csv", - ]) - baseline_run_cmd = deepcopy(common_run_cmd) - baseline_run_cmd.extend([ - f"--engine_dir={baseline_engine_dir}", - f"--output_csv={engine_dir}/baseline_output.csv", - ]) - - venv_check_call(llm_venv, run_cmd) - venv_check_call(llm_venv, baseline_run_cmd) - - print("Compare outputs ...") - with open(f"{engine_dir}/ngram_output.csv") as dt_f, open( - f"{engine_dir}/baseline_output.csv") as b_f: - for bs, (dt_request, - b_request) in enumerate(zip(csv.reader(dt_f), - csv.reader(b_f))): - assert ( - len(dt_request) == len(b_request) - ), f"Output length at ({bs=}) is different ({len(dt_request)} v.s. {len(b_request)})." - for index, (dt, b) in enumerate(zip(dt_request, b_request)): - assert ( - int(dt) == int(b) - ), f"Output at ({bs=}, {index=}) is different ({dt} v.s. {b})." - - if batch_size > 1 or streaming: # Summarize tests for only batch_size=1 and streaming=False. - return - - print("Run summarize...") - ngram_config = f"[{max_draft_len},{max_matching_ngram_size},[0]]" - - run_cmd = [ - f"{ngram_example_root}/../summarize.py", - "--test_hf", - "--test_trt_llm", - "--check_accuracy", - "--batch_size=1", - f"--hf_model_dir={ngram_root}", - f"--engine_dir={target_engine_dir}", - f"--dataset_dir={llm_datasets_root}", - f"--rouge_dir={llm_rouge_root}", - "--kv_cache_enable_block_reuse", - f"--ngram_config={ngram_config}", - "--tensorrt_llm_rouge1_threshold=20", - f"--kv_cache_free_gpu_memory_fraction=0.25", - ] - - venv_check_call(llm_venv, run_cmd) diff --git a/tests/integration/defs/examples/test_openai.py b/tests/integration/defs/examples/test_openai.py deleted file mode 100644 index 4f4776926ac2..000000000000 --- a/tests/integration/defs/examples/test_openai.py +++ /dev/null @@ -1,196 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -"""Module test_openai test openai examples.""" -import os -import subprocess # fmt: off - -import pytest -from defs.common import find_tensorrt, venv_check_call -from defs.trt_test_alternative import call, check_call, make_clean_dirs - - -@pytest.fixture(scope="module") -def openai_triton_example_root(llm_root): - "Get openai-triton example root" - example_root = os.path.join(llm_root, "examples", "openai_triton", - "manual_plugin") - - return example_root - - -@pytest.fixture(scope="module") -def openai_triton_plugingen_example_root(llm_root): - "Get openai-triton PluginGen example root" - example_root = os.path.join(llm_root, "examples", "openai_triton", - "plugin_autogen") - - return example_root - - -@pytest.fixture(scope="module") -def llm_openai_triton_model_root(llm_venv): - "prepare openai-triton model & return model root" - workspace = llm_venv.get_working_directory() - model_root = os.path.join(workspace, "triton") - commit = "d4644d6cb3ae674e1f15932cac1f28104795744f" - - call(f"git clone https://github.com/openai/triton.git {model_root}", - shell=True) - call(f"cd {model_root} && git checkout {commit}", shell=True) - llm_venv.run_cmd(["-m", "pip", "install", "cmake"]) - llm_venv.run_cmd([ - "-m", "pip", "install", - os.path.abspath(os.path.join(model_root, "python")) - ]) - - yield model_root - - llm_venv.run_cmd(["-m", "pip", "uninstall", "-y", "triton"]) - - -def test_llm_openai_triton_1gpu(openai_triton_example_root, - llm_openai_triton_model_root, llm_venv, - engine_dir, trt_config, is_trt_environment): - aot_path = os.path.join(openai_triton_example_root, "aot") - aot_fp16_path = os.path.join(aot_path, "fp16") - aot_fp32_path = os.path.join(aot_path, "fp32") - call(f"mkdir -p {aot_fp16_path}", shell=True) - call(f"mkdir -p {aot_fp32_path}", shell=True) - - num_stages = "2" - - # yapf: disable - # Kernel for data type=float16, BLOCK_M=128, BLOCK_DMODEL=64, BLOCK_N=128 - compile_cmd = [ - f"{llm_openai_triton_model_root}/python/triton/tools/compile.py", - f"{openai_triton_example_root}/fmha_triton.py", - "-n", "fused_attention_kernel", - "-o", f"{aot_fp16_path}/fmha_kernel_d64_fp16", - "--out-name", "fmha_d64_fp16", "-w", "4", "-ns", num_stages, - "-s", "*fp16:16, *fp32:16, *fp32:16, *fp16:16, *fp16:16, *fp16:16, fp32, i32, i32, i32, 128, 64, 128", - "-g", "(seq_len + 127) / 128, batch_size * num_heads, 1" - ] - venv_check_call(llm_venv, compile_cmd) - - # Kernel for data type=float32, BLOCK_M=64, BLOCK_DMODEL=64, BLOCK_N=64 - compile_cmd = [ - f"{llm_openai_triton_model_root}/python/triton/tools/compile.py", - f"{openai_triton_example_root}/fmha_triton.py", - "-n", "fused_attention_kernel", - "-o", f"{aot_fp32_path}/fmha_kernel_d64_fp32", - "--out-name", "fmha_d64_fp32", "-w", "4", "-ns", num_stages, - "-s", "*fp32:16, *fp32:16, *fp32:16, *fp32:16, *fp32:16, *fp32:16, fp32, i32, i32, i32, 64, 64, 64", - "-g", "(seq_len + 63) / 64, batch_size * num_heads, 1" - ] - venv_check_call(llm_venv, compile_cmd) - - # Link generated headers and create dispatchers. - check_call( - f"python3 {llm_openai_triton_model_root}/python/triton/tools/link.py " - f"{aot_fp16_path}/*.h -o {aot_path}/fmha_kernel_fp16", - shell=True) - check_call( - f"python3 {llm_openai_triton_model_root}/python/triton/tools/link.py " - f"{aot_fp32_path}/*.h -o {aot_path}/fmha_kernel_fp32", - shell=True) - - build_path = os.path.join(openai_triton_example_root, "build") - # yapf: enable - - # make files - make_clean_dirs(build_path) - cmake_args = [] - try: - import trt_test # noqa - except ImportError: - pass - else: - trt_include_dir, trt_lib_dir = find_tensorrt( - trt_config["new_ld_library_path"]) - - if trt_include_dir: - cmake_args.append(f"-DTRT_INCLUDE_DIR={trt_include_dir}") - - if trt_lib_dir: - cmake_args.append(f"-DTRT_LIB_DIR={trt_lib_dir}") - - if is_trt_environment: - cmake_args.append(f"-DCMAKE_C_FLAGS='-D_GLIBCXX_USE_CXX11_ABI=0'") - - cmake_args = " ".join(cmake_args) - check_call(f"cd {build_path} && cmake {cmake_args} .. && make", shell=True) - - # build engine - build_cmd = [ - f"{openai_triton_example_root}/build.py", "--num_heads=32", - "--head_size=64", "--max_batch_size=8", "--max_seq_len=512", - "--dtype=float16", f"--output={engine_dir}" - ] - venv_check_call(llm_venv, build_cmd) - - # run inference - run_cmd = [ - f"{openai_triton_example_root}/run.py", "--num_heads=32", - "--head_size=64", "--batch_size=8", "--seq_len=512", - "--log_level=verbose", "--benchmark", f"--engine_dir={engine_dir}" - ] - venv_check_call(llm_venv, run_cmd) - - -# TODO[chunweiy]: Enable it later -def test_llm_openai_triton_plugingen_1gpu(openai_triton_plugingen_example_root, - openai_triton_example_root, - llm_openai_triton_model_root, - plugin_gen_path, llm_venv, - trt_config): - # copy the triton kernel definition - subprocess.run( - f"cp {openai_triton_example_root}/fmha_triton.py {openai_triton_plugingen_example_root}/fmha_triton.py" - .split(), - check=True) - - # generate plugin - cmd = [ - plugin_gen_path, - "--workspace", - "./tmp", - "--kernel_config", - os.path.join(openai_triton_plugingen_example_root, "kernel_config.py"), - ] - try: - import trt_test # noqa - except ImportError: - pass - else: - trt_include_dir, trt_lib_dir = find_tensorrt( - trt_config["new_ld_library_path"]) - if trt_lib_dir is not None: - cmd.append(f'--trt_lib_dir={trt_lib_dir}') - if trt_include_dir is not None: - cmd.append(f'--trt_include_dir={trt_include_dir}') - - venv_check_call(llm_venv, cmd) - - # build engine - cmd = [ - os.path.join(openai_triton_plugingen_example_root, "build_engine.py"), - ] - venv_check_call(llm_venv, cmd) - - # run engine - cmd = [ - os.path.join(openai_triton_plugingen_example_root, "run_engine.py"), - ] - venv_check_call(llm_venv, cmd) diff --git a/tests/integration/defs/examples/test_phi.py b/tests/integration/defs/examples/test_phi.py index 62f67e4bae65..a8c18f731e94 100644 --- a/tests/integration/defs/examples/test_phi.py +++ b/tests/integration/defs/examples/test_phi.py @@ -16,9 +16,8 @@ import defs.ci_profiler import pytest -from defs.common import test_llm_torch_multi_lora_support, venv_check_call -from defs.conftest import get_sm_version, skip_post_blackwell, skip_pre_ada -from defs.trt_test_alternative import check_call +from defs.common import test_llm_torch_multi_lora_support +from defs.conftest import get_sm_version, skip_pre_ada # skip trt flow cases on post-Blackwell-Ultra if get_sm_version() >= 103: @@ -39,75 +38,6 @@ def phi_example_root(llm_root, llm_venv): return example_root -@skip_pre_ada -@pytest.mark.parametrize("data_type", ['float16', 'bfloat16']) -@pytest.mark.parametrize("qformat", ['fp8']) -@pytest.mark.parametrize("llm_phi_model_root", [ - pytest.param("phi-2", marks=skip_post_blackwell), - pytest.param("Phi-3-mini-128k-instruct", marks=skip_post_blackwell), - pytest.param("Phi-3-small-128k-instruct", marks=skip_post_blackwell), - pytest.param("Phi-3.5-mini-instruct", marks=skip_post_blackwell), - "Phi-3.5-MoE-instruct", "Phi-4-mini-instruct" -], - indirect=True) -def test_llm_phi_quantization_1gpu(data_type, llm_phi_model_root, llm_venv, - cmodel_dir, engine_dir, phi_example_root, - llm_datasets_root, llm_rouge_root, qformat): - "Run phi quantization tests" - # Workaround for Modelopt can't convert Phi-3 on multi GPUs. - gpu_constraint = {"CUDA_VISIBLE_DEVICES": "0"} - - print("Convert checkpoint by modelopt...") - convert_cmd = [ - f"{phi_example_root}/../../../quantization/quantize.py", - f"--model_dir={llm_phi_model_root}", - f"--calib_dataset={llm_datasets_root}/cnn_dailymail", - f"--dtype={data_type}", - f"--qformat={qformat}", - f"--kv_cache_dtype={qformat}", - f"--output_dir={cmodel_dir}", - ] - venv_check_call(llm_venv, convert_cmd, env=gpu_constraint) - - print("Build engines...") - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={cmodel_dir}", - f"--output_dir={engine_dir}", - "--max_input_len=3000", - "--max_seq_len=3100", - f"--max_batch_size={16}", - ] - - build_env = { - **llm_venv._new_env, - **gpu_constraint - } if llm_venv._new_env else gpu_constraint - check_call(" ".join(build_cmd), shell=True, env=build_env) - - print("Run summarize...") - threshold_score = 24.0 - model_name = os.path.basename(llm_phi_model_root) - if model_name == "phi-2": - threshold_score = 22.0 - - summary_cmd = [ - f"{phi_example_root}/../../../summarize.py", - "--test_trt_llm", - f"--hf_model_dir={llm_phi_model_root}", - f"--tokenizer_dir={llm_phi_model_root}", - f"--engine_dir={engine_dir}", - "--check_accuracy", - f"--tensorrt_llm_rouge1_threshold={threshold_score}", - "--max_ite=40", - f"--batch_size={16}", - f"--dataset_dir={llm_datasets_root}", - f"--rouge_dir={llm_rouge_root}", - ] - - venv_check_call(llm_venv, summary_cmd, env=gpu_constraint) - - @pytest.mark.skip( reason="TODO: Resolve an import issue with transformers's LossKwargs") @skip_pre_ada diff --git a/tests/integration/defs/examples/test_qwen.py b/tests/integration/defs/examples/test_qwen.py deleted file mode 100644 index 35352b20e279..000000000000 --- a/tests/integration/defs/examples/test_qwen.py +++ /dev/null @@ -1,242 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -"""Module test_qwen test qwen examples.""" - -import csv -import os - -import pytest -from defs.common import (convert_weights, test_multi_lora_support, - venv_check_call) -from defs.conftest import get_sm_version, skip_pre_ada -from defs.trt_test_alternative import check_call - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - -# Delete this case refer to https://nvbugs/5072417 -# @pytest.mark.parametrize("llm_lora_model_root", ["Ko-QWEN-7B-Chat-LoRA"], -# indirect=True) -# @pytest.mark.parametrize("llm_qwen_model_root", ["qwen_7b_chat"], indirect=True) -# def test_llm_qwen_7b_single_gpu_lora( -# qwen_example_root, -# llm_qwen_model_root, -# llm_venv, -# cmodel_dir, -# engine_dir, -# llm_lora_model_root, -# ): -# "run Qwen lora test on single gpu." -# print("Build engines...") -# dtype = 'float16' -# model_name = os.path.basename(llm_qwen_model_root) -# ckpt_dir = convert_weights(llm_venv=llm_venv, -# example_root=qwen_example_root, -# cmodel_dir=cmodel_dir, -# model=model_name, -# model_path=llm_qwen_model_root, -# data_type=dtype) - -# print("Build engines...") -# build_cmd = [ -# "trtllm-build", -# f"--checkpoint_dir={ckpt_dir}", -# f"--output_dir={engine_dir}", -# "--lora_plugin=auto", -# "--gemm_plugin=auto", -# f"--lora_dir={llm_lora_model_root}", -# ] -# check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - -# ref_1 = [ -# 151644, 8948, 198, 2610, 525, 264, 10950, 17847, 13, 151645, 198, -# 151644, 872, 198, 126246, 144370, 91145, 11, 137601, 29326, 86034, -# 12802, 5140, 98734, 19391, 35711, 30, 151645, 198, 151644, 77091, 198, -# 126246, 144370, 91145, 0, 134561, 58677, 78125, 21329, 66019, 124685, -# 134619, 94152, 28626, 17380, 11, 134637, 20401, 138520, 19391, 143603 -# ] -# ref_2 = [ -# 151644, 8948, 198, 2610, 525, 264, 10950, 17847, 13, 151645, 198, -# 151644, 872, 198, 126246, 144370, 91145, 11, 137601, 29326, 86034, -# 12802, 5140, 98734, 19391, 35711, 30, 151645, 198, 151644, 77091, 198, -# 126246, 144370, 91145, 0, 134561, 330, 48, 1103, 54, 268, 1, 78952, 13, -# 151645, 198, 151643, 151643, 151643, 151643, 151643 -# ] - -# input_text = "안녕하세요, 혹시 이름이 뭐에요?" -# print("Run inference with lora id 0...") -# venv_check_call(llm_venv, [ -# f"{qwen_example_root}/../run.py", -# "--max_output_len=20", -# f"--input_text={input_text}", -# "--lora_task_uids=0", -# f"--tokenizer_dir={llm_qwen_model_root}", -# f"--engine_dir={engine_dir}", -# f"--output_csv={llm_venv.get_working_directory()}/use_lora.csv", -# "--use_py_session", -# ]) - -# with open(f"{llm_venv.get_working_directory()}/use_lora.csv") as f: -# predict = csv.reader(f) -# predict = next(predict) -# predict = [int(p) for p in predict] -# assert ref_1 == predict - -# print("Run inference with lora id -1...") -# venv_check_call(llm_venv, [ -# f"{qwen_example_root}/../run.py", -# "--max_output_len=20", -# f"--input_text={input_text}", -# "--lora_task_uids=-1", -# f"--tokenizer_dir={llm_qwen_model_root}", -# f"--engine_dir={engine_dir}", -# f"--output_csv={llm_venv.get_working_directory()}/no_lora.csv", -# "--use_py_session", -# ]) - -# with open(f"{llm_venv.get_working_directory()}/no_lora.csv") as f: -# predict = csv.reader(f) -# predict = next(predict) -# predict = [int(p) for p in predict] -# assert ref_2 == predict - - -@pytest.mark.parametrize("llm_lora_model_root", ["Qwen1.5-7B-Chat-750Mb-lora"], - indirect=True) -@pytest.mark.parametrize("llm_qwen_model_root", ["qwen1.5_7b_chat"], - indirect=True) -def test_llm_qwen1_5_7b_single_gpu_lora( - qwen_example_root, - llm_qwen_model_root, - llm_venv, - cmodel_dir, - engine_dir, - llm_lora_model_root, -): - "run Qwen1.5 lora test on single gpu." - print("Build engines...") - dtype = 'float16' - model_name = os.path.basename(llm_qwen_model_root) - ckpt_dir = convert_weights(llm_venv=llm_venv, - example_root=qwen_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=llm_qwen_model_root, - data_type=dtype) - - print("Build engines...") - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={ckpt_dir}", - f"--output_dir={engine_dir}", - "--lora_plugin=auto", - "--gemm_plugin=auto", - f"--lora_dir={llm_lora_model_root}", - ] - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - ref_1 = [ - 151644, 8948, 198, 2610, 525, 264, 10950, 17847, 13, 151645, 198, - 151644, 872, 198, 3838, 374, 697, 829, 30, 151645, 198, 151644, 77091, - 198, 40, 2776, 458, 15235, 7881, 553, 5264, 15469, 11, 773, 358, 1513, - 944, 614, 264, 829, 304, 279, 8606, 5530 - ] - ref_2 = [ - 151644, 8948, 198, 2610, 525, 264, 10950, 17847, 13, 151645, 198, - 151644, 872, 198, 3838, 374, 697, 829, 30, 151645, 198, 151644, 77091, - 198, 40, 1079, 1207, 16948, 11, 264, 3460, 4128, 1614, 3465, 553, 54364, - 14817, 13, 151645, 151645, 151645, 151645, 151645, 151645 - ] - - input_text = "What is your name?" - print("Run inference with lora id 0...") - venv_check_call(llm_venv, [ - f"{qwen_example_root}/../../../run.py", - "--max_output_len=20", - f"--input_text={input_text}", - "--lora_task_uids=0", - f"--tokenizer_dir={llm_qwen_model_root}", - f"--engine_dir={engine_dir}", - f"--output_csv={llm_venv.get_working_directory()}/use_lora.csv", - "--use_py_session", - ]) - - with open(f"{llm_venv.get_working_directory()}/use_lora.csv") as f: - predict = csv.reader(f) - predict = next(predict) - predict = [int(p) for p in predict] - assert ref_1 == predict - - print("Run inference with lora id -1...") - venv_check_call(llm_venv, [ - f"{qwen_example_root}/../../../run.py", - "--max_output_len=20", - f"--input_text={input_text}", - "--lora_task_uids=-1", - f"--tokenizer_dir={llm_qwen_model_root}", - f"--engine_dir={engine_dir}", - f"--output_csv={llm_venv.get_working_directory()}/no_lora.csv", - "--use_py_session", - ]) - - with open(f"{llm_venv.get_working_directory()}/no_lora.csv") as f: - predict = csv.reader(f) - predict = next(predict) - predict = [int(p) for p in predict] - assert ref_2 == predict - - -@skip_pre_ada -@pytest.mark.parametrize( - "llm_qwen_model_root", - ["qwen2_0.5b_instruct", "qwen2.5_0.5b_instruct", "qwen2.5_1.5b_instruct"], - indirect=True) -def test_llm_hf_qwen_multi_lora_1gpu(llm_qwen_model_root, - llm_venv, - cmodel_dir, - engine_dir, - qwen_example_root, - llm_datasets_root, - qformat='fp8', - dtype='bfloat16'): - "Run Qwen models with multiple dummy LoRAs." - - print("Convert checkpoint by modelopt...") - convert_cmd = [ - f"{qwen_example_root}/../../../quantization/quantize.py", - f"--model_dir={llm_qwen_model_root}", - f"--calib_dataset={llm_datasets_root}/cnn_dailymail", - f"--dtype={dtype}", - f"--qformat={qformat}", - f"--kv_cache_dtype={qformat}", - f"--output_dir={cmodel_dir}", - ] - venv_check_call(llm_venv, convert_cmd) - - test_multi_lora_support( - hf_model_dir=llm_qwen_model_root, - tllm_ckpt_dir=cmodel_dir, - engine_dir=engine_dir, - llm_venv=llm_venv, - example_root=qwen_example_root, - num_loras=2, - lora_rank=8, - target_hf_modules=["q_proj", "k_proj", "v_proj"], - target_trtllm_modules=["attn_q", "attn_k", "attn_v"], - zero_lora_weights=True, - ) diff --git a/tests/integration/defs/examples/test_qwen2audio.py b/tests/integration/defs/examples/test_qwen2audio.py deleted file mode 100644 index 9e911cca97e8..000000000000 --- a/tests/integration/defs/examples/test_qwen2audio.py +++ /dev/null @@ -1,108 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -"""Module test_qwen test qwen2audio examples.""" - -import os -import re - -import pytest -from defs.common import venv_check_call, venv_check_output -from defs.trt_test_alternative import check_call - - -@pytest.fixture(scope="module") -def qwen2audio_example_root(llm_root, llm_venv): - "Get qwen2audio example root" - example_root = os.path.join(llm_root, "examples", "models", "core", - "qwen2audio") - llm_venv.run_cmd([ - "-m", "pip", "install", "-r", - os.path.join(example_root, "requirements.txt") - ]) - - return example_root - - -@pytest.mark.parametrize("llm_qwen_model_root", ["qwen2_audio_7b_instruct"], - indirect=True) -def test_llm_qwen2audio_single_gpu(qwen2audio_example_root, llm_qwen_model_root, - llm_venv, engine_dir): - "Build & run qwen2audio on 1 gpu." - workspace = llm_venv.get_working_directory() - - print("Generate audio engine...") - audio_engine_dir = f"{engine_dir}/audio" - audio_cmd = [ - f"{qwen2audio_example_root}/../multimodal/build_multimodal_engine.py", - f"--model_type=qwen2_audio", - f"--model_path={llm_qwen_model_root}", - f"--max_batch_size=32", - f"--output_dir={audio_engine_dir}", - ] - - venv_check_call(llm_venv, audio_cmd) - - print("Convert checkpoint...") - convert_cmd = [ - f"{qwen2audio_example_root}/../qwen/convert_checkpoint.py", - f"--model_dir={llm_qwen_model_root}", - f"--output_dir={workspace}/Qwen2-Audio", - f"--dtype=float16", - ] - - venv_check_call(llm_venv, convert_cmd) - - print("Build TRT-LLM engine...") - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={workspace}/Qwen2-Audio", - f"--gemm_plugin=float16", - f"--gpt_attention_plugin=float16", - f"--max_prompt_embedding_table_size=4096", - f"--output_dir={engine_dir}", - f"--max_batch_size={1}", - ] - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Run engine...") - audio_url = f"{qwen2audio_example_root}/audio/glass-breaking-151256.mp3" - - run_cmd = [ - f"{qwen2audio_example_root}/run.py", - f"--tokenizer_dir={llm_qwen_model_root}", - f"--engine_dir={engine_dir}", - f"--audio_engine_path={audio_engine_dir}/model.engine", - f"--audio_url={audio_url}", - ] - - output = venv_check_output(llm_venv, run_cmd) - output = [line for line in output.split("\n") if "Output:" in line] - print(output) - - print("Verify the output...") - results = [] - for item in output: - match = re.search(r"Output: \"(.*)", item) - if match: - results.append(match.group(1)) - - for item in results: - # check the output if it contains key words - item = item.lower() - if ("glass" in item) and ("shatter" in item or "break" in item): - pass - else: - assert False, f"output is: {item}" diff --git a/tests/integration/defs/examples/test_qwenvl.py b/tests/integration/defs/examples/test_qwenvl.py deleted file mode 100644 index d4ff5765aa33..000000000000 --- a/tests/integration/defs/examples/test_qwenvl.py +++ /dev/null @@ -1,39 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -"""Module test_qwen test qwenvl examples.""" - -import os - -import pytest -from defs.conftest import get_sm_version - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -@pytest.fixture(scope="module") -def qwenvl_example_root(llm_root, llm_venv): - "Get qwenvl example root" - example_root = os.path.join(llm_root, "examples", "models", "core", - "qwenvl") - llm_venv.run_cmd([ - "-m", "pip", "install", "-r", - os.path.join(example_root, "requirements.txt") - ]) - - return example_root diff --git a/tests/integration/defs/examples/test_redrafter.py b/tests/integration/defs/examples/test_redrafter.py deleted file mode 100644 index ce9a62d097b2..000000000000 --- a/tests/integration/defs/examples/test_redrafter.py +++ /dev/null @@ -1,99 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. - -import pytest -from defs.common import convert_weights, venv_check_call -from defs.conftest import get_sm_version, skip_post_blackwell -from defs.trt_test_alternative import check_call - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -@skip_post_blackwell -@pytest.mark.parametrize("batch_size", [8], ids=['bs8']) -@pytest.mark.parametrize("redrafter_num_beams", [5, 8], ids=['nb5', 'nb8']) -@pytest.mark.parametrize("redrafter_draft_len_per_beam", [5], ids=['dl5']) -@pytest.mark.parametrize("data_type", ['bfloat16']) -@pytest.mark.parametrize("redrafter_model_roots", ["redrafter-vicuna-7b-v1.3"], - indirect=True) -@pytest.mark.parametrize("use_py_session", [False, True], - ids=["use_cpp_session", "use_py_session"]) -def test_llm_redrafter_1gpu(batch_size, data_type, redrafter_model_roots, - redrafter_num_beams, redrafter_draft_len_per_beam, - redrafter_example_root, llama_example_root, - llm_datasets_root, llm_rouge_root, llm_venv, - cmodel_dir, cmodel_base_dir, engine_dir, - use_py_session): - print("Build engines...") - model_name = "redrafter" - base_model_name = "llama" - base_example_root = llama_example_root - - base_model_dir = convert_weights(llm_venv=llm_venv, - example_root=base_example_root, - cmodel_dir=cmodel_base_dir, - model=base_model_name, - model_path=redrafter_model_roots[0], - data_type=data_type) - - redrafter_convert_roots = (base_model_dir, redrafter_model_roots[1]) - - model_dir = convert_weights( - llm_venv=llm_venv, - example_root=redrafter_example_root, - cmodel_dir=cmodel_dir, - model=model_name, - model_path=redrafter_convert_roots, - data_type=data_type, - redrafter_num_beams=redrafter_num_beams, - redrafter_draft_len_per_beam=redrafter_draft_len_per_beam) - - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={model_dir}", - f"--output_dir={engine_dir}", - f"--gpt_attention_plugin={data_type}", - f"--gemm_plugin={data_type}", - f"--max_beam_width=1", - "--remove_input_padding=enable", - "--context_fmha=enable", - "--max_input_len=1024", - "--max_seq_len=1536", - f"--max_batch_size={batch_size}", - "--kv_cache_type=paged", - '--speculative_decoding_mode=explicit_draft_tokens', - ] - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - print("Run summarize...") - - summary_cmd = [ - f"{redrafter_example_root}/../summarize.py", "--test_trt_llm", - "--hf_model_dir", f"{redrafter_model_roots[0]}", "--tokenizer_dir", - f"{redrafter_model_roots[0]}", f"--engine_dir={engine_dir}", - "--check_accuracy", "--tensorrt_llm_rouge1_threshold=24", - f"--temperature=1.0", f"--max_ite=40", f"--batch_size={batch_size}", - f"--dataset_dir={llm_datasets_root}", f"--rouge_dir={llm_rouge_root}" - ] - - if use_py_session: - summary_cmd.append("--use_py_session") - - venv_check_call(llm_venv, summary_cmd) diff --git a/tests/integration/defs/examples/test_whisper.py b/tests/integration/defs/examples/test_whisper.py deleted file mode 100644 index 5b7b6859ffaf..000000000000 --- a/tests/integration/defs/examples/test_whisper.py +++ /dev/null @@ -1,298 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 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. - -import os - -import pytest -from defs.common import convert_weights, venv_check_call -from defs.conftest import get_sm_version, llm_models_root, skip_post_blackwell -from defs.trt_test_alternative import check_call - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -@skip_post_blackwell -@pytest.mark.parametrize("use_cpp_runtime", [True, False], - ids=["use_cpp_runtime", "use_python_runtime"]) -@pytest.mark.parametrize("num_beams", [1, 4], - ids=lambda num_beams: f'nb:{num_beams}') -@pytest.mark.parametrize("data_type", ['float16']) -@pytest.mark.parametrize("weight_only_precision", [ - 'disable_weight_only', - pytest.param('int8', marks=skip_post_blackwell), - pytest.param('int4', marks=skip_post_blackwell) -]) -@pytest.mark.parametrize( - "use_attention_plugin", [True, False], - ids=["enable_attention_plugin", "disable_attention_plugin"]) -@pytest.mark.parametrize("use_gemm_plugin", [True, False], - ids=["enable_gemm_plugin", "disable_gemm_plugin"]) -@pytest.mark.parametrize("whisper_model_root", ['large-v3', 'large-v2'], - indirect=True) -def test_llm_whisper_general(llm_venv, engine_dir, data_type, - weight_only_precision, use_attention_plugin, - use_gemm_plugin, whisper_example_root, - whisper_model_root, num_beams, use_cpp_runtime, - whisper_example_audio_file, llm_datasets_root): - print("Locate model checkpoints in test storage...") - tllm_model_name, model_ckpt_dir = whisper_model_root - - if any((not use_attention_plugin, use_gemm_plugin, 'v3' - not in tllm_model_name)) and use_cpp_runtime: - pytest.skip(f"Plugins might not support C++ runtime. Skip the test...") - - whisper_engine_dir = f"{engine_dir}/{tllm_model_name}/{data_type}_{weight_only_precision}" - - if 'int' in weight_only_precision: - use_weight_only = True - else: - use_weight_only = False - weight_only_precision = None - converted_weight_dir = convert_weights( - llm_venv=llm_venv, - example_root=whisper_example_root, - cmodel_dir=whisper_engine_dir, - model=tllm_model_name, - model_path=model_ckpt_dir, - use_weight_only=use_weight_only, - weight_only_precision=weight_only_precision) - print("Build engines...") - for component in ["encoder", "decoder"]: - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={converted_weight_dir}/{component}", - f"--output_dir={whisper_engine_dir}/{component}", - "--paged_kv_cache=disable", - "--moe_plugin=disable", - "--max_batch_size=8", - ] - if use_cpp_runtime: - build_cmd.extend( - ("--paged_kv_cache enable", "--remove_input_padding enable")) - else: - build_cmd.append("--remove_input_padding=disable") - - if component == "encoder": - build_cmd.append( - f"--max_input_len=3000" - ) # check against actual encoder features length (3000,...) in C++ runtime - build_cmd.append(f"--max_seq_len=3000") - if component == "decoder": - build_cmd.append(f"--max_input_len=14") - build_cmd.append(f"--max_seq_len=114") - build_cmd.append(f"--max_encoder_input_len=3000") - build_cmd.append(f"--max_beam_width={num_beams}") - - if use_gemm_plugin: - build_cmd.append(f"--gemm_plugin={data_type}") - else: - build_cmd.append(f"--gemm_plugin=disable") - - if use_attention_plugin: - build_cmd.append(f"--bert_attention_plugin={data_type}") - build_cmd.append(f"--gpt_attention_plugin={data_type}") - else: - build_cmd.append(f"--bert_attention_plugin=disable") - build_cmd.append(f"--gpt_attention_plugin=disable") - - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - if use_cpp_runtime: - print("Run inference using Python bindings of C++ runtime...") - run_cmd = [ - f'{whisper_example_root}/../../../run.py', - f'--tokenizer_dir={llm_models_root()}/whisper-{tllm_model_name}', - f'--multimodal_input_file={whisper_example_audio_file}', - f'--engine_dir={whisper_engine_dir}', - f'--max_output_len=96', - ] - else: - print("Run inference using Whisper's custom Python runtime...") - run_cmd = [ - f"{whisper_example_root}/run.py", - f"--dataset={llm_datasets_root}/hf-internal-testing/librispeech_asr_dummy", - f"--engine_dir={whisper_engine_dir}", - f"--assets_dir={model_ckpt_dir}", - f"--num_beams={num_beams}", - f"--dtype={data_type}", - f"--use_py_session", - f"--accuracy_check", - ] - # https://nvbugs/4658787 - # WAR before whisper tests can work offline - env = {"HF_DATASETS_OFFLINE": "0"} - venv_check_call(llm_venv, run_cmd, env=env) - - -@skip_post_blackwell -@pytest.mark.parametrize("num_beams", [4], - ids=lambda num_beams: f'nb:{num_beams}') -@pytest.mark.parametrize("whisper_model_root", ['large-v3'], indirect=True) -def test_whisper_beam_search_generation_logits(llm_venv, engine_dir, - whisper_example_root, - whisper_model_root, num_beams, - whisper_example_audio_file): - """Verify that generation_logits are reordered to match the final beam paths. - - With beam search (num_beams > 1), generation_logits must be reindexed by - parentIds so that argmax(generation_logits[beam][t]) == output_ids[beam][t] - at every position. Without reordering, logits are indexed by beam slot - rather than by the final beam path, producing incorrect probabilities. - """ - tllm_model_name, model_ckpt_dir = whisper_model_root - - whisper_engine_dir = f"{engine_dir}/{tllm_model_name}/float16_disable_weight_only" - - converted_weight_dir = convert_weights(llm_venv=llm_venv, - example_root=whisper_example_root, - cmodel_dir=whisper_engine_dir, - model=tllm_model_name, - model_path=model_ckpt_dir, - use_weight_only=False, - weight_only_precision=None) - - print("Build engines for beam search generation logits test...") - for component in ["encoder", "decoder"]: - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={converted_weight_dir}/{component}", - f"--output_dir={whisper_engine_dir}/{component}", - "--paged_kv_cache=enable", - "--remove_input_padding=enable", - "--moe_plugin=disable", - "--max_batch_size=1", - ] - if component == "encoder": - build_cmd.append("--max_input_len=3000") - build_cmd.append("--max_seq_len=3000") - if component == "decoder": - build_cmd.append("--max_input_len=14") - build_cmd.append("--max_seq_len=114") - build_cmd.append("--max_encoder_input_len=3000") - build_cmd.append(f"--max_beam_width={num_beams}") - build_cmd.append("--gemm_plugin=float16") - build_cmd.append("--bert_attention_plugin=float16") - build_cmd.append("--gpt_attention_plugin=float16") - - check_call(build_cmd, env=llm_venv._new_env) - - print("Run generation logits beam search validation...") - validation_script = os.path.join(os.path.dirname(__file__), - "validate_whisper_beam_logits.py") - run_cmd = [ - validation_script, - f"--engine_dir={whisper_engine_dir}", - f"--assets_dir={model_ckpt_dir}", - f"--input_file={whisper_example_audio_file}", - f"--num_beams={num_beams}", - ] - env = { - "HF_DATASETS_OFFLINE": - "0", - "PYTHONPATH": - os.pathsep.join( - filter(None, [whisper_example_root, - os.environ.get("PYTHONPATH")])), - } - venv_check_call(llm_venv, run_cmd, env=env) - - -@skip_post_blackwell -@pytest.mark.parametrize("num_beams", [4], - ids=lambda num_beams: f'nb:{num_beams}') -@pytest.mark.parametrize("batch_size", [4], - ids=lambda batch_size: f'bs:{batch_size}') -@pytest.mark.parametrize("whisper_model_root", ['large-v3'], indirect=True) -def test_whisper_log_probs_determinism(llm_venv, engine_dir, llm_root, - whisper_example_root, whisper_model_root, - num_beams, batch_size, - llm_datasets_root): - """Regression test for nMaxBatchSize stride bug in beam search log_probs. - - Sends a batch of requests with different audio samples so each item produces - a different number of output tokens. When the batch drains unevenly the - active batch dimension shrinks; the bug caused beamStage3 to use that - shrinking value as the logProbsTiled stride, producing non-deterministic - log_probs across runs. - - Runs inference num_runs times and asserts log_probs are bit-identical. - """ - tllm_model_name, model_ckpt_dir = whisper_model_root - - whisper_engine_dir = (f"{engine_dir}/{tllm_model_name}" - f"/float16_bs{batch_size}_nb{num_beams}_logprobs_det") - - converted_weight_dir = convert_weights(llm_venv=llm_venv, - example_root=whisper_example_root, - cmodel_dir=whisper_engine_dir, - model=tllm_model_name, - model_path=model_ckpt_dir, - use_weight_only=False, - weight_only_precision=None) - - print("Build engines for log_probs determinism test...") - for component in ["encoder", "decoder"]: - build_cmd = [ - "trtllm-build", - f"--checkpoint_dir={converted_weight_dir}/{component}", - f"--output_dir={whisper_engine_dir}/{component}", - "--paged_kv_cache=enable", - "--remove_input_padding=enable", - "--moe_plugin=disable", - f"--max_batch_size={batch_size}", - "--gemm_plugin=float16", - "--bert_attention_plugin=float16", - "--gpt_attention_plugin=float16", - ] - if component == "encoder": - build_cmd.append("--max_input_len=3000") - build_cmd.append("--max_seq_len=3000") - if component == "decoder": - build_cmd.append("--max_input_len=14") - build_cmd.append("--max_seq_len=114") - build_cmd.append("--max_encoder_input_len=3000") - build_cmd.append(f"--max_beam_width={num_beams}") - - check_call(build_cmd, env=llm_venv._new_env) - - print("Run log_probs determinism validation...") - validation_script = os.path.join( - os.path.dirname(__file__), "validate_whisper_log_probs_determinism.py") - librispeech_dir = os.path.join(llm_datasets_root, - "hf-internal-testing/librispeech_asr_dummy") - run_cmd = [ - validation_script, - f"--engine_dir={whisper_engine_dir}", - f"--assets_dir={model_ckpt_dir}", - f"--dataset_dir={librispeech_dir}", - f"--num_beams={num_beams}", - f"--batch_size={batch_size}", - "--num_runs=5", - ] - env = { - # llm_root is listed first so the worktree's tensorrt_llm takes priority - # over any system-installed version when the validation subprocess runs. - "PYTHONPATH": - os.pathsep.join( - filter( - None, - [llm_root, whisper_example_root, - os.environ.get("PYTHONPATH")])), - } - venv_check_call(llm_venv, run_cmd, env=env) diff --git a/tests/integration/defs/examples/validate_whisper_beam_logits.py b/tests/integration/defs/examples/validate_whisper_beam_logits.py deleted file mode 100644 index faed875b9b6c..000000000000 --- a/tests/integration/defs/examples/validate_whisper_beam_logits.py +++ /dev/null @@ -1,205 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 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. -"""Validate that generation_logits are correctly reordered for beam search. - -With beam search (num_beams > 1), generation_logits must be reindexed to match -the final beam paths after gatherTree finalization. This script runs whisper -inference via ModelRunnerCpp and checks that each output token has a reasonable -probability under its corresponding generation logits. - -Exits with non-zero status if any output token has near-zero probability -(log P < -10), which indicates logits from a different beam's context. -""" - -import argparse -import json -import sys -from collections import OrderedDict -from pathlib import Path - -import numpy as np -import torch - -import tensorrt_llm -from tensorrt_llm._utils import str_dtype_to_torch -from tensorrt_llm.runtime import ModelRunnerCpp - - -def read_config(component, engine_dir): - config_path = engine_dir / component / "config.json" - with open(config_path, "r") as f: - config = json.load(f) - model_config = OrderedDict() - model_config.update(config["pretrained_config"]) - model_config.update(config["build_config"]) - return model_config - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument("--engine_dir", type=str, required=True) - parser.add_argument("--assets_dir", type=str, required=True) - parser.add_argument("--input_file", type=str, default=None) - parser.add_argument("--num_beams", type=int, default=4) - parser.add_argument("--max_new_tokens", type=int, default=96) - return parser.parse_args() - - -def get_mel_input(input_file, n_mels, mel_filters_dir): - from whisper_utils import log_mel_spectrogram - - if input_file: - mel, _ = log_mel_spectrogram( - input_file, n_mels, device="cuda", return_duration=True, mel_filters_dir=mel_filters_dir - ) - else: - from datasets import load_dataset - - dataset = load_dataset( - "hf-internal-testing/librispeech_asr_dummy", - "clean", - split="validation", - trust_remote_code=True, - ) - speech = dataset[0]["audio"]["array"].astype(np.float32) - waveform = torch.from_numpy(speech) - mel = log_mel_spectrogram(waveform, n_mels, device="cuda", mel_filters_dir=mel_filters_dir) - - mel = mel.type(str_dtype_to_torch("float16")) - mel = mel.unsqueeze(0) - if mel.shape[2] % 2: - mel = torch.nn.functional.pad(mel, (0, 1)) - return mel - - -def validate_logits_alignment(output_ids, generation_logits, input_len, eot_id): - """Check that output tokens have reasonable probability under generation_logits. - - Returns True if all output tokens have log P > -10 across all beams, False otherwise. - """ - LOG_PROB_THRESHOLD = -10.0 - batch_size = output_ids.shape[0] - num_beams = output_ids.shape[1] - all_aligned = True - - for b in range(batch_size): - for beam in range(num_beams): - gen_tokens = output_ids[b, beam, input_len:] - eot_positions = (gen_tokens == eot_id).nonzero(as_tuple=True)[0] - gen_len = eot_positions[0].item() if len(eot_positions) > 0 else gen_tokens.shape[0] - - if gen_len == 0: - continue - - gen_tokens = gen_tokens[:gen_len] - logits = generation_logits[b, beam, :gen_len, :] - - log_probs = torch.nn.functional.log_softmax(logits, dim=-1) - actual_logprobs = log_probs.gather(1, gen_tokens.unsqueeze(1)).squeeze(1) - - min_logprob = actual_logprobs.min().item() - near_zero = (actual_logprobs < LOG_PROB_THRESHOLD).sum().item() - - argmax_matches = (logits.argmax(dim=-1) == gen_tokens).sum().item() - - print( - f" Batch {b}, beam {beam}: argmax match {argmax_matches}/{gen_len}, " - f"min log P = {min_logprob:.4f}, " - f"near-zero positions = {near_zero}/{gen_len}" - ) - - if near_zero > 0: - all_aligned = False - print(f" FAIL: {near_zero} positions have near-zero probability") - - return all_aligned - - -def main(): - args = parse_arguments() - tensorrt_llm.logger.set_level("warning") - - engine_dir = Path(args.engine_dir) - encoder_config = read_config("encoder", engine_dir) - decoder_config = read_config("decoder", engine_dir) - - n_mels = encoder_config["n_mels"] - is_multilingual = decoder_config["vocab_size"] >= 51865 - - from tokenizer import get_tokenizer - - tokenizer_name = "multilingual" if is_multilingual else "gpt2" - tokenizer = get_tokenizer( - name=tokenizer_name, - num_languages=encoder_config["num_languages"], - tokenizer_dir=args.assets_dir, - ) - eot_id = tokenizer.encode("<|endoftext|>", allowed_special=tokenizer.special_tokens_set)[0] - - runner = ModelRunnerCpp.from_dir( - engine_dir=engine_dir, - is_enc_dec=True, - max_batch_size=1, - max_input_len=3000, - max_output_len=args.max_new_tokens, - max_beam_width=args.num_beams, - kv_cache_free_gpu_memory_fraction=0.9, - cross_kv_cache_fraction=0.5, - gather_generation_logits=True, - ) - - mel = get_mel_input(args.input_file, n_mels, args.assets_dir) - mel_input_lengths = torch.full((1,), mel.shape[2], dtype=torch.int32, device="cuda") - - prompt_text = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>" - prompt_ids = tokenizer.encode(prompt_text, allowed_special=tokenizer.special_tokens_set) - decoder_input_ids = torch.tensor(prompt_ids).unsqueeze(0) - input_len = decoder_input_ids.shape[1] - - with torch.no_grad(): - outputs = runner.generate( - batch_input_ids=decoder_input_ids, - encoder_input_features=mel.transpose(1, 2), - encoder_output_lengths=mel_input_lengths // 2, - max_new_tokens=args.max_new_tokens, - end_id=eot_id, - pad_id=eot_id, - num_beams=args.num_beams, - num_return_sequences=args.num_beams, - output_sequence_lengths=True, - output_generation_logits=True, - return_dict=True, - ) - torch.cuda.synchronize() - - generation_logits = outputs["generation_logits"] - assert generation_logits.shape[1] == args.num_beams, ( - f"Expected generation_logits beam dimension to be {args.num_beams}, " - f"got {generation_logits.shape[1]}" - ) - - passed = validate_logits_alignment( - outputs["output_ids"].cpu(), generation_logits.cpu(), input_len, eot_id - ) - - if passed: - print("PASS: generation_logits aligned with output_ids") - else: - print("FAIL: generation_logits misaligned with output_ids") - sys.exit(1) - - -if __name__ == "__main__": - main() diff --git a/tests/integration/defs/examples/validate_whisper_log_probs_determinism.py b/tests/integration/defs/examples/validate_whisper_log_probs_determinism.py deleted file mode 100644 index 78d950880613..000000000000 --- a/tests/integration/defs/examples/validate_whisper_log_probs_determinism.py +++ /dev/null @@ -1,218 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 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. -"""Validate that log_probs are deterministic across runs with beam search and batch_size > 1. - -Reproduces the nMaxBatchSize stride bug: when batch items finish at different -times the active batch dimension shrinks, causing beamStage3 to read logProbsTiled -with the wrong stride and produce different log_probs values across runs. - -Different LibriSpeech samples produce different decoder output lengths, creating -an uneven batch where some items finish before others — the exact scenario that -triggers the bug. -""" - -import argparse -import json -import sys -from collections import OrderedDict -from pathlib import Path - -import numpy as np -import torch - -import tensorrt_llm -from tensorrt_llm._utils import str_dtype_to_torch -from tensorrt_llm.runtime import ModelRunnerCpp - - -def read_config(component, engine_dir): - config_path = engine_dir / component / "config.json" - with open(config_path, "r") as f: - config = json.load(f) - model_config = OrderedDict() - model_config.update(config["pretrained_config"]) - model_config.update(config["build_config"]) - return model_config - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument("--engine_dir", type=str, required=True) - parser.add_argument("--assets_dir", type=str, required=True) - parser.add_argument( - "--dataset_dir", - type=str, - required=True, - help="Local path to hf-internal-testing/librispeech_asr_dummy dataset", - ) - parser.add_argument("--num_beams", type=int, default=4) - parser.add_argument("--batch_size", type=int, default=4) - parser.add_argument("--max_new_tokens", type=int, default=96) - parser.add_argument("--num_runs", type=int, default=5) - args = parser.parse_args() - if args.num_runs < 2: - parser.error("--num_runs must be >= 2 to compare log_probs across runs") - return args - - -def main(): - args = parse_arguments() - tensorrt_llm.logger.set_level("warning") - - engine_dir = Path(args.engine_dir) - encoder_config = read_config("encoder", engine_dir) - decoder_config = read_config("decoder", engine_dir) - - n_mels = encoder_config["n_mels"] - is_multilingual = decoder_config["vocab_size"] >= 51865 # multilingual vocab size threshold - - from tokenizer import get_tokenizer - from whisper_utils import log_mel_spectrogram - - tokenizer_name = "multilingual" if is_multilingual else "gpt2" - tokenizer = get_tokenizer( - name=tokenizer_name, - num_languages=encoder_config["num_languages"], - tokenizer_dir=args.assets_dir, - ) - eot_id = tokenizer.encode("<|endoftext|>", allowed_special=tokenizer.special_tokens_set)[0] - - runner = ModelRunnerCpp.from_dir( - engine_dir=engine_dir, - is_enc_dec=True, - max_batch_size=args.batch_size, - max_input_len=3000, - max_output_len=args.max_new_tokens, - max_beam_width=args.num_beams, - kv_cache_free_gpu_memory_fraction=0.9, - cross_kv_cache_fraction=0.5, - gather_generation_logits=True, - ) - - # Use different LibriSpeech samples so each batch item produces a different - # number of output tokens — this creates the uneven-finish condition that - # triggers the nMaxBatchSize stride bug. - from datasets import load_dataset - - dataset = load_dataset(args.dataset_dir, "clean", split="validation", trust_remote_code=True) - - mel_list = [] - for i in range(args.batch_size): - speech = dataset[i]["audio"]["array"].astype(np.float32) - waveform = torch.from_numpy(speech) - m = log_mel_spectrogram(waveform, n_mels, device="cuda", mel_filters_dir=args.assets_dir) - m = m.type(str_dtype_to_torch("float16")) - if m.shape[1] % 2: - m = torch.nn.functional.pad(m, (0, 1)) - mel_list.append(m) - - max_mel_len = max(m.shape[1] for m in mel_list) - mel_batched = torch.zeros( - args.batch_size, - mel_list[0].shape[0], - max_mel_len, - dtype=mel_list[0].dtype, - device=mel_list[0].device, - ) - for i, m in enumerate(mel_list): - mel_batched[i, :, : m.shape[1]] = m - - mel_input_lengths = torch.full( - (args.batch_size,), max_mel_len, dtype=torch.int32, device="cuda" - ) - - prompt_text = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>" - prompt_ids = tokenizer.encode(prompt_text, allowed_special=tokenizer.special_tokens_set) - decoder_input_ids = torch.tensor(prompt_ids).unsqueeze(0).repeat(args.batch_size, 1) - - all_log_probs = [] - ref_output_ids = None - ref_gen_logits = None - input_len = decoder_input_ids.shape[1] - for run_idx in range(args.num_runs): - with torch.no_grad(): - outputs = runner.generate( - batch_input_ids=decoder_input_ids, - encoder_input_features=mel_batched.transpose(1, 2), - encoder_output_lengths=mel_input_lengths // 2, - max_new_tokens=args.max_new_tokens, - end_id=eot_id, - pad_id=eot_id, - num_beams=args.num_beams, - num_return_sequences=1, - output_sequence_lengths=True, - output_generation_logits=True, - output_log_probs=True, - return_dict=True, - ) - torch.cuda.synchronize() - all_log_probs.append(outputs["log_probs"][:, 0, :].cpu()) - if run_idx == 0: - ref_output_ids = outputs["output_ids"].cpu() - ref_gen_logits = outputs["generation_logits"].cpu() - - for i in range(1, args.num_runs): - if all_log_probs[i].shape != all_log_probs[0].shape: - print( - f"FAIL: log_probs shape mismatch between run 1 {all_log_probs[0].shape} " - f"and run {i + 1} {all_log_probs[i].shape}" - ) - sys.exit(1) - - max_diff = max( - (all_log_probs[0] - all_log_probs[i]).abs().max().item() for i in range(1, args.num_runs) - ) - - if max_diff >= 1e-6: - print(f"FAIL: log_probs are non-deterministic (max diff: {max_diff:.6f})") - sys.exit(1) - - # Correctness check: verify log_probs[b][t] matches log_softmax(generation_logits[b][t])[token] - # for every batch slot b (including b > 0, exercising the gatherTree batchSlot offset fix). - # Uses a loose tolerance because log_probs come from float32 beam search bookkeeping while - # generation_logits are the raw fp16->fp32 decoder outputs. - LOG_PROB_ATOL = 0.5 - all_aligned = True - for b in range(args.batch_size): - gen_tokens = ref_output_ids[b, 0, input_len:] - eot_pos = (gen_tokens == eot_id).nonzero(as_tuple=True)[0] - gen_len = eot_pos[0].item() if len(eot_pos) > 0 else gen_tokens.shape[0] - if gen_len == 0: - continue - logits = ref_gen_logits[b, 0, :gen_len, :] - log_probs_from_logits = torch.nn.functional.log_softmax(logits.float(), dim=-1) - expected = log_probs_from_logits.gather(1, gen_tokens[:gen_len].unsqueeze(1)).squeeze(1) - actual = all_log_probs[0][b, :gen_len] - max_lp_diff = (actual - expected).abs().max().item() - if max_lp_diff > LOG_PROB_ATOL: - print( - f"FAIL: log_probs[batch={b}] deviate from generation_logits " - f"(max diff: {max_lp_diff:.4f} > {LOG_PROB_ATOL}); " - "likely caused by wrong logProbsTiled batchSlot offset in gatherTree." - ) - all_aligned = False - - if not all_aligned: - sys.exit(1) - - print( - f"PASS: log_probs are deterministic across {args.num_runs} runs " - f"(max diff: {max_diff:.2e}) and aligned with generation_logits for all batch slots." - ) - sys.exit(0) - - -if __name__ == "__main__": - main() diff --git a/tests/integration/defs/perf/build.py b/tests/integration/defs/perf/build.py deleted file mode 100644 index a12169ba1ff6..000000000000 --- a/tests/integration/defs/perf/build.py +++ /dev/null @@ -1,1524 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 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. - -import argparse -import multiprocessing as mp -import os -import time -from collections import OrderedDict - -# isort: off -import torch -import tensorrt as trt -# isort: on - -from allowed_configs import (get_allowed_models, get_build_config, - get_model_config, get_model_family) - -import tensorrt_llm -from tensorrt_llm._utils import str_dtype_to_trt -from tensorrt_llm.builder import BuildConfig, Builder, build -from tensorrt_llm.functional import LayerNormPositionType, LayerNormType -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import PretrainedConfig -from tensorrt_llm.models.modeling_utils import QuantConfig, optimize_model -from tensorrt_llm.network import net_guard -from tensorrt_llm.plugin.plugin import ContextFMHAType -from tensorrt_llm.quantization import QuantAlgo, QuantMode -from tensorrt_llm.quantization.quantize import quantize - -WEIGHT_STREAMING_DISABLED_VAL = "1.0" - - -def parse_arguments(): - parser = argparse.ArgumentParser(description='Build TensorRT LLM models.') - parser.add_argument('-m', - '--model', - type=str, - required=True, - choices=get_allowed_models(), - help='Specify model you want to build.') - parser.add_argument( - '--mode', - type=str, - default="plugin", - choices=['ootb', 'plugin', 'plugin-ifb', 'ootb-except-mha'], - help= - ('Choose mode between ootb/plugin/ootb-except-mha. ' - '\"ootb\" means the engines will be built without any plugins, ' - '\"plugin\" means the engines will be built with tuned recipe of using plugins.' - '\"plugin-ifb\" will include additional options required for inflight batching.' - '\"ootb-except-mha\" means the engines will be built with only attention plugins.' - )) - - parser.add_argument( - '--dtype', - type=str, - default='float16', - choices=['float16', 'bfloat16', 'float32'], - help='Choose data type between float16/bfloat16/float32.') - parser.add_argument( - '--quantization', - type=str, - default=None, - choices=[ - 'fp8', 'fp8_gemm', 'fp8_kv_cache', 'int8_sq_per_tensor', - 'int8_sq_per_token_channel', 'int8_weight_only', 'int4_weight_only', - 'int4_weight_only_awq', 'int4_weight_only_gptq' - ], - help="Optimize the model with specified quantization recipe") - - parser.add_argument( - '--input_timing_cache', - type=str, - default=None, - help= - 'The path to read timing cache, will be ignored if the file does not exist' - ) - parser.add_argument('--output_timing_cache', - type=str, - default='model.cache', - help='The path to write timing cache') - - parser.add_argument( - '--profiling_verbosity', - type=str, - default='layer_names_only', - choices=['layer_names_only', 'detailed', 'none'], - help= - 'The profiling verbosity for the generated TRT engine. Set to detailed can inspect tactic choices and kernel parameters.' - ) - parser.add_argument( - '--log_level', - type=str, - default="error", - choices=['verbose', 'info', 'warning', 'error', 'internal_error'], - help= - 'Choose log level between verbose/info/warning/error/internal_error.') - - parser.add_argument( - '--output_dir', - type=str, - required=True, - help='TensorRT engines will be saved to the specified path.') - - parser.add_argument( - '--max_beam_width', - type=int, - default=None, - help= - ('If this option is specified, it will override the max beam width of ' - 'TRT engines to the specified value instead of using pre-defined one')) - parser.add_argument( - '--max_input_len', - type=int, - default=None, - help= - ('If this option is specified, it will override the max input len of ' - 'TRT engines to the specified value instead of using pre-defined one')) - parser.add_argument( - '--max_seq_len', - '--max_decoder_seq_len', - dest='max_seq_len', - type=int, - default=None, - help= - ('If this option is specified, it will override the max sequence len of ' - 'TRT engines to the specified value instead of using pre-defined one')) - parser.add_argument( - '--max_batch_size', - type=int, - default=None, - help= - ('If this option is specified, it will override the max batch size of ' - 'TRT engines to the specified value instead of using pre-defined one')) - parser.add_argument('--force_num_layer_1', - default=False, - action='store_true', - help='Quick sanity check with num_layer=1.') - parser.add_argument('--serial_build', - default=False, - action='store_true', - help="Build engines serially") - parser.add_argument( - '--multiple_profiles', - default=False, - action='store_true', - help= - 'This option will benefit performance, but will increase the engine build time.' - ) - - parser.add_argument( - '--weight_streaming', - default=False, - action='store_true', - help= - 'Specify whether offloading weights to CPU and streaming loading at runtime.', - ) - - parser.add_argument( - '--monitor_memory', - default=False, - action='store_true', - help='Specify whether turning on the memory monitor flag.', - ) - - parser.add_argument( - '--rank', - type=int, - default=None, - help= - "The rank of the model to be built, only used when --serial_build is specified" - ) - parser.add_argument( - '--world_size', - type=int, - default=None, - help= - "The number of gpus to be used for inference, only used when --serial_build is specified" - ) - parser.add_argument( - '--opt_batch_size', - type=int, - default=None, - help= - "If opt_batch_size option is specified, it will override the opt batch size." - "This flag only takes effect when `--mode=ootb` is added. For other modes, please use --opt_num_tokens to replace it." - ) - - parser.add_argument( - '--opt_num_tokens', - type=int, - default=None, - help="It equals to max_batch_size*max_beam_width by default, set this " - "value as close as possible to the actual number of tokens on your workload. " - "Note that this argument might be removed in the future." - "This flag only takes effect when `--mode` is not `ootb`. For ootb mode, please use --opt_batch_size to replace it." - ) - - return parser.parse_args() - - -def serialize_engine(engine, path): - logger.info(f'Serializing engine to {path}...') - tik = time.time() - with open(path, 'wb') as f: - # engine object is already complies with python buffer protocol, no need to - # convert it to bytearray before write, converting to bytearray consumes lots of memory - f.write(engine) - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Engine serialized. Total time: {t}') - - -def get_quant_config(quantization: str): - if quantization == "fp8": - return QuantConfig(quant_algo=QuantAlgo.FP8, - kv_cache_quant_algo=QuantAlgo.FP8) - elif quantization == "fp8_gemm": - return QuantConfig(quant_algo=QuantAlgo.FP8) - elif quantization == "fp8_kv_cache": - return QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) - elif quantization == "int8_sq_per_tensor": - return QuantConfig(quant_algo=QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN) - elif quantization == "int8_sq_per_token_channel": - return QuantConfig( - quant_algo=QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN) - elif quantization == "int8_sq_per_channel_ootb": - return QuantConfig(quant_algo=QuantAlgo.W8A8_SQ_PER_CHANNEL) - elif quantization == "int8_weight_only": - return QuantConfig(quant_algo=QuantAlgo.W8A16) - elif quantization == "int4_weight_only": - return QuantConfig(quant_algo=QuantAlgo.W4A16) - elif quantization == "int4_weight_only_awq": - return QuantConfig(quant_algo=QuantAlgo.W4A16_AWQ) - elif quantization == "int4_weight_only_gptq": - return QuantConfig(quant_algo=QuantAlgo.W4A16_GPTQ) - elif quantization is None: - return QuantConfig() - else: - raise Exception(f"Unexpected quantization: {quantization}") - - -def build_gpt(args): - build_config = get_build_config(args.model) - build_config = BuildConfig(**build_config) - model_config = get_model_config(args.model) - if args.force_num_layer_1: - model_config['num_layers'] = 1 - - # More parameters - if args.serial_build and args.rank is not None and args.world_size is not None: - runtime_rank = args.rank - world_size = args.world_size - else: - runtime_rank = tensorrt_llm.mpi_rank() - world_size = tensorrt_llm.mpi_world_size() - if not args.serial_build: - torch.cuda.set_device(runtime_rank) - - if args.profiling_verbosity != "layer_names_only": - build_config.profiling_verbosity = args.profiling_verbosity - - if args.max_batch_size is not None: - build_config.max_batch_size = args.max_batch_size - if args.max_input_len is not None: - build_config.max_input_len = args.max_input_len - if args.max_seq_len is not None: - build_config.max_seq_len = args.max_seq_len - if args.max_beam_width is not None: - build_config.max_beam_width = args.max_beam_width - if args.opt_batch_size is not None: - build_config.opt_batch_size = args.opt_batch_size - if args.opt_num_tokens is not None: - build_config.opt_num_tokens = args.opt_num_tokens - build_config.weight_streaming = getattr(args, "weight_streaming", False) - build_config.max_num_tokens = build_config.max_batch_size * max( - build_config.max_input_len, build_config.max_beam_width) - - if args.mode != "ootb" and args.opt_batch_size is not None: - raise Exception( - f'--opt_batch_size only used when mode is ootb. Please using --opt_num_tokens instead it.' - ) - if args.mode == "ootb" and args.opt_num_tokens is not None: - raise Exception( - f'--opt_num_tokens does not support ootb mode. Please using --opt_batch_size instead it.' - ) - - quant_config = get_quant_config(args.quantization) - quant_algo = quant_config.quant_algo - kv_cache_quant_algo = quant_config.kv_cache_quant_algo - quant_mode = quant_config.quant_mode - - # Initialize Module - family = get_model_family(args.model) - if family == "gpt": - if model_config['num_kv_heads'] is None: - model_config['num_kv_heads'] = model_config['num_heads'] - if model_config['inter_size'] is None: - model_config['inter_size'] = model_config['hidden_size'] * 4 - if model_config['position_embedding_type'] is None: - model_config['position_embedding_type'] = 'learned_absolute' - - config = { - 'architecture': 'GPTForCausalLM', - 'dtype': args.dtype, - 'num_hidden_layers': model_config['num_layers'], - 'num_attention_heads': model_config['num_heads'], - 'num_key_value_heads': model_config['num_kv_heads'], - 'hidden_size': model_config['hidden_size'], - 'intermediate_size': model_config['inter_size'], - 'norm_epsilon': 1e-05, - 'vocab_size': model_config['vocab_size'], - 'position_embedding_type': model_config['position_embedding_type'], - 'max_position_embeddings': model_config['n_positions'], - 'hidden_act': model_config['hidden_act'], - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo, - 'group_size': 128, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': world_size, - 'rank': runtime_rank - }, - 'bias': model_config['bias'], - 'apply_query_key_layer_scaling': False, - 'rotary_pct': model_config['rotary_pct'], - 'moe': { - 'num_experts': model_config["moe_num_experts"], - 'top_k': model_config["moe_top_k"], - }, - } - config = PretrainedConfig.from_dict(config) - tensorrt_llm_model = tensorrt_llm.models.GPTForCausalLM(config) - - elif family == "opt": - config = { - 'architecture': 'OPTForCausalLM', - 'dtype': args.dtype, - 'vocab_size': model_config['vocab_size'], - 'hidden_size': model_config['hidden_size'], - 'num_hidden_layers': model_config['num_layers'], - 'num_attention_heads': model_config['num_heads'], - 'hidden_act': model_config['hidden_act'], - 'max_position_embeddings': model_config['n_positions'], - 'mapping': { - 'world_size': world_size, - 'tp_size': world_size, - 'rank': runtime_rank - }, - 'use_parallel_embedding': False, - 'embedding_sharding_dim': 0, - 'do_layer_norm_before': model_config['do_layer_norm_before'], - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo, - 'group_size': 128 - } - } - config = PretrainedConfig.from_dict(config) - tensorrt_llm_model = tensorrt_llm.models.OPTForCausalLM(config) - - elif family == "llama": - config = { - 'architecture': - 'LlamaForCausalLM', - 'dtype': - args.dtype, - 'logits_dtype': - 'float32', - 'num_hidden_layers': - model_config['num_layers'], - 'num_attention_heads': - model_config['num_heads'], - 'num_key_value_heads': - model_config['num_heads'] if model_config['num_kv_heads'] is None - else model_config['num_kv_heads'], - 'hidden_size': - model_config['hidden_size'], - 'intermediate_size': - model_config['inter_size'], - 'vocab_size': - model_config['vocab_size'], - 'position_embedding_type': - 'rope_gpt_neox', - 'max_position_embeddings': - model_config['n_positions'], - 'hidden_act': - model_config['hidden_act'], - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo, - 'group_size': 128 - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': world_size, - 'moe_tp_size': world_size, - 'moe_ep_size': 1, - 'rank': runtime_rank - }, - 'moe': { - 'num_experts': model_config["moe_num_experts"], - 'top_k': model_config["moe_top_k"], - } - } - config = PretrainedConfig.from_dict(config) - tensorrt_llm_model = tensorrt_llm.models.LLaMAForCausalLM(config) - tensorrt_llm_model = optimize_model(tensorrt_llm_model, - use_fused_mlp=True) - elif family == "gptj": - config = { - 'architecture': 'GPTJForCausalLM', - 'dtype': args.dtype, - 'vocab_size': model_config['vocab_size'], - 'hidden_size': model_config['hidden_size'], - 'num_hidden_layers': model_config['num_layers'], - 'num_attention_heads': model_config['num_heads'], - 'hidden_act': model_config['hidden_act'], - 'max_position_embeddings': model_config['n_positions'], - 'position_embedding_type': 'rope_gptj', - 'rotary_dim': model_config['rotary_dim'], - 'mapping': { - 'world_size': world_size, - 'tp_size': world_size, - 'rank': runtime_rank - }, - 'use_parallel_embedding': False, - 'embedding_sharding_dim': 0, - 'do_layer_norm_before': model_config['do_layer_norm_before'], - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo, - 'group_size': 128 - } - } - config = PretrainedConfig.from_dict(config) - tensorrt_llm_model = tensorrt_llm.models.GPTJForCausalLM(config) - - elif family == "gptneox": - config = { - 'architecture': - 'GPTNeoXForCausalLM', - 'dtype': - args.dtype, - 'num_hidden_layers': - model_config['num_layers'], - 'num_attention_heads': - model_config['num_heads'], - 'hidden_size': - model_config['hidden_size'], - 'vocab_size': - model_config['vocab_size'], - 'position_embedding_type': - 'rope_gpt_neox', - 'max_position_embeddings': - model_config['n_positions'], - 'rotary_emb_base': - 10000, - 'rotary_pct': - 1.0 * model_config['rotary_dim'] * model_config['num_heads'] / - model_config['hidden_size'], - 'hidden_act': - model_config['hidden_act'], - 'mapping': { - 'world_size': world_size, - 'tp_size': world_size, - 'rank': runtime_rank - }, - 'use_parallel_embedding': - False, - 'embedding_sharding_dim': - 0, - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo, - 'group_size': 128, - } - } - config = PretrainedConfig.from_dict(config) - tensorrt_llm_model = tensorrt_llm.models.GPTNeoXForCausalLM(config) - - elif family == "chatglm": - config = { - 'architecture': 'ChatGLMModel', - 'dtype': args.dtype, - 'num_hidden_layers': model_config['num_layers'], - 'num_attention_heads': model_config['num_heads'], - 'num_key_value_heads': model_config['num_kv_heads'], - 'hidden_size': model_config['hidden_size'], - 'intermediate_size': model_config['inter_size'], - 'norm_epsilon': 1e-5, - 'vocab_size': model_config['vocab_size'], - 'position_embedding_type': 'chatglm', - 'max_position_embeddings': model_config['n_positions'], - 'hidden_act': model_config['hidden_act'], - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': world_size, - 'rank': runtime_rank - }, - 'chatglm_version': 'chatglm', - 'add_bias_linear': True, - 'add_qkv_bias': True, - 'apply_query_key_layer_scaling': False, - 'apply_residual_connection_post_layernorm': False, - 'rmsnorm': False, - } - config = PretrainedConfig.from_dict(config) - tensorrt_llm_model = tensorrt_llm.models.ChatGLMForCausalLM(config) - - elif family in ["chatglm2", "chatglm3"]: - config = { - 'architecture': 'ChatGLMModel', - 'dtype': args.dtype, - 'num_hidden_layers': model_config['num_layers'], - 'num_attention_heads': model_config['num_heads'], - 'num_key_value_heads': model_config['num_kv_heads'], - 'hidden_size': model_config['hidden_size'], - 'intermediate_size': model_config['inter_size'], - 'norm_epsilon': 1e-5, - 'vocab_size': model_config['vocab_size'], - 'position_embedding_type': 'rope_gptj', - 'max_position_embeddings': model_config['n_positions'], - 'hidden_act': model_config['hidden_act'], - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': world_size, - 'rank': runtime_rank - }, - 'chatglm_version': family, - 'add_bias_linear': False, - 'add_qkv_bias': True, - 'apply_query_key_layer_scaling': False, - 'apply_residual_connection_post_layernorm': False, - 'rmsnorm': True, - } - config = PretrainedConfig.from_dict(config) - tensorrt_llm_model = tensorrt_llm.models.ChatGLMForCausalLM(config) - - elif family == "glm": - config = { - 'architecture': 'GLMModel', - 'dtype': args.dtype, - 'num_hidden_layers': model_config['num_layers'], - 'num_attention_heads': model_config['num_heads'], - 'num_key_value_heads': model_config['num_kv_heads'], - 'hidden_size': model_config['hidden_size'], - 'intermediate_size': model_config['inter_size'], - 'norm_epsilon': 1e-5, - 'vocab_size': model_config['vocab_size'], - 'position_embedding_type': 'learned_absolute', - 'max_position_embeddings': model_config['n_positions'], - 'hidden_act': model_config['hidden_act'], - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': world_size, - 'rank': runtime_rank - }, - 'chatglm_version': 'glm', - 'add_bias_linear': True, - 'add_qkv_bias': True, - 'apply_query_key_layer_scaling': False, - 'apply_residual_connection_post_layernorm': False, - 'rmsnorm': False, - } - config = PretrainedConfig.from_dict(config) - tensorrt_llm_model = tensorrt_llm.models.ChatGLMForCausalLM(config) - - elif family == "bloom": - config = { - 'architecture': 'BloomForCausalLM', - 'dtype': args.dtype, - 'vocab_size': model_config['vocab_size'], - 'hidden_size': model_config['hidden_size'], - 'num_hidden_layers': model_config['num_layers'], - 'num_attention_heads': model_config['num_heads'], - 'hidden_act': model_config['hidden_act'], - 'max_position_embeddings': model_config['n_positions'], - 'mapping': { - 'world_size': world_size, - 'tp_size': world_size, - 'rank': runtime_rank - }, - 'use_parallel_embedding': (args.model == 'bloom_176b'), - 'embedding_sharding_dim': 0, - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo, - 'group_size': 128 - } - } - config = PretrainedConfig.from_dict(config) - tensorrt_llm_model = tensorrt_llm.models.BloomForCausalLM(config) - elif family == "falcon": - config = { - 'architecture': - 'FalconForCausalLM', - 'dtype': - args.dtype, - 'num_hidden_layers': - model_config['num_layers'], - 'num_attention_heads': - model_config['num_heads'], - 'num_key_value_heads': - model_config['num_heads'] if model_config['num_kv_heads'] is None - else model_config['num_kv_heads'], - 'hidden_size': - model_config['hidden_size'], - 'vocab_size': - model_config['vocab_size'], - 'position_embedding_type': - 'alibi_with_scale' - if model_config['use_alibi'] else 'rope_gpt_neox', - 'max_position_embeddings': - model_config['n_positions'], - 'hidden_act': - model_config['hidden_act'], - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo, - 'group_size': 128 - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': world_size, - 'rank': runtime_rank - }, - 'bias': - model_config['bias'], - 'parallel_attention': - model_config['parallel_attention'], - 'new_decoder_architecture': - model_config['new_decoder_architecture'], - } - if quant_mode.is_weight_only() and quant_mode.has_per_group_scaling(): - config['quantization'].update({ - 'has_zero_point': False, - 'pre_quant_scale': True, - }) - config = PretrainedConfig.from_dict(config) - tensorrt_llm_model = tensorrt_llm.models.FalconForCausalLM(config) - - elif family == "baichuan": - config = { - 'architecture': - 'BaichuanForCausalLM', - 'dtype': - args.dtype, - 'logits_dtype': - 'float32', - 'vocab_size': - model_config['vocab_size'], - 'max_position_embeddings': - model_config['n_positions'], - 'hidden_size': - model_config['hidden_size'], - 'num_hidden_layers': - model_config['num_layers'], - 'num_attention_heads': - model_config['num_heads'], - 'num_key_value_heads': - model_config['num_heads'], - 'hidden_act': - model_config['hidden_act'], - 'intermediate_size': - model_config['inter_size'], - 'position_embedding_type': - 'alibi_with_scale' if '7b' in args.model else 'rope_gpt_neox', - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo, - 'group_size': 128 - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': world_size, - 'rank': runtime_rank - }, - } - config = PretrainedConfig.from_dict(config) - tensorrt_llm_model = tensorrt_llm.models.BaichuanForCausalLM(config) - - elif family == "internlm": - config = { - 'architecture': - 'LlamaForCausalLM', - 'dtype': - args.dtype, - 'num_hidden_layers': - model_config['num_layers'], - 'num_attention_heads': - model_config['num_heads'], - 'num_key_value_heads': - model_config['num_heads'] if model_config['num_kv_heads'] is None - else model_config['num_kv_heads'], - 'hidden_size': - model_config['hidden_size'], - 'intermediate_size': - model_config['inter_size'], - 'vocab_size': - model_config['vocab_size'], - 'position_embedding_type': - 'rope_gpt_neox', - 'max_position_embeddings': - model_config['n_positions'], - 'hidden_act': - model_config['hidden_act'], - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': world_size, - 'rank': runtime_rank - }, - 'attn_bias': - model_config['bias'], - } - if quant_mode.is_weight_only(): - if 'awq' in args.quantization: - config['quantization'].update({ - "group_size": 128, - "has_zero_point": False, - "pre_quant_scale": True, - }) - elif 'gptq' in args.quantization: - config['quantization'].update({ - "group_size": 128, - "has_zero_point": True, - "pre_quant_scale": False, - }) - config = PretrainedConfig.from_dict(config) - tensorrt_llm_model = tensorrt_llm.models.LLaMAForCausalLM(config) - - elif family == "qwen": - config = { - 'architecture': - 'QWenForCausalLM', - 'dtype': - args.dtype, - 'num_hidden_layers': - model_config['num_layers'], - 'num_attention_heads': - model_config['num_heads'], - 'num_key_value_heads': - model_config['num_heads'] if model_config['num_kv_heads'] is None - else model_config['num_kv_heads'], - 'seq_length': - model_config['n_positions'], - 'hidden_size': - model_config['hidden_size'], - 'intermediate_size': - model_config['inter_size'], - 'vocab_size': - model_config['vocab_size'], - 'position_embedding_type': - 'rope_gpt_neox', - 'max_position_embeddings': - model_config['n_positions'], - 'hidden_act': - model_config['hidden_act'], - 'quantization': { - 'group_size': 128, - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': world_size, - 'rank': runtime_rank - }, - 'moe': { - 'num_experts': model_config["moe_num_experts"], - 'top_k': model_config["moe_top_k"], - }, - 'qwen_type': - 'qwen', - } - config = PretrainedConfig.from_dict(config) - tensorrt_llm_model = tensorrt_llm.models.QWenForCausalLM(config) - elif family == "qwen2": - config = { - 'architecture': - 'QWenForCausalLM', - 'dtype': - args.dtype, - 'num_hidden_layers': - model_config['num_layers'], - 'num_attention_heads': - model_config['num_heads'], - 'num_key_value_heads': - model_config['num_heads'] if model_config['num_kv_heads'] is None - else model_config['num_kv_heads'], - 'seq_length': - model_config['n_positions'], - 'hidden_size': - model_config['hidden_size'], - 'intermediate_size': - model_config['inter_size'], - 'vocab_size': - model_config['vocab_size'], - 'position_embedding_type': - 'rope_gpt_neox', - 'max_position_embeddings': - model_config['n_positions'], - 'hidden_act': - model_config['hidden_act'], - 'quantization': { - 'group_size': 128, - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': world_size, - 'rank': runtime_rank - }, - 'moe': { - 'num_experts': model_config["moe_num_experts"], - 'top_k': model_config["moe_top_k"], - }, - 'qwen_type': - 'qwen2', - } - config = PretrainedConfig.from_dict(config) - tensorrt_llm_model = tensorrt_llm.models.QWenForCausalLM(config) - elif family == "mamba": - config = { - 'architecture': 'MambaForCausalLM', - 'dtype': args.dtype, - 'vocab_size': model_config['vocab_size'], - 'hidden_size': model_config['hidden_size'], - 'num_hidden_layers': model_config['num_layers'], - 'num_attention_heads': model_config['num_heads'], - 'hidden_act': model_config['hidden_act'], - 'state_size': model_config['state_size'], - 'conv_kernel': model_config['conv_kernel'], - 'layer_types': model_config['layer_types'], - 'rnn_hidden_size': model_config['rnn_hidden_size'], - 'rnn_head_size': model_config['rnn_head_size'], - 'rnn_conv_dim_size': model_config['rnn_conv_dim_size'], - 'rms_norm': True, - 'residual_in_fp32': True, - 'pad_vocab_size_multiple': 8, - 'use_bias': model_config['use_bias'], - 'mamba_version': model_config['mamba_version'], - 'ssm_rmsnorm': model_config['ssm_rmsnorm'], - 'ngroups': model_config['ngroups'], - 'chunk_size': model_config['chunk_size'], - } - config = PretrainedConfig.from_dict(config) - tensorrt_llm_model = tensorrt_llm.models.MambaForCausalLM(config) - elif family == "recurrentgemma": - config = { - 'architecture': 'RecurrentGemmaForCausalLM', - 'dtype': args.dtype, - 'vocab_size': model_config['vocab_size'], - 'hidden_size': model_config['hidden_size'], - 'num_hidden_layers': model_config['num_layers'], - 'num_attention_heads': model_config['num_heads'], - 'num_key_value_heads': model_config['num_kv_heads'], - 'hidden_act': model_config['hidden_act'], - 'intermediate_size': model_config['inter_size'], - 'rms_norm': True, - 'norm_epsilon': 1e-6, - 'quantization': { - 'group_size': 128, - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': world_size, - 'rank': runtime_rank - }, - 'position_embedding_type': model_config['position_embedding_type'], - 'rotary_percentage': model_config['rotary_pct'], - 'max_position_embeddings': model_config['n_positions'], - 'conv_kernel': model_config['conv_kernel'], - 'state_size': model_config['state_size'], - 'layer_types': model_config['layer_types'], - 'rnn_hidden_size': model_config['rnn_hidden_size'], - 'rnn_head_size': model_config['rnn_head_size'], - 'rnn_conv_dim_size': model_config['rnn_conv_dim_size'], - 'logits_soft_cap': model_config['logits_soft_cap'], - 'rotary_pct': model_config['rotary_pct'], - } - config = PretrainedConfig.from_dict(config) - tensorrt_llm_model = tensorrt_llm.models.RecurrentGemmaForCausalLM( - config) - tensorrt_llm_model = optimize_model(tensorrt_llm_model, - use_fused_mlp=True, - use_fused_rg_lru=True) - elif family == "phi3": - config = { - 'architecture': - 'PhiForCausalLM', - 'dtype': - args.dtype, - 'rotary_base': - 10000.0, - 'num_hidden_layers': - model_config['num_layers'], - 'num_attention_heads': - model_config['num_heads'], - 'num_key_value_heads': - model_config['num_heads'] if model_config['num_kv_heads'] is None - else model_config['num_kv_heads'], - 'hidden_size': - model_config['hidden_size'], - 'intermediate_size': - model_config['inter_size'], - 'vocab_size': - model_config['vocab_size'], - 'position_embedding_type': - 'rope_gpt_neox', - 'max_position_embeddings': - model_config['n_positions'], - 'hidden_act': - model_config['hidden_act'], - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo, - 'group_size': 128 - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': world_size, - 'rank': runtime_rank - }, - } - config = PretrainedConfig.from_dict(config) - tensorrt_llm_model = tensorrt_llm.models.Phi3ForCausalLM(config) - - else: - raise Exception(f'Unexpected model: {args.model}') - - # Plugins - build_config.plugin_config.to_legacy_setting() - if args.mode in ['plugin', 'plugin-ifb']: - build_config.plugin_config.gpt_attention_plugin = args.dtype - build_config.plugin_config.set_context_fmha(ContextFMHAType.enabled) - build_config.plugin_config.remove_input_padding = True - build_config.plugin_config.moe_plugin = args.dtype - build_config.plugin_config.mamba_conv1d_plugin = args.dtype - - if args.quantization is None or "fp8" not in args.quantization: - build_config.plugin_config.gemm_plugin = args.dtype - - # Quantization plugins. - use_weight_only = quant_mode.is_weight_only() - if use_weight_only: - build_config.plugin_config.weight_only_quant_matmul_plugin = args.dtype - - use_smooth_quant = quant_mode.has_act_and_weight_quant() - if use_smooth_quant: - build_config.plugin_config.set_smooth_quant_plugins( - dtype=args.dtype) - - use_qserve = quant_mode.has_act_and_weight_quant() and quant_mode._any( - QuantMode.INT4_WEIGHTS) - if use_qserve: - build_config.plugin_config.set_qserve_plugins(dtype=args.dtype) - - # Inflight batching - if args.mode == 'plugin-ifb': - build_config.plugin_config.enable_paged_kv_cache() - build_config.plugin_config.paged_state = True - elif args.mode == 'ootb-except-mha': - build_config.plugin_config.gpt_attention_plugin = args.dtype - build_config.plugin_config.set_context_fmha(ContextFMHAType.enabled) - build_config.plugin_config.remove_input_padding = True - - if args.mode not in ('plugin', 'plugin-ifb'): - build_config.plugin_config.smooth_quant_plugins = False - - if world_size > 1: - build_config.plugin_config.set_nccl_plugin(dtype=args.dtype) - - if args.multiple_profiles: - build_config.plugin_config.multiple_profiles = True - - # Enable trt monitor memory for perf tests - build_config.monitor_memory = args.monitor_memory - - start = time.time() - engine = build(tensorrt_llm_model, build_config) - assert engine.engine is not None, f'Failed to build engine for rank {runtime_rank}' - build_time = round(time.time() - start, 2) - - engine.save(args.output_dir) - - return engine, build_time - - -def build_bert(args): - family = get_model_family(args.model) - build_config = get_build_config(args.model, return_dict=False) - model_config = get_model_config(args.model) - if args.force_num_layer_1: - model_config['num_layers'] = 1 - - # More parameters - if args.serial_build and args.rank is not None and args.world_size is not None: - runtime_rank = args.rank - world_size = args.world_size - else: - runtime_rank = tensorrt_llm.mpi_rank() - world_size = tensorrt_llm.mpi_world_size() - if not args.serial_build: - torch.cuda.set_device(runtime_rank) - - num_kv_heads = model_config['num_heads'] \ - if model_config['num_kv_heads'] is None else model_config['num_kv_heads'] - max_batch_size = build_config.max_batch_size \ - if args.max_batch_size is None else args.max_batch_size - max_input_len = build_config.max_input_len \ - if args.max_input_len is None else args.max_input_len - bs_range = [1, (max_batch_size + 1) // 2, max_batch_size] - inlen_range = [1, (max_input_len + 1) // 2, max_input_len] - - is_weight_streaming = getattr(args, "weight_streaming", False) - - builder = Builder() - builder_config = builder.create_builder_config( - name=args.model, - precision=args.dtype, - timing_cache=args.input_timing_cache, - profiling_verbosity=args.profiling_verbosity, - tensor_parallel=world_size, # TP only - parallel_build=True, - num_layers=model_config['num_layers'], - num_heads=model_config['num_heads'], - num_kv_heads=num_kv_heads, - hidden_size=model_config['hidden_size'], - vocab_size=model_config['vocab_size'], - hidden_act=model_config['hidden_act'], - max_position_embeddings=model_config['n_positions'], - max_batch_size=max_batch_size, - max_input_len=max_input_len, - strongly_typed=True, - weight_streaming=is_weight_streaming, - monitor_memory=args.monitor_memory, - ) - engine_name = '{}_{}_tp{}_rank{}.engine'.format(args.model, args.dtype, - world_size, runtime_rank) - - # Initialize model - config = { - 'architecture': 'BertModel', - 'dtype': args.dtype, - 'num_hidden_layers': model_config['num_layers'], - 'num_attention_heads': model_config['num_heads'], - 'hidden_size': model_config['hidden_size'], - 'vocab_size': model_config['vocab_size'], - 'position_embedding_type': 'learned_absolute', - 'max_position_embeddings': model_config['n_positions'], - 'hidden_act': model_config['hidden_act'], - 'type_vocab_size': model_config['type_vocab_size'], - 'pad_token_id': - None if family == 'bert' else 1, # hard code for RoBERTa here - 'is_roberta': (family == 'roberta'), - 'mapping': { - 'world_size': world_size, - 'tp_size': world_size, - 'rank': runtime_rank - }, - } - config = PretrainedConfig.from_dict(config) - tensorrt_llm_bert = tensorrt_llm.models.BertModel(config) - - # Module -> Network - network = builder.create_network() - network.trt_network.name = engine_name - network.plugin_config.to_legacy_setting() - - # Plugins - if args.mode == 'plugin': - network.plugin_config.bert_attention_plugin = args.dtype - network.plugin_config.gemm_plugin = args.dtype - network.plugin_config.set_context_fmha(ContextFMHAType.enabled) - elif args.mode == 'ootb-except-mha': - network.plugin_config.bert_attention_plugin = args.dtype - network.plugin_config.set_context_fmha(ContextFMHAType.enabled) - - if world_size > 1: - network.plugin_config.set_nccl_plugin(dtype=args.dtype) - - with net_guard(network): - # Prepare - network.set_named_parameters(tensorrt_llm_bert.named_parameters()) - - # Forward - input_ids = tensorrt_llm.Tensor( - name='input_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([('batch_size', [bs_range]), - ('input_len', [inlen_range])]), - ) - input_lengths = tensorrt_llm.Tensor(name='input_lengths', - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([ - ('batch_size', [bs_range]) - ])) - hidden_states = tensorrt_llm_bert(input_ids=input_ids, - input_lengths=input_lengths) - - # Mark outputs - hidden_states_dtype = str_dtype_to_trt(args.dtype) - hidden_states.mark_output('hidden_states', hidden_states_dtype) - - # Network -> Engine - start = time.time() - engine = builder.build_engine(network, builder_config) - assert engine is not None, f'Failed to build engine for rank {runtime_rank}' - build_time = round(time.time() - start, 2) - - if args.output_dir is not None: - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - serialize_path = os.path.join(args.output_dir, engine_name) - serialize_engine(engine, serialize_path) - if runtime_rank == 0: - config_path = os.path.join(args.output_dir, 'config.json') - builder_config.plugin_config = network.plugin_config - builder.save_config(builder_config, config_path) - if args.output_timing_cache: - # Save timing cache to output_dir if not absolute path - timing_cache_path = args.output_timing_cache if os.path.isabs( - args.output_timing_cache) else os.path.join( - args.output_dir, args.output_timing_cache) - ok = builder.save_timing_cache(builder_config, - timing_cache_path) - if not ok: - logger.warning("Failed to save timing cache.") - - return engine, build_time - - -def enc_dec_build_helper(component, build_config, model_config, args): - # More parameters - if args.serial_build and args.rank is not None and args.world_size is not None: - runtime_rank = args.rank - world_size = args.world_size - else: - runtime_rank = tensorrt_llm.mpi_rank() - world_size = tensorrt_llm.mpi_world_size() - if not args.serial_build: - torch.cuda.set_device(runtime_rank) - - family = get_model_family(args.model) - logits_dtype = 'float32' - if family == 'bart': - q_scaling = 1.0 - has_attention_qkvo_bias = True - has_mlp_bias = True - has_model_final_layernorm = False - has_position_embedding = True - has_embedding_layernorm = True - layernorm_type = LayerNormType.LayerNorm - relative_attention = False - layernorm_position = LayerNormPositionType.pre_layernorm if model_config.get( - 'normalize_before', True) else LayerNormPositionType.post_layernorm - rescale_before_lm_head = False - elif family == 'whisper': - q_scaling = 1.0 - has_position_embedding = True - relative_attention = False - has_embedding_layernorm = False - has_attention_qkvo_bias = True - has_mlp_bias = True - has_model_final_layernorm = True - layernorm_position = LayerNormPositionType.pre_layernorm - layernorm_type = LayerNormType.LayerNorm - rescale_before_lm_head = False - logits_dtype = args.dtype - model_config['n_mels'] - else: - q_scaling = 1 / model_config['head_size']**.5 - has_attention_qkvo_bias = False - has_mlp_bias = False - has_model_final_layernorm = True - has_position_embedding = False - has_embedding_layernorm = False - layernorm_type = LayerNormType.RmsNorm - relative_attention = True - layernorm_position = LayerNormPositionType.pre_layernorm - if family == 't5': - rescale_before_lm_head = True - else: - rescale_before_lm_head = False - - quant_config = get_quant_config(args.quantization) - quant_mode = quant_config.quant_mode - use_weight_only = quant_mode.is_weight_only() - - # Plugins - build_config.plugin_config.to_legacy_setting() - if args.mode in ['plugin', 'plugin-ifb']: - build_config.plugin_config.bert_attention_plugin = args.dtype - build_config.plugin_config.gpt_attention_plugin = args.dtype - build_config.plugin_config.set_context_fmha(ContextFMHAType.enabled) - build_config.plugin_config.gemm_plugin = args.dtype - build_config.plugin_config.remove_input_padding = True - build_config.plugin_config.enable_paged_kv_cache() - build_config.plugin_config.paged_state = True - if use_weight_only: - build_config.plugin_config.weight_only_quant_matmul_plugin = args.dtype - elif args.mode == 'ootb-except-mha': - build_config.plugin_config.bert_attention_plugin = args.dtype - build_config.plugin_config.gpt_attention_plugin = args.dtype - build_config.plugin_config.set_context_fmha(ContextFMHAType.enabled) - - if world_size > 1: - build_config.plugin_config.set_nccl_plugin(dtype=args.dtype) - - # build engine - mapping = Mapping(world_size=world_size, - rank=runtime_rank, - tp_size=world_size, - pp_size=1) # TP only - - if component == 'encoder': - if family == 'whisper': - pretrained_config = PretrainedConfig.from_dict({ - 'architecture': - "WhisperEncoder", - 'dtype': - args.dtype, - 'num_hidden_layers': - model_config['num_layers'], - 'num_attention_heads': - model_config['num_heads'], - 'hidden_size': - model_config['hidden_size'], - 'has_position_embedding': - has_position_embedding, - 'n_mels': - model_config['n_mels'], - 'max_position_embeddings': - 1500, - 'vocab_size': - model_config['vocab_size'], - 'hidden_act': - "gelu", - 'num_languages': - 100, - 'mapping': { - 'world_size': mapping.world_size, - 'tp_size': mapping.tp_size, - 'pp_size': mapping.pp_size, - 'rank': mapping.rank, - }, - }) - tllm_model = tensorrt_llm.models.WhisperEncoder(pretrained_config) - if use_weight_only: - tllm_model = quantize(tllm_model, quant_config) - else: - pretrained_config = PretrainedConfig.from_dict({ - 'architecture': - "EncoderModel", - 'dtype': - args.dtype, - 'logits_dtype': - logits_dtype, - 'num_hidden_layers': - model_config['num_layers'], - 'num_attention_heads': - model_config['num_heads'], - 'hidden_size': - model_config['hidden_size'], - 'norm_epsilon': - 1e-6, - 'vocab_size': - model_config['vocab_size'], - 'hidden_act': - model_config['hidden_act'], - 'mapping': { - 'world_size': mapping.world_size, - 'tp_size': mapping.tp_size, - 'pp_size': mapping.pp_size, - 'rank': mapping.rank, - }, - 'use_parallel_embedding': - False, - 'embedding_sharding_dim': - 0, - 'max_position_embeddings': - model_config.get('n_positions', 0), - 'use_prompt_tuning': - False, - 'head_size': - model_config['head_size'], - 'has_position_embedding': - has_position_embedding, - 'layernorm_type': - layernorm_type, - 'has_attention_qkvo_bias': - has_attention_qkvo_bias, - 'has_mlp_bias': - has_mlp_bias, - 'has_model_final_layernorm': - has_model_final_layernorm, - 'has_embedding_layernorm': - has_embedding_layernorm, - 'has_embedding_scale': - model_config.get('has_embedding_scale', False), - 'intermediate_size': - model_config['ffn_hidden_size'], - 'q_scaling': - q_scaling, - 'layernorm_position': - layernorm_position, - 'relative_attention': - relative_attention, - 'max_distance': - model_config.get('max_distance', 0), - 'num_buckets': - model_config.get('num_buckets', 0), - 'model_type': - family, - }) - tllm_model = tensorrt_llm.models.EncoderModel(pretrained_config) - elif component == 'decoder': - pretrained_config = PretrainedConfig.from_dict({ - 'architecture': - "DecoderModel", - 'dtype': - args.dtype, - 'logits_dtype': - logits_dtype, - 'num_hidden_layers': - model_config['num_layers'], - 'num_attention_heads': - model_config['num_heads'], - 'hidden_size': - model_config['hidden_size'], - 'norm_epsilon': - 1e-6, - 'vocab_size': - model_config['vocab_size'], - 'hidden_act': - model_config['hidden_act'], - 'mapping': { - 'world_size': mapping.world_size, - 'tp_size': mapping.tp_size, - 'pp_size': mapping.pp_size, - 'rank': mapping.rank, - }, - 'use_parallel_embedding': - False, - 'embedding_sharding_dim': - 0, - 'max_position_embeddings': - model_config.get('n_positions', 0), - 'use_prompt_tuning': - False, - 'head_size': - model_config['head_size'], - 'has_position_embedding': - has_position_embedding, - 'layernorm_type': - layernorm_type, - 'has_attention_qkvo_bias': - has_attention_qkvo_bias, - 'has_mlp_bias': - has_mlp_bias, - 'has_model_final_layernorm': - has_model_final_layernorm, - 'has_embedding_layernorm': - has_embedding_layernorm, - 'has_embedding_scale': - model_config.get('has_embedding_scale', False), - 'intermediate_size': - model_config['ffn_hidden_size'], - 'q_scaling': - q_scaling, - 'layernorm_position': - layernorm_position, - 'relative_attention': - relative_attention, - 'max_distance': - model_config.get('max_distance', 0), - 'num_buckets': - model_config.get('num_buckets', 0), - 'model_type': - family, - 'rescale_before_lm_head': - rescale_before_lm_head, - 'encoder_hidden_size': - model_config['hidden_size'], - 'encoder_num_heads': - model_config['num_heads'], - 'encoder_head_size': - model_config['head_size'], - 'skip_cross_kv': - model_config['skip_cross_kv'], - 'use_implicit_relative_attention': - model_config['use_implicit_relative_attention'], - 'decoder_start_token_id': - model_config['decoder_start_token_id'], - }) - tllm_model = tensorrt_llm.models.DecoderModel(pretrained_config) - if use_weight_only and family == 'whisper': - tllm_model = quantize(tllm_model, quant_config) - - tllm_model.precompute_relative_attention_bias(build_config) - - start = time.time() - engine = build(tllm_model, build_config) - assert engine.engine is not None, f'Failed to build engine for rank {runtime_rank}' - build_time = round(time.time() - start, 2) - - engine.save(os.path.join(args.output_dir, component)) - - return engine, model_config, build_time - - -def build_enc_dec(args): - build_config = get_build_config(args.model) - build_config = BuildConfig(**build_config) - model_config = get_model_config(args.model) - if args.force_num_layer_1: - model_config['num_layers'] = 1 - - if args.profiling_verbosity != "layer_names_only": - build_config.profiling_verbosity = args.profiling_verbosity - - if args.max_batch_size is not None: - build_config.max_batch_size = args.max_batch_size - if args.max_input_len is not None: - build_config.max_encoder_input_len = args.max_input_len - build_config.max_input_len = args.max_input_len - if args.max_seq_len is not None: - build_config.max_seq_len = args.max_seq_len - if args.max_beam_width is not None: - build_config.max_beam_width = args.max_beam_width - if args.opt_batch_size is not None: - build_config.opt_batch_size = args.opt_batch_size - if args.opt_num_tokens is not None: - build_config.opt_num_tokens = args.opt_num_tokens - build_config.max_num_tokens = build_config.max_batch_size * max( - build_config.max_encoder_input_len, build_config.max_beam_width) - - encoder_max_seq_len = build_config.max_encoder_input_len - decoder_max_seq_len = build_config.max_seq_len - - # Enable trt monitor memory for perf tests - build_config.monitor_memory = args.monitor_memory - - # for encoder, input len and output len both equal to max_encoder_input_len - build_config.max_input_len = encoder_max_seq_len - build_config.max_seq_len = encoder_max_seq_len - encoder_engine, encoder_model_config, encoder_build_time = enc_dec_build_helper( - component='encoder', - build_config=build_config, - model_config=model_config, - args=args) - - # for decoder, input len equals to 1 and output len equals to max_seq_len - build_config.max_input_len = 1 - build_config.max_seq_len = decoder_max_seq_len - decoder_engine, decoder_model_config, decoder_build_time = enc_dec_build_helper( - component='decoder', - build_config=build_config, - model_config=model_config, - args=args) - - return encoder_engine, decoder_engine, encoder_model_config, decoder_model_config, encoder_build_time, decoder_build_time - - -def main(args): - logger.set_level(args.log_level) - if args.model in get_allowed_models(benchmark_type="gpt"): - engine = build_gpt(args)[0] - engine_size = engine.engine.nbytes - elif args.model in get_allowed_models(benchmark_type="bert"): - engine = build_bert(args)[0] - engine_size = engine.nbytes - elif args.model in get_allowed_models(benchmark_type="enc_dec"): - encoder_engine, decoder_engine = build_enc_dec(args)[:2] - engine_size = encoder_engine.engine.nbytes + decoder_engine.engine.nbytes - else: - raise Exception(f'Unexpected model: {args.model}') - - # Print engine size for CI/CD to track. - logger.info( - f"Total engine size per GPU is {engine_size / 1048576:.2f} MiB.") - - -if __name__ == '__main__': - mp.set_start_method('spawn') - args = parse_arguments() - main(args) diff --git a/tests/integration/defs/perf/test_perf.py b/tests/integration/defs/perf/test_perf.py index c1281f8c9088..08c709891474 100644 --- a/tests/integration/defs/perf/test_perf.py +++ b/tests/integration/defs/perf/test_perf.py @@ -959,9 +959,6 @@ def set_runtime_configs(self, else: raise RuntimeError(f"Invalid runtime {self._config.runtime}.") - allowed_configs = import_allowed_perf_config() - allowed_models = allowed_configs.get_allowed_models() - if self._config.runtime == "bench": build_script = "trtllm-bench" elif self._config.runtime == "serve": @@ -970,12 +967,10 @@ def set_runtime_configs(self, build_script = None elif self._config.runtime == "multi_node_disagg_server": build_script = None - elif self._config.pp_size > 1 or self._config.model_name not in allowed_models: - build_script = "trtllm-build" else: - # build.py is used to build engines for both python and cpp runtime - build_script = os.path.join(llm_root, - "tests/integration/defs/perf/build.py") + raise RuntimeError( + f"Invalid runtime {self._config.runtime}: engine-build flows " + "were removed with the legacy TensorRT backend.") self._build_script = build_script self._benchmark_script = benchmark_script @@ -985,56 +980,6 @@ def set_runtime_configs(self, self._llm_root = llm_root self._gpu_clock_lock = gpu_clock_lock - def get_trtllm_build_command(self, engine_dir, checkpoint_dir) -> list: - build_cmd = [ - self._build_script, f"--output_dir={engine_dir}", - f"--checkpoint_dir={checkpoint_dir}", - f"--workers={self._config.tp_size}", - f"--use_paged_context_fmha=enable", f"--monitor_memory", - f"--max_batch_size={self._config.max_batch_size}" - ] - # For Multiple Profiles - if self._config.multiple_profiles: - build_cmd.append(f"--multiple_profiles=enable") - else: - build_cmd.append(f"--multiple_profiles=disable") - num_beams = self._config.num_beams - if num_beams > 1: - build_cmd.append(f"--max_beam_width={num_beams}") - gpu_percent = self._config.gpu_weights_percent - if gpu_percent != -1: - build_cmd += [f"--weight_streaming"] - # For engine inspector - build_cmd.append("--profiling_verbosity=layer_names_only") - if self._config.num_loras > 0: - if "mixtral" in self._config.model_name: - build_cmd.append(f"--lora_plugin=auto") - build_cmd.append(f"--moe_plugin=auto") - build_cmd.append(f"--lora_target_modules") - build_cmd.append(f"attn_q") - build_cmd.append(f"attn_k") - build_cmd.append(f"attn_v") - build_cmd.append(f"attn_dense") - build_cmd.append(f"moe_h_to_4h") - build_cmd.append(f"moe_4h_to_h") - build_cmd.append(f"moe_gate") - build_cmd.append(f"moe_router") - elif "llama" in self._config.model_name: - build_cmd.append(f"--lora_plugin=float16") - build_cmd.append(f"--lora_target_modules") - build_cmd.append(f"attn_q") - build_cmd.append(f"attn_k") - build_cmd.append(f"attn_v") - build_cmd.append(f"attn_dense") - build_cmd.append(f"mlp_h_to_4h") - build_cmd.append(f"mlp_4h_to_h") - build_cmd.append(f"mlp_gate") - if TIMING_CACHE_DIR and not self._config.build_only: - timing_cache = os.path.join(TIMING_CACHE_DIR, "model.cache") - build_cmd.append(f"--input_timing_cache={timing_cache}") - build_cmd.append(f"--output_timing_cache={timing_cache}") - return build_cmd - def get_trtllm_bench_model(self): return get_model_dir(self._config.model_name) @@ -1071,8 +1016,6 @@ def get_trtllm_bench_build_command(self, engine_dir) -> list: def get_prepare_data_command(self, engine_dir, input_len, output_len) -> list: data_cmd = [] - prepare_data_script = os.path.join(self._llm_root, "benchmarks", - "prepare_dataset.py") if self._config.model_name in MODEL_PATH_DICT.keys(): tokenizer_dir = os.path.join( @@ -1158,13 +1101,9 @@ def get_prepare_data_command(self, engine_dir, input_len, f"--input-stdev={istdev}", f"--output-stdev={ostdev}" ] else: - data_cmd += [ - "python3", prepare_data_script, f"--output={dataset_path}", - f"--tokenizer={tokenizer_dir}", f"token-norm-dist", - f"--num-requests={self._config.num_reqs}", - f"--input-mean={input_len}", f"--output-mean={output_len}", - f"--input-stdev={istdev}", f"--output-stdev={ostdev}" - ] + raise RuntimeError( + f"Unsupported build script {self._build_script} for " + "dataset preparation.") return data_cmd diff --git a/tests/integration/defs/test_e2e.py b/tests/integration/defs/test_e2e.py index 643e72f46622..215731702435 100644 --- a/tests/integration/defs/test_e2e.py +++ b/tests/integration/defs/test_e2e.py @@ -40,35 +40,6 @@ _MEM_FRACTION_95 = 0.95 -def test_gpt3_175b_1layers_build_only(llm_root, llm_venv, engine_dir): - """Build GPT-3 175B: 96 layer w/ plugins""" - example_root = os.path.join(llm_root, "examples", "models", "core", "gpt") - engine_dir = os.path.join(engine_dir, "gpt-175-96layers-build-only") - - dtype = 'float16' - convert_cmd = [ - f"{example_root}/../../../generate_checkpoint_config.py", - f"--output_path={engine_dir}/ckpt_config.json", - "--architecture=GPTForCausalLM", f"--dtype={dtype}", - "--num_hidden_layers=1", "--num_attention_heads=96", - "--hidden_size=12288", "--vocab_size=51200", "--tp_size=8" - ] - venv_check_call(llm_venv, convert_cmd) - - print("Building engines...") - build_cmd = [ - "trtllm-build", - f"--model_config={engine_dir}/ckpt_config.json", - f"--output_dir={engine_dir}", - "--max_batch_size=256", - "--max_input_len=200", - "--max_seq_len=400", - "--max_beam_width=1", - f"--gpt_attention_plugin={dtype}", - ] - check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env) - - @pytest.mark.parametrize("model_name,model_path", [ ("DeepSeek-R1-Distill-Qwen-1.5B", "DeepSeek-R1-Distill-Qwen-1.5B"), ]) diff --git a/tests/integration/defs/triton_server/conftest.py b/tests/integration/defs/triton_server/conftest.py index 18f6087ff6b9..932585184826 100644 --- a/tests/integration/defs/triton_server/conftest.py +++ b/tests/integration/defs/triton_server/conftest.py @@ -172,11 +172,6 @@ def tensorrt_llm_gptj_example_root(llm_backend_root): return os.path.join(llm_backend_root, "../examples/models/contrib/gptj") -@pytest.fixture(scope="session") -def tensorrt_llm_multimodal_example_root(llm_backend_root): - return os.path.join(llm_backend_root, "../examples/models/core/multimodal") - - @pytest.fixture(scope="session") def tensorrt_llm_opt_example_root(llm_backend_root): return os.path.join(llm_backend_root, "../examples/models/contrib/opt") @@ -207,21 +202,11 @@ def tensorrt_llm_llama_example_root(llm_backend_root): return os.path.join(llm_backend_root, "../examples/models/core/llama") -@pytest.fixture(scope="session") -def tensorrt_llm_qwen_example_root(llm_backend_root): - return os.path.join(llm_backend_root, "../examples/models/core/qwen") - - @pytest.fixture(scope="session") def tensorrt_llm_mllama_example_root(llm_backend_root): return os.path.join(llm_backend_root, "../examples/models/core/mllama") -@pytest.fixture(scope="session") -def tensorrt_llm_mixtral_example_root(llm_backend_root): - return os.path.join(llm_backend_root, "../examples/models/core/mixtral") - - @pytest.fixture(scope="session") def inflight_batcher_llm_client_root(llm_backend_root): inflight_batcher_llm_client_root = os.path.join(llm_backend_root, diff --git a/tests/integration/test_lists/qa/README.md b/tests/integration/test_lists/qa/README.md index c82f055ed020..e242584e64f3 100644 --- a/tests/integration/test_lists/qa/README.md +++ b/tests/integration/test_lists/qa/README.md @@ -90,7 +90,7 @@ cd tests/integration/defs # Run all fp8 functional test pytest --no-header -vs --test-list=../test_lists/qa/llm_function_full.txt -k fp8 # Run a single test case -pytest -vs accuracy/test_cli_flow.py::TestLlama3_1_8B::test_auto_dtype +pytest -vs accuracy/test_llm_api_pytorch.py::TestLlama3_1_8B::test_auto_dtype ``` ### Automated Execution diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 02c43b89f8dd..f9949c97654a 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -693,6 +693,7 @@ accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_parallelism[TP accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_block_reuse[DEP4_MTP_ON] accuracy/test_llm_api_pytorch.py::TestPhi4MiniInstruct::test_auto_dtype accuracy/test_llm_api_pytorch.py::TestQwen2_7BInstruct::test_auto_dtype +accuracy/test_llm_api_pytorch.py::TestQwen2_7BInstruct::test_tp2 accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_fp8[latency] accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_fp8[throughput_latency] accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4[latency_moe_cutlass] diff --git a/tests/integration/test_lists/test-db/README.md b/tests/integration/test_lists/test-db/README.md index f962e9053f0a..74e8a137cceb 100644 --- a/tests/integration/test_lists/test-db/README.md +++ b/tests/integration/test_lists/test-db/README.md @@ -31,8 +31,8 @@ l0_e2e: - '*h100*' linux_distribution_name: ubuntu* tests: - - examples/test_llama.py::test_llm_llama_v3_1_1node_multi_gpus[llama-3.1-8b-enable_fp8] - - examples/test_llama.py::test_llm_llama_v3_1_1node_multi_gpus[llama-3.1-70b-enable_fp8] + - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=False-attn_backend=TRTLLM-torch_compile=False] + - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=True-attn_backend=TRTLLM-torch_compile=False] ``` ## Generating Test Lists diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index bf8c6d131d50..f4f8f324075e 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -180,20 +180,6 @@ l0_a10: - thirdparty/test_git_modules.py::test_gitmodules # helper-script unit tests (CPU-only, ~3s) - unittest/scripts -- condition: - ranges: - system_gpu_count: - gte: 1 - lte: 1 - wildcards: - gpu: - - '*a10*' - linux_distribution_name: ubuntu* - terms: - stage: pre_merge - backend: tensorrt - tests: - # ------------- TRT tests --------------- - unittest/kv_cache_manager_v2_tests # 4 min - unittest/dynamo - unittest/api_stability @@ -209,30 +195,16 @@ l0_a10: linux_distribution_name: ubuntu* terms: stage: post_merge - backend: tensorrt + backend: pytorch tests: - # ------------- Move from pre_merge to post_merge --------------- + - stress_test/stress_test.py::test_run_stress_test[llama-v3-8b-instruct-hf_tp1-stress_time_300s_timeout_450s-GUARANTEED_NO_EVICT-pytorch-stress-test] + - stress_test/stress_test.py::test_run_stress_test[llama-v3-8b-instruct-hf_tp1-stress_time_300s_timeout_450s-MAX_UTILIZATION-pytorch-stress-test] - llmapi/test_llm_examples.py::test_llmapi_chat_example - llmapi/test_llm_examples.py::test_llmapi_server_example - llmapi/test_llm_examples.py::test_llmapi_kv_cache_connector[Qwen2-0.5B] - test_e2e.py::test_openai_health - test_e2e.py::test_trtllm_serve_example - llmapi/test_llm_examples.py::test_llmapi_quickstart_atexit -- condition: - ranges: - system_gpu_count: - gte: 1 - lte: 1 - wildcards: - gpu: - - '*a10*' - linux_distribution_name: ubuntu* - terms: - stage: post_merge - backend: pytorch - tests: - - stress_test/stress_test.py::test_run_stress_test[llama-v3-8b-instruct-hf_tp1-stress_time_300s_timeout_450s-GUARANTEED_NO_EVICT-pytorch-stress-test] - - stress_test/stress_test.py::test_run_stress_test[llama-v3-8b-instruct-hf_tp1-stress_time_300s_timeout_450s-MAX_UTILIZATION-pytorch-stress-test] - condition: ranges: system_gpu_count: diff --git a/tests/integration/test_lists/test-db/l0_a100.yml b/tests/integration/test_lists/test-db/l0_a100.yml index a61443e641ad..fe9ec737f01c 100644 --- a/tests/integration/test_lists/test-db/l0_a100.yml +++ b/tests/integration/test_lists/test-db/l0_a100.yml @@ -52,7 +52,7 @@ l0_a100: linux_distribution_name: ubuntu* terms: stage: post_merge - backend: tensorrt + backend: pytorch tests: - unittest/llmapi/test_llm.py -m "part0" - unittest/llmapi/test_llm.py -m "not part0" TIMEOUT (90) diff --git a/tests/integration/test_lists/test-db/l0_a30.yml b/tests/integration/test_lists/test-db/l0_a30.yml index d406ca8829a1..14ddddd8c619 100644 --- a/tests/integration/test_lists/test-db/l0_a30.yml +++ b/tests/integration/test_lists/test-db/l0_a30.yml @@ -64,39 +64,6 @@ l0_a30: - cpp/test_unit_tests.py::test_unit_tests[layers-80] - cpp/test_unit_tests.py::test_unit_tests[runtime-80] - cpp/test_unit_tests.py::test_unit_tests[thop-80] -- condition: - ranges: - system_gpu_count: - gte: 1 - lte: 1 - wildcards: - gpu: - - '*a30*' - linux_distribution_name: ubuntu* - terms: - stage: post_merge - backend: tensorrt - tests: - # ------------- TRT tests --------------- - - examples/test_qwen2audio.py::test_llm_qwen2audio_single_gpu[qwen2_audio_7b_instruct] - - examples/test_granite.py::test_llm_granite[granite-3.0-1b-a400m-instruct-bfloat16] # 5 mins - - examples/test_internlm.py::test_llm_internlm2_7b_1node_1gpu[bfloat16-enable_context_fmha-enable_gemm_plugin-enable_attention_plugin-nb:2] # 5 mins - - examples/test_ngram.py::test_llm_ngram_1gpu[streaming-gpt2-use_cpp_session-use_tokens-max_matching_ngram_size_2-max_draft_len_8-float16-bs2] # 1 min -- condition: - ranges: - system_gpu_count: - gte: 1 - lte: 1 - wildcards: - gpu: - - '*a30*' - linux_distribution_name: ubuntu* - terms: - stage: post_merge - backend: tensorrt - tests: - - examples/test_granite.py::test_llm_granite[granite-3.0-2b-instruct-bfloat16] # 5 mins - - examples/test_ngram.py::test_llm_ngram_1gpu[no_streaming-gpt2-use_cpp_session-use_tokens-max_matching_ngram_size_2-max_draft_len_8-float16-bs2] # 1 min # ------------- AutoDeploy Backend Stages --------------- - condition: ranges: diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 3bdfda75c3bb..d812dcf5d84e 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -259,9 +259,8 @@ l0_b200: linux_distribution_name: ubuntu* terms: stage: post_merge - backend: tensorrt + backend: pytorch tests: - # ------------- TRT tests --------------- - unittest/llmapi/test_llm_quant.py # 3.5 mins on B200 - unittest/disaggregated/test_openai_server_info.py - condition: diff --git a/tests/integration/test_lists/test-db/l0_dgx_h200.yml b/tests/integration/test_lists/test-db/l0_dgx_h200.yml index cf7863c91f27..f707f7e61036 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h200.yml @@ -142,27 +142,6 @@ l0_dgx_h200: - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507::test_skip_softmax_attention_4gpus[target_sparsity_0.5-fp8kv=True] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507::test_skip_softmax_attention_4gpus[target_sparsity_0.9-fp8kv=False] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507::test_skip_softmax_attention_4gpus[target_sparsity_0.9-fp8kv=True] -- condition: - ranges: - system_gpu_count: - gte: 4 - lte: 4 - wildcards: - gpu: - - '*h200*' - linux_distribution_name: ubuntu* - cpu: x86_64 - terms: - stage: post_merge - backend: tensorrt - tests: - # ------------- TRT tests --------------- - unittest/llmapi/test_llm_kv_cache_events.py::test_llm_api_attention_dp_kv_events - - examples/test_nemotron_nas.py::test_nemotron_nas_summary_2gpu[DeciLM-7B] - llmapi/test_llm_examples.py::test_llmapi_example_distributed_tp2 - - examples/test_multimodal.py::test_llm_multimodal_general[Llama-3.2-11B-Vision-pp:1-tp:2-bfloat16-bs:1-cpp_e2e:False-nb:1] - # ------------- TRT tests --------------- - - examples/test_bert.py::test_llm_bert_general[compare_hf-enable_remove_input_padding-use_attention_plugin-enable_context_fmha-tp:2-pp:1-float16-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] - - examples/test_bert.py::test_llm_bert_general[compare_hf-enable_remove_input_padding-use_attention_plugin-enable_context_fmha-tp:2-pp:1-float16-RobertaForQuestionAnswering-bert/roberta-base-squad2] - - examples/test_bert.py::test_llm_bert_general[compare_hf-enable_remove_input_padding-disable_attention_plugin-disable_context_fmha-tp:2-pp:1-float16-RobertaForSequenceClassification-bert/twitter-roberta-base-emotion] - unittest/llmapi/apps/_test_openai_multi_gpu.py -m "part0" diff --git a/tests/integration/test_lists/test-db/l0_gb203.yml b/tests/integration/test_lists/test-db/l0_gb203.yml index 11d692a85fdb..9fe58993698a 100644 --- a/tests/integration/test_lists/test-db/l0_gb203.yml +++ b/tests/integration/test_lists/test-db/l0_gb203.yml @@ -12,11 +12,8 @@ l0_gb203: linux_distribution_name: ubuntu* terms: stage: pre_merge - backend: tensorrt + backend: pytorch tests: - # ------------- TRT tests --------------- - # - examples/test_qwen.py::test_llm_qwen1_5_7b_single_gpu_lora[qwen1.5_7b_chat-Qwen1.5-7B-Chat-750Mb-lora] # https://nvbugs/5234573 - # - examples/test_qwen.py::test_llm_qwen_single_gpu_summary[qwen2.5_1.5b_instruct-enable_paged_kv_cache-enable_remove_input_padding-enable_weight_only-enable_fmha_fp32_acc] # https://nvbugs/5234573 - llmapi/test_llm_examples.py::test_llmapi_quickstart - llmapi/test_llm_examples.py::test_llmapi_example_inference - llmapi/test_llm_examples.py::test_llmapi_example_inference_async diff --git a/tests/integration/test_lists/test-db/l0_gh200.yml b/tests/integration/test_lists/test-db/l0_gh200.yml index f62424367d51..6cf8033c2d6b 100644 --- a/tests/integration/test_lists/test-db/l0_gh200.yml +++ b/tests/integration/test_lists/test-db/l0_gh200.yml @@ -12,7 +12,7 @@ l0_gh200: cpu: aarch64 terms: stage: post_merge - backend: tensorrt + backend: pytorch tests: - unittest/bindings - unittest/llmapi/test_llm_quant.py diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 2874ce258d48..4da1d645b59b 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -284,9 +284,8 @@ l0_h100: linux_distribution_name: ubuntu* terms: stage: pre_merge - backend: tensorrt + backend: pytorch tests: - # ------------- TRT tests --------------- - unittest/llmapi/test_llm_quant.py # 5.5 mins on H100 - examples/visual_gen/test_visual_gen.py::test_visual_gen_quickstart - examples/visual_gen/test_visual_gen.py::test_visual_gen_api_walkthrough @@ -332,6 +331,20 @@ l0_h100: - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-decoder-cuda-graph-on-beam2-batch2-t5-small] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-decoder-cuda-graph-on-beam2-batch2-flan-t5-small] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v2-decoder-cuda-graph-on-greedy-batch2-t5-small] + - examples/test_gpt.py::test_gpt_oss_20b_lora_torch[gpt-oss-20b-lora-adapter_NIM_r8-gpt-oss-20b] + - unittest/bindings # 8 mins on H100 + - unittest/kv_cache_manager_v2_tests # 4 min + # ------------- KV Cache Iteration Stats --------------- + - unittest/executor/test_stats_serializer.py + - unittest/metrics/test_collector.py + - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_cold_start + - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_partial_block_reuse + - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_full_block_reuse + - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_shared_prefix + - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_batch_generation + - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_long_context + - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_rapid_fire + - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_field_completeness - llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[eager] TIMEOUT (90) - llmapi/test_llm_api_pytorch_moe_lora.py::test_mixtral_moe_routed_expert_fp8_multi_lora_varying_ranks[eager] TIMEOUT (90) - llmapi/test_llm_api_pytorch_moe_lora.py::test_mixtral_moe_routed_expert_fp8_multi_lora_varying_ranks[cudagraph] TIMEOUT (90) @@ -417,42 +430,6 @@ l0_h100: - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS and None" # Documentation URL validation (CPU-only, no GPU needed) - test_doc.py::test_url_validity -- condition: - ranges: - system_gpu_count: - gte: 1 - lte: 1 - wildcards: - gpu: - - '*h100*' - linux_distribution_name: ubuntu* - terms: - stage: post_merge - backend: tensorrt - tests: - # ------------- TRT tests --------------- - - examples/test_nemotron.py::test_llm_nemotron_3_8b_1gpu[bfloat16-fp8] # 18mins - - examples/test_whisper.py::test_llm_whisper_general[large-v3-disable_gemm_plugin-disable_attention_plugin-disable_weight_only-float16-nb:1-use_python_runtime] # 8 mins - - examples/test_granite.py::test_granite_bf16_lora[granite-3.0-1b-a400m-instruct] - - examples/test_gpt.py::test_llm_minitron_fp8_with_pseudo_loras[4b] TIMEOUT (90) - - unittest/bindings # 8 mins on H100 - - examples/test_multimodal.py::test_llm_multimodal_general[Llama-3.2-11B-Vision-pp:1-tp:1-bfloat16-bs:8-cpp_e2e:False-nb:1] - - examples/test_multimodal.py::test_llm_fp8_multimodal_general[fp8-fp8-scienceqa-Llama-3.2-11B-Vision-Instruct-pp:1-tp:1-bfloat16-bs:1-cpp_e2e:False] - - examples/test_medusa.py::test_llm_medusa_with_qaunt_base_model_1gpu[fp8-use_py_session-medusa-vicuna-7b-v1.3-4-heads-float16-bs1] - - examples/test_medusa.py::test_llm_medusa_with_qaunt_base_model_1gpu[fp8-use_cpp_session-medusa-vicuna-7b-v1.3-4-heads-float16-bs1] - - examples/test_gpt.py::test_gpt_oss_20b_lora_torch[gpt-oss-20b-lora-adapter_NIM_r8-gpt-oss-20b] - - unittest/kv_cache_manager_v2_tests # 4 min - # ------------- KV Cache Iteration Stats --------------- - - unittest/executor/test_stats_serializer.py - - unittest/metrics/test_collector.py - - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_cold_start - - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_partial_block_reuse - - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_full_block_reuse - - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_shared_prefix - - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_batch_generation - - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_long_context - - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_rapid_fire - - kv_cache/test_kv_cache_iteration_stats.py::TestKvCacheIterationStats::test_field_completeness - condition: ranges: system_gpu_count: diff --git a/tests/integration/test_lists/test-db/l0_l40s.yml b/tests/integration/test_lists/test-db/l0_l40s.yml index d74675c711dd..bf71ba62fc8c 100644 --- a/tests/integration/test_lists/test-db/l0_l40s.yml +++ b/tests/integration/test_lists/test-db/l0_l40s.yml @@ -65,10 +65,8 @@ l0_l40s: linux_distribution_name: ubuntu* terms: stage: post_merge - backend: tensorrt + backend: pytorch tests: - # ------------- TRT tests --------------- - - examples/test_nemotron_nas.py::test_nemotron_nas_summary_1gpu[DeciLM-7B] - llmapi/test_llm_examples.py::test_llmapi_quickstart - examples/visual_gen/test_visual_gen.py::test_visual_gen_quickstart - examples/visual_gen/test_visual_gen.py::test_visual_gen_api_walkthrough @@ -79,22 +77,6 @@ l0_l40s: - llmapi/test_llm_examples.py::test_llmapi_example_guided_decoding - llmapi/test_llm_examples.py::test_llmapi_example_logits_processor - examples/test_llm_api_with_mpi.py::test_llm_api_single_gpu_with_mpirun[TinyLlama-1.1B-Chat-v1.0] -- condition: - ranges: - system_gpu_count: - gte: 1 - lte: 1 - wildcards: - gpu: - - '*l40s*' - linux_distribution_name: ubuntu* - terms: - stage: post_merge - backend: tensorrt - tests: - - examples/test_nemotron.py::test_llm_nemotron_3_8b_1gpu[bfloat16-fp8] # 18mins - - examples/test_whisper.py::test_llm_whisper_general[large-v3-disable_gemm_plugin-disable_attention_plugin-disable_weight_only-float16-nb:1-use_python_runtime] # 8 mins - - examples/test_granite.py::test_granite_bf16_lora[granite-3.0-1b-a400m-instruct] TIMEOUT (90) - condition: ranges: system_gpu_count: diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 4aebb18910db..e7b609836b74 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -150,26 +150,9 @@ disaggregated/test_workers.py::test_workers_conversation_router[TinyLlama-1.1B-C disaggregated/test_workers.py::test_workers_kv_cache_aware_router_deepseek_v3_lite_bf16[DeepSeek-V3-Lite-bf16] SKIP (https://nvbugs/6162322) disaggregated/test_workers.py::test_workers_kv_cache_aware_router_eviction[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6162322) examples/test_ad_speculative_decoding.py::test_autodeploy_eagle3_one_model_acceptance_rate[trtllm-torch-cudagraph] SKIP (https://nvbugs/6426841) -examples/test_bert.py::test_llm_bert_general[compare_hf-enable_remove_input_padding-disable_attention_plugin-disable_context_fmha-tp:2-pp:1-float16-RobertaForSequenceClassification-bert/twitter-roberta-base-emotion] SKIP (https://nvbugs/5234058) -examples/test_bert.py::test_llm_bert_general[compare_hf-enable_remove_input_padding-use_attention_plugin-enable_context_fmha-tp:2-pp:1-float16-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] SKIP (https://nvbugs/5234058) -examples/test_bert.py::test_llm_bert_general[compare_hf-enable_remove_input_padding-use_attention_plugin-enable_context_fmha-tp:2-pp:1-float16-RobertaForQuestionAnswering-bert/roberta-base-squad2] SKIP (https://nvbugs/5234058) -examples/test_gpt.py::test_llm_minitron_fp8_with_pseudo_loras[4b] SKIP (https://nvbugs/5606233) -examples/test_granite.py::test_granite_bf16_lora[granite-3.0-1b-a400m-instruct] SKIP (https://nvbugs/5431132) -examples/test_granite.py::test_llm_granite[granite-3.0-1b-a400m-instruct-bfloat16] SKIP (https://nvbugs/5608979) -examples/test_granite.py::test_llm_granite[granite-3.0-2b-instruct-bfloat16] SKIP (https://nvbugs/5608979) -examples/test_medusa.py::test_llm_medusa_with_qaunt_base_model_1gpu[fp8-use_cpp_session-medusa-vicuna-7b-v1.3-4-heads-float16-bs1] SKIP (https://nvbugs/5802248) -examples/test_medusa.py::test_llm_medusa_with_qaunt_base_model_1gpu[fp8-use_py_session-medusa-vicuna-7b-v1.3-4-heads-float16-bs1] SKIP (https://nvbugs/5333849) -examples/test_multimodal.py::test_llm_fp8_multimodal_general[fp8-fp8-scienceqa-Llama-3.2-11B-Vision-Instruct-pp:1-tp:1-bfloat16-bs:1-cpp_e2e:False] SKIP (https://nvbugs/5222697) -examples/test_multimodal.py::test_llm_multimodal_general[Llama-3.2-11B-Vision-pp:1-tp:1-bfloat16-bs:8-cpp_e2e:False-nb:1] SKIP (https://nvbugs/5333818) -examples/test_multimodal.py::test_llm_multimodal_general[Llama-3.2-11B-Vision-pp:1-tp:2-bfloat16-bs:1-cpp_e2e:False-nb:1] SKIP (https://nvbugs/5333818) -examples/test_nemotron.py::test_llm_nemotron_3_8b_1gpu[bfloat16-fp8] SKIP (https://nvbugs/4961624) -examples/test_nemotron_nas.py::test_nemotron_nas_summary_1gpu[DeciLM-7B] SKIP (https://nvbugs/5444636) -examples/test_nemotron_nas.py::test_nemotron_nas_summary_2gpu[DeciLM-7B] SKIP (https://nvbugs/5444636) -examples/test_qwen2audio.py::test_llm_qwen2audio_single_gpu[qwen2_audio_7b_instruct] SKIP (https://nvbugs/5447530) examples/test_ray.py::test_llm_inference_distributed_ray[pp2] SKIP (https://nvbugs/6427411) examples/test_ray.py::test_llm_inference_distributed_ray[tp2pp2] SKIP (https://nvbugs/6427411) examples/test_ray.py::test_ray_disaggregated_serving[tp2] SKIP (https://nvbugs/5612502) -examples/test_whisper.py::test_llm_whisper_general[large-v3-disable_gemm_plugin-disable_attention_plugin-disable_weight_only-float16-nb:1-use_python_runtime] SKIP (https://nvbugs/5244570) examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2i_lpips_against_golden SKIP (https://nvbugs/6418815) examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2v_lpips_against_golden SKIP (https://nvbugs/6437341) examples/visual_gen/test_visual_gen.py::test_ltx2_cuda_graph_trtllm_backend SKIP (https://nvbugs/6463822) @@ -190,9 +173,6 @@ full:B200/disaggregated/test_disaggregated.py::test_disaggregated_stress_test[in full:B200/disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-qwen3_32b_fp8_stress] SKIP (https://nvbugs/6472256) full:B200/llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[cudagraph] SKIP (https://nvbugs/6475623) full:B200/llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[eager] SKIP (https://nvbugs/6475621) -full:B200/perf/test_perf.py::test_perf[quant:int8_sq_per_tensor] SKIP (https://nvbugs/5161074) -full:B200/perf/test_perf.py::test_perf[quant:int8_sq_per_token_channel] SKIP (https://nvbugs/5161074) -full:B200/perf/test_perf.py::test_perf[quant:w4a8_awq] SKIP (https://nvbugs/5161074) full:B200/test_e2e.py::test_multi_nodes_eval[Qwen3/Qwen3-235B-A22B-tp16-mmlu] SKIP (https://nvbugs/6424188) full:B200/test_e2e.py::test_multi_nodes_eval[Qwen3/saved_models_Qwen3-235B-A22B_nvfp4_hf-tp16-mmlu] SKIP (https://nvbugs/6424188) full:B200/test_e2e.py::test_qwen_e2e_cpprunner_large_new_tokens[DeepSeek-R1-Distill-Qwen-1.5B-DeepSeek-R1-Distill-Qwen-1.5B] SKIP (https://nvbugs/6414760) @@ -259,8 +239,6 @@ full:GB300/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_ full:GB300/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6422318) full:GB300/accuracy/test_llm_api_pytorch_multimodal.py::TestGemma3_27BInstruct::test_fp8_prequantized SKIP (https://nvbugs/6479708) full:GB300/test_e2e.py::test_qwen_e2e_cpprunner_large_new_tokens[DeepSeek-R1-Distill-Qwen-1.5B-DeepSeek-R1-Distill-Qwen-1.5B] SKIP (https://nvbugs/6414760) -full:GH200/examples/test_nemotron.py::test_llm_nemotron_3_8b_1gpu[bfloat16-fp8] SKIP (arm is not supported) -full:GH200/examples/test_qwen2audio.py::test_llm_qwen2audio_single_gpu[qwen2_audio_7b_instruct] SKIP (arm is not supported) full:H100/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=False] SKIP (https://nvbugs/6313072) full:H100/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=True] SKIP (https://nvbugs/6313072) full:H100/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=False] SKIP (https://nvbugs/6422343) @@ -351,10 +329,6 @@ full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::Tes full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16 SKIP (https://nvbugs/6273850) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_dflash SKIP (https://nvbugs/6273850) full:RTX_PRO_6000_Blackwell_Server_Edition/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_fp8 SKIP (https://nvbugs/6273850) -full:RTX_PRO_6000_Blackwell_Server_Edition/perf/test_perf.py::test_perf[quant:int8_sq_per_tensor] SKIP (https://nvbugs/5161074) -full:RTX_PRO_6000_Blackwell_Server_Edition/perf/test_perf.py::test_perf[quant:int8_sq_per_token_channel] SKIP (https://nvbugs/5161074) -full:RTX_PRO_6000_Blackwell_Server_Edition/perf/test_perf.py::test_perf[quant:w4a8_awq] SKIP (https://nvbugs/5161074) -full:sm100/examples/test_nemotron.py::test_llm_nemotron_3_8b_1gpu[bfloat16-fp8] SKIP (megatron-core 0.8 is not supported in python 3.12) full:sm100/unittest/bindings SKIP (Disable for Blackwell) kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2Llama::test_chunked_prefill SKIP (https://nvbugs/6428002) kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2Llama::test_eviction_with_block_reuse SKIP (https://nvbugs/6462303) @@ -362,22 +336,6 @@ llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_ llmapi/test_llm_api_pytorch_moe_lora.py::test_mixtral_moe_routed_expert_fp8_multi_lora_varying_ranks[cudagraph] SKIP (https://nvbugs/6463829) llmapi/test_llm_api_pytorch_moe_lora.py::test_mixtral_moe_routed_expert_fp8_multi_lora_varying_ranks[eager] SKIP (https://nvbugs/6463829) llmapi/test_llm_examples.py::test_llmapi_speculative_decoding_eagle3 SKIP (https://nvbugs/6075431) -perf/test_perf.py::test_perf[bart_large_cnn-plugin-float16-bs:8-input_output_len:60,20] SKIP # (https://nvidia.slack.com/archives/C059LSY62BT/p1704525727177449) -perf/test_perf.py::test_perf[flan_t5_base-bench-float16-input_output_len:128,20] SKIP -perf/test_perf.py::test_perf[flan_t5_base-plugin-float16-bs:8-input_output_len:60,20] SKIP # (https://nvidia.slack.com/archives/C059LSY62BT/p1704525727177449) -perf/test_perf.py::test_perf[flan_t5_large-bench-float16-input_output_len:128,20-gpus:2] SKIP -perf/test_perf.py::test_perf[flan_t5_large-bench-float16-input_output_len:128,20] SKIP -perf/test_perf.py::test_perf[flan_t5_large-bench-float16-maxbs:1-input_output_len:128,20-gpus:2] SKIP -perf/test_perf.py::test_perf[gpt_20b-bench-float16-maxbs:8-input_output_len:128,128-reqs:80-gpus:8] SKIP -perf/test_perf.py::test_perf[gpt_20b-bench-float16-maxbs:8-input_output_len:512,32-reqs:80-gpus:8] SKIP -perf/test_perf.py::test_perf[mamba_2.8b-bench-float16-input_output_len:128,128] SKIP -perf/test_perf.py::test_perf[mamba_2.8b-bench-float16-input_output_len:512,32] SKIP -perf/test_perf.py::test_perf[mamba_370m-bench-float16-input_output_len:128,128] SKIP -perf/test_perf.py::test_perf[mamba_370m-bench-float16-input_output_len:512,32] SKIP -perf/test_perf.py::test_perf[t5-bench-float16-input_output_len:128,20-gpus:2] SKIP -perf/test_perf.py::test_perf[t5-bench-float16-maxbs:1-input_output_len:128,20-gpus:2] SKIP -perf/test_perf.py::test_perf[t5_base-plugin-float16-bs:8-input_output_len:60,20] SKIP # (https://nvidia.slack.com/archives/C059LSY62BT/p1704525727177449) -perf/test_perf.py::test_perf[whisper_large_v3-bench-float16-input_output_len:128,20] SKIP perf/test_perf_sanity.py::test_e2e[aggr_upload-deepseek_r1_fp4_v2_grace_blackwell-r1_fp4_v2_dep4_mtp1_1k8k] SKIP (https://nvbugs/6422339) perf/test_perf_sanity.py::test_e2e[aggr_upload-deepseek_r1_fp8_blackwell-r1_fp8_tp8_mtp3_8k1k] SKIP (https://nvbugs/6432948) perf/test_perf_sanity.py::test_e2e[aggr_upload-dynamo_gpt_oss_120b_fp4_blackwell-gpt_oss_fp4_tep4_adp_cutlass_8k1k] SKIP (https://nvbugs/6374910) diff --git a/tests/unittest/tools/test_test_to_stage_mapping.py b/tests/unittest/tools/test_test_to_stage_mapping.py index 29052059259a..f8689ff12398 100644 --- a/tests/unittest/tools/test_test_to_stage_mapping.py +++ b/tests/unittest/tools/test_test_to_stage_mapping.py @@ -262,8 +262,15 @@ def test_backend_filtering_consistency(stage_query): f"at least one stage containing '{backend.upper()}', " \ f"but got stages: {stages}" - # Check that test does NOT map to stages of other backends - other_backends = all_backends - {backend} + # Check that test does NOT map to stages of backends it is not + # declared under (tests may legitimately be listed under several + # backends across test-db files). + declared_backends = { + b.strip() + for _, _, b in stage_query.test_map[test_name] + if b and b.strip() + } + other_backends = all_backends - declared_backends for stage in stages: stage_upper = stage.upper() for other_backend in other_backends: