From bea14cde7b79b1a6a3c3276cc676283db254045c Mon Sep 17 00:00:00 2001 From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Date: Fri, 17 Nov 2023 06:14:52 -0800 Subject: [PATCH 1/3] Update TensorRT-LLM --- README.md | 61 +- benchmarks/cpp/README.md | 16 +- benchmarks/cpp/gptSessionBenchmark.cpp | 45 +- benchmarks/python/allowed_configs.py | 6 +- benchmarks/python/gpt_benchmark.py | 20 +- benchmarks/python/mem_monitor.py | 2 +- cpp/CMakeLists.txt | 18 + .../tensorrt_llm/batch_manager/GptManager.h | 3 + .../batch_manager/kvCacheManager.h | 4 +- .../tensorrt_llm/batch_manager/llmRequest.h | 73 +- .../tensorrt_llm/runtime/bufferManager.h | 24 + .../tensorrt_llm/runtime/decodingOutput.h | 2 +- .../tensorrt_llm/runtime/generationOutput.h | 6 +- cpp/include/tensorrt_llm/runtime/gptDecoder.h | 17 +- .../tensorrt_llm/runtime/gptDecoderBatch.h | 34 +- .../tensorrt_llm/runtime/gptModelConfig.h | 12 + cpp/include/tensorrt_llm/runtime/gptSession.h | 23 +- cpp/include/tensorrt_llm/runtime/iBuffer.h | 56 +- .../tensorrt_llm/runtime/iGptDecoderBatch.h | 20 +- .../runtime/iStatefulGptDecoder.h | 13 +- cpp/tensorrt_llm/CMakeLists.txt | 28 +- .../libtensorrt_llm_batch_manager_static.a | 3 - ...sorrt_llm_batch_manager_static.pre_cxx11.a | 3 - .../aarch64-linux-gnu/version.txt | 3 - .../libtensorrt_llm_batch_manager_static.a | 4 +- ...sorrt_llm_batch_manager_static.pre_cxx11.a | 4 +- .../x86_64-linux-gnu/version.txt | 4 +- cpp/tensorrt_llm/common/mpiUtils.h | 13 +- .../threadblock/epilogue_tensor_op_int32.h | 21 +- .../gemm/device/gemm_universal_base_compat.h | 438 ++++++++ cpp/tensorrt_llm/kernels/CMakeLists.txt | 6 + .../kernels/beamSearchPenaltyKernels.cu | 1 + .../fpA_intB_gemm/fpA_intB_gemm_template.h | 4 +- .../int8_gemm/int8_gemm_template.h | 5 +- .../decoderMaskedMultiheadAttentionLaunch.h | 2 +- .../decoderMaskedMultiheadAttentionTemplate.h | 44 +- .../decoderMaskedMultiheadAttentionUtils.h | 6 + cpp/tensorrt_llm/kernels/decodingKernels.cu | 36 +- cpp/tensorrt_llm/kernels/decodingKernels.h | 3 + cpp/tensorrt_llm/kernels/gptKernels.cu | 13 +- cpp/tensorrt_llm/kernels/gptKernels.h | 7 +- cpp/tensorrt_llm/kernels/kvCacheUtils.h | 11 + cpp/tensorrt_llm/kernels/layernormKernels.cu | 2 - cpp/tensorrt_llm/kernels/rmsnormKernels.cu | 2 - .../kernels/samplingTopPKernels.cu | 2 - .../kernels/unfusedAttentionKernels.cu | 1 - .../layers/baseBeamSearchLayer.cu | 31 +- .../layers/dynamicDecodeLayer.cpp | 17 +- cpp/tensorrt_llm/layers/dynamicDecodeLayer.h | 12 +- cpp/tensorrt_llm/layers/fillBuffers.h | 68 ++ .../layers/onlineBeamSearchLayer.cu | 24 +- cpp/tensorrt_llm/plugins/CMakeLists.txt | 3 +- cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp | 3 + .../bertAttentionPlugin.cpp | 141 ++- .../bertAttentionPlugin/bertAttentionPlugin.h | 4 +- .../gptAttentionCommon/gptAttentionCommon.cpp | 87 +- .../gptAttentionPlugin/gptAttentionPlugin.cpp | 2 +- .../plugins/loraPlugin/CMakeLists.txt | 21 + .../plugins/loraPlugin/loraPlugin.cpp | 573 +++++++++++ .../plugins/loraPlugin/loraPlugin.h | 161 +++ .../weightOnlyGroupwiseQuantMatmulPlugin.cpp | 9 +- .../weightOnlyQuantMatmulPlugin.cpp | 8 +- cpp/tensorrt_llm/pybind/bindings.cpp | 7 +- .../pybind/runtime/generationOutput.cpp | 9 +- cpp/tensorrt_llm/runtime/bufferManager.cpp | 49 +- cpp/tensorrt_llm/runtime/gptDecoder.cpp | 42 +- cpp/tensorrt_llm/runtime/gptDecoderBatch.cpp | 47 +- cpp/tensorrt_llm/runtime/gptJsonConfig.cpp | 6 +- cpp/tensorrt_llm/runtime/gptSession.cpp | 285 ++++-- cpp/tensorrt_llm/runtime/iBuffer.cpp | 29 + cpp/tensorrt_llm/runtime/ncclCommunicator.cpp | 33 +- cpp/tensorrt_llm/runtime/ncclCommunicator.h | 15 +- cpp/tensorrt_llm/runtime/runtimeBuffers.cpp | 41 +- cpp/tensorrt_llm/runtime/runtimeBuffers.h | 7 + cpp/tensorrt_llm/runtime/runtimeKernels.cu | 83 ++ cpp/tensorrt_llm/runtime/runtimeKernels.h | 3 + .../runtime/statefulGptDecoder.cpp | 31 +- cpp/tensorrt_llm/runtime/statefulGptDecoder.h | 17 +- .../runtime/utils/multiDeviceUtils.h | 24 +- cpp/tensorrt_llm/runtime/worldConfig.cpp | 14 +- cpp/tensorrt_llm/thop/ncclCommunicatorOp.cpp | 10 +- cpp/tensorrt_llm/thop/ncclCommunicatorOp.h | 1 - .../weightOnly/weightOnlyKernelTest.cpp | 34 +- .../scripts/build_chatglm_engines.py | 29 +- .../generate_expected_chatglm_output.py | 126 ++- cpp/tests/resources/scripts/test_cpp.py | 97 +- cpp/tests/runtime/bufferManagerTest.cpp | 19 +- cpp/tests/runtime/gptDecoderBatchTest.cpp | 58 +- cpp/tests/runtime/gptDecoderTest.cpp | 5 +- cpp/tests/runtime/gptSessionTest.cpp | 50 +- cpp/tests/runtime/tllmBuffersTest.cpp | 2 +- cpp/tests/runtime/torchTest.cpp | 2 +- docker/Dockerfile.multi | 16 +- docker/Makefile | 24 +- docker/common/install_base.sh | 4 +- docker/common/install_tensorrt.sh | 43 +- docs/source/blogs/H200launch.md | 2 +- docs/source/gpt_runtime.md | 5 + docs/source/index.rst | 1 + docs/source/installation.md | 17 +- docs/source/memory.md | 114 +++ docs/source/performance.md | 314 ++++++ docs/source/precision.md | 10 +- examples/baichuan/README.md | 31 +- examples/baichuan/build.py | 1 + examples/baichuan/requirements.txt | 1 + examples/baichuan/summarize.py | 401 -------- examples/blip2/run.py | 5 +- examples/bloom/README.md | 36 +- examples/bloom/build.py | 55 +- examples/bloom/requirements.txt | 1 + examples/bloom/summarize.py | 377 ------- examples/bloom/weight.py | 115 ++- examples/chatglm/.gitignore | 9 +- examples/chatglm/README.md | 90 +- examples/chatglm/build.py | 316 ++++-- examples/chatglm/process.py | 40 + examples/chatglm/quantize.py | 147 +++ examples/chatglm/requirements.txt | 2 +- examples/chatglm/run.py | 111 +- examples/chatglm/summarize.py | 473 --------- examples/chatglm/weight.py | 738 +++++++++---- examples/common/utils.py | 80 ++ examples/enc_dec/README.md | 106 +- examples/enc_dec/build.py | 307 ++++-- examples/enc_dec/download.py | 18 - examples/enc_dec/models/config.ini | 48 - examples/enc_dec/run.py | 430 +++++--- examples/enc_dec/t5/hf_convert.py | 222 ++++ examples/enc_dec/{ => t5}/weight.py | 243 ++++- examples/falcon/README.md | 77 +- examples/falcon/build.py | 77 +- examples/falcon/quantize.py | 2 +- examples/falcon/requirements.txt | 2 +- examples/falcon/summarize.py | 446 -------- examples/falcon/weight.py | 173 ++++ examples/gpt/README.md | 24 +- examples/gpt/build.py | 30 +- examples/gpt/requirements.txt | 2 + examples/gpt/run.py | 23 +- examples/gpt/summarize.py | 539 ---------- examples/gpt/weight.py | 75 +- examples/gptj/README.md | 23 +- examples/gptj/build.py | 4 +- examples/gptj/summarize.py | 416 -------- examples/gptneox/README.md | 78 +- examples/gptneox/build.py | 1 + examples/gptneox/requirements.txt | 1 + examples/gptneox/summarize.py | 380 ------- examples/internlm/README.md | 72 +- examples/internlm/build.py | 7 +- examples/internlm/requirements.txt | 1 + examples/internlm/summarize.py | 414 -------- examples/internlm/weight.py | 10 +- examples/llama/.gitignore | 1 + examples/llama/README.md | 127 ++- examples/llama/build.py | 59 +- examples/llama/hf_llama_convert.py | 31 +- examples/llama/requirements.txt | 1 + examples/llama/run.py | 48 +- examples/llama/smoothquant.py | 27 +- examples/llama/summarize_long.py | 74 +- examples/llama/weight.py | 227 +++- examples/mpt/README.md | 101 +- examples/mpt/build.py | 63 +- examples/mpt/convert_hf_mpt_to_ft.py | 56 +- examples/mpt/run.py | 11 +- examples/mpt/weight.py | 285 ++---- examples/opt/README.md | 71 +- examples/opt/build.py | 4 +- examples/opt/requirements.txt | 1 + examples/opt/summarize.py | 377 ------- examples/quantization/quantize.py | 143 +++ examples/qwen/README.md | 366 +++++++ examples/qwen/benchmark.py | 378 +++++++ examples/qwen/build.py | 623 +++++++++++ examples/qwen/hf_qwen_convert.py | 361 +++++++ examples/qwen/requirements.txt | 14 + examples/qwen/run.py | 315 ++++++ examples/qwen/smoothquant.py | 209 ++++ examples/{llama => qwen}/summarize.py | 319 +++--- .../qwen/utils}/__init__.py | 0 examples/qwen/utils/convert.py | 304 ++++++ examples/qwen/utils/utils.py | 134 +++ examples/qwen/weight.py | 524 ++++++++++ examples/summarize.py | 560 ++++++++++ requirements-dev.txt | 2 + scripts/build_wheel.py | 13 +- setup.py | 2 +- tensorrt_llm/_utils.py | 47 +- tensorrt_llm/functional.py | 130 ++- tensorrt_llm/graph_rewriting.py | 5 +- tensorrt_llm/layers/__init__.py | 3 + tensorrt_llm/layers/attention.py | 133 ++- tensorrt_llm/layers/lora.py | 66 ++ tensorrt_llm/mapping.py | 3 + tensorrt_llm/models/__init__.py | 11 +- tensorrt_llm/models/bloom/model.py | 1 + tensorrt_llm/models/chatglm/model.py | 161 ++- tensorrt_llm/models/enc_dec/model.py | 969 +++++++++++------- tensorrt_llm/models/generation_mixin.py | 34 +- tensorrt_llm/models/gpt/model.py | 99 +- tensorrt_llm/models/internlm/model.py | 427 -------- tensorrt_llm/models/llama/model.py | 35 +- tensorrt_llm/models/quantized/ammo.py | 4 + tensorrt_llm/models/quantized/quant.py | 54 +- tensorrt_llm/models/qwen/__init__.py | 14 + tensorrt_llm/models/qwen/model.py | 640 ++++++++++++ tensorrt_llm/module.py | 32 +- tensorrt_llm/parameter.py | 5 +- tensorrt_llm/plugin/plugin.py | 5 + tensorrt_llm/profiler.py | 2 + tensorrt_llm/quantization/layers.py | 8 +- tensorrt_llm/runtime/__init__.py | 7 +- tensorrt_llm/runtime/generation.py | 197 +++- tensorrt_llm/runtime/lora_manager.py | 74 ++ tensorrt_llm/runtime/model_runner.py | 335 ++++++ tensorrt_llm/runtime/session.py | 20 + tests/bindings/test_bindings.py | 13 + tests/bindings/test_gpt_session.py | 216 ++++ tests/model/test_gpt.py | 32 + tests/model/test_gpt_e2e.py | 3 +- tests/tools/plugin_gen/test_plugin_gen.py | 5 + 223 files changed, 14005 insertions(+), 6906 deletions(-) delete mode 100644 cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/libtensorrt_llm_batch_manager_static.a delete mode 100644 cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/libtensorrt_llm_batch_manager_static.pre_cxx11.a delete mode 100644 cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/version.txt create mode 100644 cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/device/gemm_universal_base_compat.h create mode 100644 cpp/tensorrt_llm/layers/fillBuffers.h create mode 100644 cpp/tensorrt_llm/plugins/loraPlugin/CMakeLists.txt create mode 100644 cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.cpp create mode 100644 cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.h create mode 100644 docs/source/memory.md delete mode 100644 examples/baichuan/summarize.py delete mode 100644 examples/bloom/summarize.py create mode 100644 examples/chatglm/process.py create mode 100644 examples/chatglm/quantize.py delete mode 100644 examples/chatglm/summarize.py create mode 100644 examples/common/utils.py delete mode 100644 examples/enc_dec/download.py delete mode 100644 examples/enc_dec/models/config.ini create mode 100644 examples/enc_dec/t5/hf_convert.py rename examples/enc_dec/{ => t5}/weight.py (55%) delete mode 100644 examples/falcon/summarize.py delete mode 100644 examples/gpt/summarize.py delete mode 100644 examples/gptj/summarize.py delete mode 100644 examples/gptneox/summarize.py delete mode 100644 examples/internlm/summarize.py delete mode 100644 examples/opt/summarize.py create mode 100644 examples/quantization/quantize.py create mode 100644 examples/qwen/README.md create mode 100644 examples/qwen/benchmark.py create mode 100644 examples/qwen/build.py create mode 100644 examples/qwen/hf_qwen_convert.py create mode 100644 examples/qwen/requirements.txt create mode 100644 examples/qwen/run.py create mode 100644 examples/qwen/smoothquant.py rename examples/{llama => qwen}/summarize.py (58%) rename {tensorrt_llm/models/internlm => examples/qwen/utils}/__init__.py (100%) create mode 100644 examples/qwen/utils/convert.py create mode 100644 examples/qwen/utils/utils.py create mode 100644 examples/qwen/weight.py create mode 100644 examples/summarize.py create mode 100644 tensorrt_llm/layers/lora.py delete mode 100644 tensorrt_llm/models/internlm/model.py create mode 100644 tensorrt_llm/models/qwen/__init__.py create mode 100644 tensorrt_llm/models/qwen/model.py create mode 100644 tensorrt_llm/runtime/lora_manager.py create mode 100644 tensorrt_llm/runtime/model_runner.py create mode 100644 tests/bindings/test_gpt_session.py diff --git a/README.md b/README.md index 894d9d94e845..aadebab5cecd 100644 --- a/README.md +++ b/README.md @@ -43,17 +43,22 @@ H200 FP8 achieves 11,819 tok/s on Llama2-13B on a single GPU, and is up to 1.9x - [Installation](#installation) - [Quick Start](#quick-start) - [Support Matrix](#support-matrix) + - [Devices](#devices) + - [Precision](#precision) + - [Key Features](#key-features) + - [Models](#models) - [Performance](#performance) - [Advanced Topics](#advanced-topics) - [Quantization](#quantization) - [In-flight Batching](#in-flight-batching) - [Attention](#attention) - [Graph Rewriting](#graph-rewriting) - - [Benchmarking](#benchmarking) + - [Benchmark](#benchmark) - [Troubleshooting](#troubleshooting) -- [Release Notes](#release-notes) - - [Changelog](#changelog) - - [Known issues](#known-issues) +- [Release notes](#release-notes) + - [Change Log](#change-log) + - [Known Issues](#known-issues) + - [Report Issues](#report-issues) ## TensorRT-LLM Overview @@ -154,14 +159,14 @@ See the BLOOM [example](examples/bloom) for more details and options regarding t ***3. Run*** -The `summarize.py` script can be used to perform the summarization of articles +The `../summarize.py` script can be used to perform the summarization of articles from the CNN Daily dataset: ```python -python summarize.py --test_trt_llm \ - --hf_model_location ./bloom/560M/ \ - --data_type fp16 \ - --engine_dir ./bloom/560M/trt_engines/fp16/1-gpu/ +python ../summarize.py --test_trt_llm \ + --hf_model_dir ./bloom/560M/ \ + --data_type fp16 \ + --engine_dir ./bloom/560M/trt_engines/fp16/1-gpu/ ``` More details about the script and how to run the BLOOM model can be found in @@ -237,19 +242,26 @@ The list of supported models is: * [Bert](examples/bert) * [Blip2](examples/blip2) * [BLOOM](examples/bloom) -* [ChatGLM](examples/chatglm), including ChatGLM-6B, ChatGLM2-6B, ChatGLM2-6B-32k, ChatGLM3-6B, ChatGLM3-6B-32k +* [ChatGLM](examples/chatglm) * [Falcon](examples/falcon) +* [Flan-T5](examples/enc_dec) * [GPT](examples/gpt) * [GPT-J](examples/gptj) * [GPT-Nemo](examples/gpt) * [GPT-NeoX](examples/gptneox) +* [InternLM](examples/internlm) * [LLaMA](examples/llama) * [LLaMA-v2](examples/llama) +* [Mistral](examples/llama) * [MPT](examples/mpt) * [OPT](examples/opt) +* [Qwen](examples/qwen) +* [Replit Code](examples/mpt) * [SantaCoder](examples/gpt) * [StarCoder](examples/gpt) -* [InternLM](examples/internlm) +* [T5](examples/enc_dec) + +Note: [Encoder-Decoder](examples/enc_dec/) provides general encoder-decoder support that contains many encoder-decoder models such as T5, Flan-T5, etc. We unroll the exact model names in the list above to let users find specific models easiler. ## Performance @@ -311,6 +323,33 @@ may happen. One possible solution is to reduce the amount of memory needed by reducing the maximum batch size, input and output lengths. Another option is to enable plugins, for example: `--use_gpt_attention_plugin`. +* MPI + Slurm + +TensorRT-LLM is a [MPI](https://en.wikipedia.org/wiki/Message_Passing_Interface)-aware package that uses [`mpi4py`](https://mpi4py.readthedocs.io/en/stable/). If you are running scripts in a [Slurm](https://slurm.schedmd.com/) environment, you might encounter interferences: +``` +-------------------------------------------------------------------------- +PMI2_Init failed to initialize. Return code: 14 +-------------------------------------------------------------------------- +-------------------------------------------------------------------------- +The application appears to have been direct launched using "srun", +but OMPI was not built with SLURM's PMI support and therefore cannot +execute. There are several options for building PMI support under +SLURM, depending upon the SLURM version you are using: + + version 16.05 or later: you can use SLURM's PMIx support. This + requires that you configure and build SLURM --with-pmix. + + Versions earlier than 16.05: you must use either SLURM's PMI-1 or + PMI-2 support. SLURM builds PMI-1 by default, or you can manually + install PMI-2. You must then build Open MPI using --with-pmi pointing + to the SLURM PMI library location. + +Please configure as appropriate and try again. +-------------------------------------------------------------------------- +``` +As a rule of thumb, if you are running TensorRT-LLM interactively on a Slurm node, prefix your commands with `mpirun -n 1` to run TensorRT-LLM in a dedicated MPI environment, not the one provided by your Slurm allocation. +For example: `mpirun -n 1 python3 examples/gpt/build.py ...` + ## Release notes * TensorRT-LLM requires TensorRT 9.1.0.4 and 23.08 containers. diff --git a/benchmarks/cpp/README.md b/benchmarks/cpp/README.md index 70c322ef6c85..6b21694d5922 100644 --- a/benchmarks/cpp/README.md +++ b/benchmarks/cpp/README.md @@ -7,18 +7,14 @@ multiple GPUs or multiple nodes with multiple GPUs. ### 1. Build TensorRT-LLM and benchmarking source code -Please follow the [`installation document`](../../../README.md) to build TensorRT-LLM. +Please follow the [`installation document`](../../docs/source/installation.md) to build TensorRT-LLM. + +Note that the benchmarking source code for C++ runtime is not built by default, you can use the argument `--benchmarks` in [`build_wheel.py`](../../scripts/build_wheel.py) to build that. Windows users: Follow the -[`Windows installation document`](../../../windows/README.md) +[`Windows installation document`](../../windows/README.md) instead, and be sure to set DLL paths as specified in -[Extra Steps for C++ Runtime Usage](../../../windows/README.md#extra-steps-for-c-runtime-usage). - -After that, you can build benchmarking source code for C++ runtime -``` -cd cpp/build -make -j benchmarks -``` +[Extra Steps for C++ Runtime Usage](../../windows/README.md#extra-steps-for-c-runtime-usage). ### 2. Launch C++ benchmarking (Fixed BatchSize/InputLen/OutputLen) @@ -59,6 +55,8 @@ mpirun -n 8 ./benchmarks/gptSessionBenchmark \ # [BENCHMARK] batch_size 1 input_length 60 output_length 20 latency(ms) 792.14 ``` +If you want to obtain context and generation logits, you could build an enigne with `--gather_all_token_logits` and run gptSessionBenchmark with `--print_all_logits`. This will print a large number of logit values and has a certain impact on performance. + *Please note that the expected outputs in that document are only for reference, specific performance numbers depend on the GPU you're using.* ### 3. Launch Batch Manager benchmarking (Inflight/V1 batching) diff --git a/benchmarks/cpp/gptSessionBenchmark.cpp b/benchmarks/cpp/gptSessionBenchmark.cpp index aff5bf0e3c1c..845101e2ae09 100644 --- a/benchmarks/cpp/gptSessionBenchmark.cpp +++ b/benchmarks/cpp/gptSessionBenchmark.cpp @@ -18,6 +18,7 @@ #include "tensorrt_llm/plugins/api/tllmPlugin.h" #include "tensorrt_llm/runtime/gptJsonConfig.h" #include "tensorrt_llm/runtime/gptSession.h" +#include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/memoryCounters.h" #include "tensorrt_llm/runtime/tllmLogger.h" @@ -37,7 +38,7 @@ namespace void benchmarkGptSession(std::string const& modelName, std::filesystem::path const& dataPath, std::vector const& batchSizes, int beamWidth, std::vector> const& inOutLen, std::shared_ptr const& logger, int warmUp, int numRuns, int duration, - GptSession::Config& sessionConfig, bool cudaGraphMode) + GptSession::Config& sessionConfig, bool cudaGraphMode, bool printAllLogits) { std::string modelNameHyphen = modelName; @@ -60,7 +61,6 @@ void benchmarkGptSession(std::string const& modelName, std::filesystem::path con SamplingConfig samplingConfig{beamWidth}; samplingConfig.temperature = std::vector{1.0f}; - samplingConfig.minLength = std::vector{1}; samplingConfig.randomSeed = std::vector{42ull}; samplingConfig.topK = std::vector{1}; samplingConfig.topP = std::vector{0.0f}; @@ -77,6 +77,7 @@ void benchmarkGptSession(std::string const& modelName, std::filesystem::path con auto const maxNewTokens = inOut[1]; sessionConfig.maxSequenceLength = maxInputLength + maxNewTokens; + samplingConfig.minLength = std::vector{maxNewTokens}; GptSession session{sessionConfig, modelConfig, worldConfig, enginePath.string(), logger}; @@ -102,6 +103,7 @@ void benchmarkGptSession(std::string const& modelName, std::filesystem::path con // copy inputs and wrap into shared_ptr GenerationInput::TensorPtr inputIds; std::vector inputsHost(batchSize * maxInputLength, padId); + if (inputPacked) { inputIds = bufferManager.copyFrom( @@ -123,6 +125,17 @@ void benchmarkGptSession(std::string const& modelName, std::filesystem::path con bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32), bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32)}; + if (session.getModelConfig().computeContextLogits()) + { + generationOutput.contextLogits + = bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kFLOAT); + } + if (session.getModelConfig().computeGenerationLogits()) + { + generationOutput.generationLogits + = bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kFLOAT); + bufferManager.setZero(*generationOutput.generationLogits); + } TLLM_LOG_INFO(memoryCounter.toString()); for (auto r = 0; r < warmUp; ++r) @@ -168,6 +181,30 @@ void benchmarkGptSession(std::string const& modelName, std::filesystem::path con "%.2f\n", batchSize, maxInputLength, maxNewTokens, averageLatency, tokensPerSec); } + + // logits are store in last rank + if (worldConfig.getRank() == worldConfig.getSize() - 1) + { + if (session.getModelConfig().computeContextLogits() && printAllLogits) + { + std::cout << "generationOutput.contextLogits.shape: " + << generationOutput.contextLogits->getShape() + << std::endl; // (batchsize, prompt_len, vocabsize) + std::cout << "generationOutput.contextLogits" << *generationOutput.contextLogits << std::endl; + } + + if (session.getModelConfig().computeGenerationLogits() && printAllLogits) + { + std::cout << "generationOutput.generationLogits.shape: " + << generationOutput.generationLogits->getShape() + << std::endl; // (batchsize, beamwidth, maxNewTokens-1, vocabsize) + generationOutput.generationLogits->reshape(ITensor::makeShape({batchSize * beamWidth, + maxNewTokens - 1, modelConfig.getVocabSizePadded(worldConfig.getSize())})); + + std::cout << "generationOutput.generationLogits: " << *generationOutput.generationLogits + << std::endl; + } + } } catch (std::runtime_error& e) { @@ -231,6 +268,7 @@ int main(int argc, char* argv[]) "kv_cache_free_gpu_mem_fraction", "K-V Cache Free Gpu Mem Fraction.", cxxopts::value()); options.add_options()("enable_cuda_graph", "Execute GPT session with CUDA graph."); + options.add_options()("print_all_logits", "Print all context and generation logits."); auto result = options.parse(argc, argv); @@ -328,6 +366,7 @@ int main(int argc, char* argv[]) // Argument: Enable CUDA graph auto enableCudaGraph = result.count("enable_cuda_graph") > 0; + auto printAllLogits = result.count("print_all_logits") > 0; initTrtLlmPlugins(logger.get()); @@ -335,7 +374,7 @@ int main(int argc, char* argv[]) { benchmarkGptSession(result["model"].as(), result["engine_dir"].as(), batchSizes, beamWidth, inOutLen, logger, result["warm_up"].as(), result["num_runs"].as(), - result["duration"].as(), sessionConfig, enableCudaGraph); + result["duration"].as(), sessionConfig, enableCudaGraph, printAllLogits); } catch (const std::exception& e) { diff --git a/benchmarks/python/allowed_configs.py b/benchmarks/python/allowed_configs.py index 586845631849..961d40879cc4 100644 --- a/benchmarks/python/allowed_configs.py +++ b/benchmarks/python/allowed_configs.py @@ -353,7 +353,7 @@ class ModelConfig(BaseModel): builder_opt=None, )), "chatglm_6b": - ModelConfig(name="chatglm-6b", + ModelConfig(name="chatglm_6b", family="chatglm", benchmark_type="gpt", build_config=BuildConfig( @@ -370,7 +370,7 @@ class ModelConfig(BaseModel): remove_input_padding=False, )), "chatglm2_6b": - ModelConfig(name="chatglm2-6b", + ModelConfig(name="chatglm2_6b", family="chatglm2", benchmark_type="gpt", build_config=BuildConfig( @@ -387,7 +387,7 @@ class ModelConfig(BaseModel): remove_input_padding=False, )), "chatglm3_6b": - ModelConfig(name="chatglm3-6b", + ModelConfig(name="chatglm3_6b", family="chatglm3", benchmark_type="gpt", build_config=BuildConfig( diff --git a/benchmarks/python/gpt_benchmark.py b/benchmarks/python/gpt_benchmark.py index 88ec9f744396..7cd4f6e0b732 100644 --- a/benchmarks/python/gpt_benchmark.py +++ b/benchmarks/python/gpt_benchmark.py @@ -143,7 +143,7 @@ def __init__(self, quant_mode=self.quant_mode, use_custom_all_reduce=self.enable_custom_all_reduce, ) - if model_name == 'chatglm-6b': + if model_name == 'chatglm_6b': self.sampling_config = tensorrt_llm.runtime.SamplingConfig( end_id=130005, pad_id=3, @@ -152,16 +152,7 @@ def __init__(self, top_p=top_p) self.decoder = tensorrt_llm.runtime.ChatGLMGenerationSession( model_config, engine_buffer, self.runtime_mapping) - elif model_name == 'chatglm2-6b': - self.sampling_config = tensorrt_llm.runtime.SamplingConfig( - end_id=2, - pad_id=0, - num_beams=num_beams, - top_k=top_k, - top_p=top_p) - self.decoder = tensorrt_llm.runtime.GenerationSession( - model_config, engine_buffer, self.runtime_mapping) - elif model_name == 'chatglm3-6b': + elif model_name in ['chatglm2_6b', 'chatglm3_6b']: self.sampling_config = tensorrt_llm.runtime.SamplingConfig( end_id=2, pad_id=0, @@ -402,7 +393,7 @@ def build(self): apply_query_key_layer_scaling=builder_config. apply_query_key_layer_scaling, quant_mode=self.quant_mode, - model_version="1") + model_name="chatglm_6b") elif family == "chatglm2": tensorrt_llm_model = tensorrt_llm.models.ChatGLMHeadModel( num_layers=self.num_layers, @@ -418,7 +409,7 @@ def build(self): apply_query_key_layer_scaling=builder_config. apply_query_key_layer_scaling, quant_mode=self.quant_mode, - model_version="2") + model_name="chatglm2_6b") elif family == "chatglm3": tensorrt_llm_model = tensorrt_llm.models.ChatGLMHeadModel( num_layers=self.num_layers, @@ -434,7 +425,7 @@ def build(self): apply_query_key_layer_scaling=builder_config. apply_query_key_layer_scaling, quant_mode=self.quant_mode, - model_version="3") + model_name="chatglm3_6b") elif family == "bloom": tensorrt_llm_model = tensorrt_llm.models.BloomForCausalLM( num_layers=self.num_layers, @@ -458,6 +449,7 @@ def build(self): max_position_embeddings=self.n_positions, dtype=kv_dtype, bias=self.bias, + quant_mode=self.quant_mode, use_alibi=self.use_alibi, new_decoder_architecture=self.new_decoder_architecture, parallel_attention=self.parallel_attention, diff --git a/benchmarks/python/mem_monitor.py b/benchmarks/python/mem_monitor.py index 132e23c5b1f5..24f28dfe9392 100644 --- a/benchmarks/python/mem_monitor.py +++ b/benchmarks/python/mem_monitor.py @@ -22,7 +22,7 @@ def get_memory_info(handle): version=pynvml.nvmlMemory_v2) total = round(mem_info.total / 1024 / 1024 / 1024, 2) used = round(mem_info.used / 1024 / 1024 / 1024, 2) - free = round(mem_info.used / 1024 / 1024 / 1024, 2) + free = round(mem_info.free / 1024 / 1024 / 1024, 2) return total, used, free diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 98ad371c7b73..7cb341301e4f 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -237,6 +237,24 @@ if(WIN32) set(CMAKE_CXX_FLAGS "/DNOMINMAX ${CMAKE_CXX_FLAGS}") endif() +if((MSVC)) + if((MSVC_VERSION GREATER_EQUAL 1914)) + # MSVC does not apply the correct __cplusplus version per the C++ standard + # by default. This is required for compiling CUTLASS 3.0 kernels on windows + # with C++-17 constexpr enabled. The 2017 15.7 MSVC adds /Zc:__cplusplus to + # set __cplusplus to 201703 with std=c++17. See + # https://learn.microsoft.com/en-us/cpp/build/reference/zc-cplusplus for + # more info. + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /Zc:__cplusplus") + set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler /Zc:__cplusplus") + else() + message( + FATAL_ERROR + "Build is only supported with Visual Studio 2017 version 15.7 or higher" + ) + endif() +endif() + set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-extended-lambda") set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr") if(FAST_MATH) diff --git a/cpp/include/tensorrt_llm/batch_manager/GptManager.h b/cpp/include/tensorrt_llm/batch_manager/GptManager.h index 945095fe67b2..2ed4e2c45cf7 100644 --- a/cpp/include/tensorrt_llm/batch_manager/GptManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/GptManager.h @@ -121,10 +121,13 @@ class GptManager inline static const std::string kMinLengthTensorName_ = "min_length"; inline static const std::string kPresencePenaltyTensorName_ = "presence_penalty"; inline static const std::string kRandomSeedTensorName_ = "random_seed"; + inline static const std::string kReturnLogProbsTensorName_ = "return_log_probs"; inline static const std::string kPromptEmbeddingTableName_ = "prompt_embedding_table"; inline static const std::string kPromptVocabSizeName_ = "prompt_vocab_size"; inline static const std::string kOutputIdsTensorName_ = "output_ids"; inline static const std::string kSequenceLengthTensorName_ = "sequence_length"; + inline static const std::string kLogProbsTensorName_ = "output_log_probs"; + inline static const std::string kCumLogProbsTensorName_ = "cum_log_probs"; std::shared_ptr mLogger{}; }; diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 6bccd129b1fd..3f913e88bbce 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -309,8 +309,8 @@ class KVCacheManager } [[nodiscard]] static SizeType getMaxNumTokens(KvCacheConfig const& config, nvinfer1::DataType dtype, - tensorrt_llm::runtime::GptModelConfig const& modelConfig, - tensorrt_llm::runtime::WorldConfig const& worldConfig); + tensorrt_llm::runtime::GptModelConfig const& modelConfig, tensorrt_llm::runtime::WorldConfig const& worldConfig, + runtime::BufferManager const& bufferManager); private: void resetBlockPointers(SizeType batchSlotIdx, SizeType beamWidth); diff --git a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h index c577151f5aa7..d8db278adb02 100644 --- a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h +++ b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h @@ -43,6 +43,7 @@ class LlmRequest using TokenIdType = runtime::TokenIdType; using RequestIdType = std::uint64_t; using BeamTokens = std::vector>; + using VecLogProbs = std::vector; using TensorPtr = runtime::ITensor::SharedPtr; LlmRequest(RequestIdType requestId, SizeType maxNewTokens, std::shared_ptr> input_tokens, @@ -50,7 +51,7 @@ class LlmRequest std::optional padId = std::nullopt, std::optional embeddingBias = std::nullopt, std::optional badWordsList = std::nullopt, std::optional stopWordsList = std::nullopt, std::optional promptEmbeddingTable = std::nullopt, - std::optional promptVocabSize = std::nullopt) + std::optional promptVocabSize = std::nullopt, bool returnLogProbs = false) : mRequestId(requestId) , mPromptLen(input_tokens->size()) , mMaxNewTokens(maxNewTokens) @@ -60,11 +61,15 @@ class LlmRequest , mEndId(endId) , mPadId(padId) , mBatchSlot(-1) + , mOrigPromptLen(input_tokens->size()) , mEmbeddingBias(embeddingBias) , mBadWordsList(badWordsList) , mStopWordsList(stopWordsList) , mPromptEmbeddingTable(promptEmbeddingTable) , mPromptVocabSize(promptVocabSize) + , mReturnLogProbs(returnLogProbs) + , mLogProbs(samplingConfig.beamWidth) + , mCumLogProbs(samplingConfig.beamWidth) { mMaxSentTokenPos = mPromptLen - 1; // Scatter the input tokens to other beam @@ -168,17 +173,29 @@ class LlmRequest // As a temporary solution, we currently reset the tokens to the prompt if (mSamplingConfig.beamWidth > 1) { - for (auto& beamTokens : *mTokens) + for (std::size_t beam = 0; beam < mTokens->size(); ++beam) { + auto& beamTokens = mTokens->at(beam); beamTokens.resize(mPromptLen); + if (mReturnLogProbs) + { + mLogProbs.at(beam).clear(); + } } } else { SizeType newPromptLen = std::min(maxInputLen, mPromptLen + getMaxNumGeneratedTokens()); - for (auto& beamTokens : *mTokens) + for (std::size_t beam = 0; beam < mTokens->size(); ++beam) { + auto& beamTokens = mTokens->at(beam); beamTokens.resize(newPromptLen); + + if (mReturnLogProbs) + { + auto& logProb = mLogProbs.at(beam); + logProb.resize(newPromptLen - mPromptLen); + } } mMaxNewTokens -= (newPromptLen - mPromptLen); mPromptLen = newPromptLen; @@ -187,16 +204,16 @@ class LlmRequest mBatchSlot = -1; } - /// @brief Get the maximum position of the tokens returned to the client. Use to ensure we don't return to client - /// duplicated token positions. + /// @brief Get the maximum position of the tokens returned to the client. Use to ensure we don't return to + /// client duplicated token positions. /// @return The maximum position of the tokens sent to the client SizeType getMaxSentTokenPos() const { return mMaxSentTokenPos; } - /// @brief Sets the maximum position of the tokens returned to the client. Use to ensure we don't return to client - /// duplicated token positions. + /// @brief Sets the maximum position of the tokens returned to the client. Use to ensure we don't return to + /// client duplicated token positions. /// @param pos The maximum position void setMaxSentTokenPos(SizeType pos) { @@ -243,6 +260,42 @@ class LlmRequest return mStopWordsList; } + bool returnLogProbs() const + { + return mReturnLogProbs; + } + + std::vector const& getLogProbs() const + { + return mLogProbs; + } + + VecLogProbs const& getLogProbs(SizeType beam) const + { + return mLogProbs.at(beam); + } + + void setLogProbs(VecLogProbs const& logProbs, SizeType beam) + { + mLogProbs.at(beam).resize(mPromptLen - mOrigPromptLen); + mLogProbs.at(beam).insert(mLogProbs.at(beam).end(), logProbs.begin(), logProbs.end()); + } + + VecLogProbs const& getCumLogProbs() const + { + return mCumLogProbs; + } + + void setCumLogProb(float cumLogProb, SizeType beam) + { + mCumLogProbs.at(beam) = cumLogProb; + } + + SizeType getOrigPromptLen() const + { + return mOrigPromptLen; + } + RequestIdType mRequestId; SizeType mPromptLen; SizeType mMaxNewTokens; @@ -255,6 +308,7 @@ class LlmRequest SizeType mBatchSlot; private: + SizeType mOrigPromptLen; std::shared_ptr mTokens; SizeType mMaxSentTokenPos; @@ -264,6 +318,11 @@ class LlmRequest std::optional mPromptEmbeddingTable; std::optional mPromptVocabSize; + + bool mReturnLogProbs; + + std::vector mLogProbs; // [beamSize, seqLen] + VecLogProbs mCumLogProbs; // [beamSize] }; } // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/runtime/bufferManager.h b/cpp/include/tensorrt_llm/runtime/bufferManager.h index 08f4035a3d19..379945f2e052 100644 --- a/cpp/include/tensorrt_llm/runtime/bufferManager.h +++ b/cpp/include/tensorrt_llm/runtime/bufferManager.h @@ -147,9 +147,33 @@ class BufferManager //! \brief Get the underlying cuda stream. [[nodiscard]] CudaStream const& getStream() const; + //! \brief The current size of the memory reserved by the memory pool. + [[nodiscard]] std::size_t memoryPoolReserved() const; + + //! \brief The current size of the memory used by the memory pool. + [[nodiscard]] std::size_t memoryPoolUsed() const; + + //! \brief The current size of the memory free in the memory pool. + [[nodiscard]] std::size_t memoryPoolFree() const; + + //! \brief Try to trim the memory reserved by the pool to `size` bytes. This synchronizes implicitly with the + //! stream. + void memoryPoolTrimTo(std::size_t size); + private: void static initMemoryPool(int device); + std::size_t static memoryPoolReserved(int device); + + std::size_t static memoryPoolUsed(int device); + + std::size_t static memoryPoolFree(int device) + { + return memoryPoolReserved(device) - memoryPoolUsed(device); + } + + void static memoryPoolTrimTo(int device, std::size_t size); + CudaStreamPtr mStream; }; diff --git a/cpp/include/tensorrt_llm/runtime/decodingOutput.h b/cpp/include/tensorrt_llm/runtime/decodingOutput.h index a52b87f99856..eda27b4bb273 100644 --- a/cpp/include/tensorrt_llm/runtime/decodingOutput.h +++ b/cpp/include/tensorrt_llm/runtime/decodingOutput.h @@ -70,9 +70,9 @@ class DecodingOutput TensorPtr finished; // [batchSize, beamWidth], mandatory in beam search and to determine whether to stop // according to DecodingInput.sequenceLimitLength, on gpu TensorPtr finishedSum; // [1], the sum of finished sequences, in pinned memory - TensorPtr logProbs; // [maxNewTokens, batchSize, beamWidth], must be float*, on gpu // mandatory parameters for beam search + TensorPtr logProbs; // [batchSize, beamWidth, maxSeqLen], must be float*, on gpu TensorPtr cumLogProbs; // [batchSize, beamWidth], optional for sampling, on gpu TensorPtr parentIds; // [batchSize, beamWidth, maxSeqLen], on gpu TensorPtr lengths; // [batchSize, beamWidth], total sequence lengths including padding, on gpu diff --git a/cpp/include/tensorrt_llm/runtime/generationOutput.h b/cpp/include/tensorrt_llm/runtime/generationOutput.h index 33b7d7272e2f..9dbb1d536183 100644 --- a/cpp/include/tensorrt_llm/runtime/generationOutput.h +++ b/cpp/include/tensorrt_llm/runtime/generationOutput.h @@ -46,8 +46,10 @@ class GenericGenerationOutput TensorPtr lengths; // [batchSize, beamWidth] // optional parameters - TensorPtr logProbs; // [request_output_length, batch_size * beam_width], must be float*, on gpu - TensorPtr contextLogits; // [batch_size, max_input_length, vocab_size_padded] + TensorPtr cumLogProbs; // [batchSize, beamWidth], must be float*, on gpu + TensorPtr logProbs; // [batchSize, beamWidth, maxInputLength + maxNewTokens], must be float*, on gpu + TensorPtr contextLogits; // [batch_size, max_input_length, vocab_size_padded] + TensorPtr generationLogits; // [batch_size, beam_width, max_output_length-1, vocab_size_padded] // callbacks Callback onTokenGenerated; diff --git a/cpp/include/tensorrt_llm/runtime/gptDecoder.h b/cpp/include/tensorrt_llm/runtime/gptDecoder.h index c7bb6d2ce7b5..139c6e9c6134 100644 --- a/cpp/include/tensorrt_llm/runtime/gptDecoder.h +++ b/cpp/include/tensorrt_llm/runtime/gptDecoder.h @@ -45,14 +45,15 @@ class IGptDecoder public: virtual ~IGptDecoder() = default; - virtual void setup(SamplingConfig const& samplingConfig, size_t batchSize) = 0; + virtual void setup(SamplingConfig const& samplingConfig, size_t batchSize, SizeType maxSequenceLength) = 0; virtual bool forward(DecodingOutput& output, DecodingInput const& input) = 0; virtual void forwardAsync(DecodingOutput& output, DecodingInput const& input) = 0; - static void gatherTree(ITensor& finalOutputIds, DecodingOutput const& decodingOutput, - DecodingInput const& decodingInput, BufferManager const& manager); + virtual void gatherTree(ITensor& finalOutputIds, DecodingOutput const& decodingOutput, + DecodingInput const& decodingInput, BufferManager const& manager) + = 0; static std::unique_ptr create( nvinfer1::DataType dtype, size_t vocabSize, size_t vocabSizePadded, BufferManager::CudaStreamPtr const& stream); @@ -64,19 +65,27 @@ class GptDecoder : public virtual IGptDecoder public: using CudaStreamPtr = BufferManager::CudaStreamPtr; + using TensorPtr = std::shared_ptr; GptDecoder(size_t vocabSize, size_t vocabSizePadded, CudaStreamPtr const& stream); - void setup(SamplingConfig const& samplingConfig, size_t batchSize) override; + void setup(SamplingConfig const& samplingConfig, size_t batchSize, SizeType maxSequenceLength) override; bool forward(DecodingOutput& output, DecodingInput const& input) override; void forwardAsync(DecodingOutput& output, DecodingInput const& input) override; + void gatherTree(ITensor& finalOutputIds, DecodingOutput const& decodingOutput, DecodingInput const& decodingInput, + BufferManager const& manager) override; + private: BufferManager mManager; + common::CudaAllocator mAllocator; std::shared_ptr> mDynamicDecodeLayer; + + TensorPtr mLogProbsTiled; // Buffer used to store the transpose of the logProbs. Needed because the kernels have + // been written to use that shape. }; inline std::unique_ptr IGptDecoder::create( diff --git a/cpp/include/tensorrt_llm/runtime/gptDecoderBatch.h b/cpp/include/tensorrt_llm/runtime/gptDecoderBatch.h index 7ea7748805eb..8919d0aff677 100644 --- a/cpp/include/tensorrt_llm/runtime/gptDecoderBatch.h +++ b/cpp/include/tensorrt_llm/runtime/gptDecoderBatch.h @@ -20,6 +20,7 @@ #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/cudaEvent.h" #include "tensorrt_llm/runtime/cudaStream.h" +#include "tensorrt_llm/runtime/generationOutput.h" #include "tensorrt_llm/runtime/gptDecoder.h" #include "tensorrt_llm/runtime/iGptDecoderBatch.h" #include "tensorrt_llm/runtime/iTensor.h" @@ -51,7 +52,8 @@ class GptDecoderBatch : public IGptDecoderBatch void newRequest( SizeType batchIdx, decoder_batch::Request const& request, SamplingConfig const& samplingConfig) override; - void newBatch(GenerationInput const& inputs, SamplingConfig const& samplingConfig) override; + void newBatch( + GenerationInput const& inputs, GenerationOutput const& outputs, SamplingConfig const& samplingConfig) override; TokenPtr forwardAsync(decoder_batch::Output& output, decoder_batch::Input const& input) override; @@ -85,14 +87,10 @@ class GptDecoderBatch : public IGptDecoderBatch //! @brief Gather final beam search results for request `batchIdx`. //! Result will only be available after event returned. - //! @returns [maxBeamWidth, maxInputLength + maxNewTokens], contains input token ids and generated token ids without - //! padding for request `batchIdx`, on gpu - [[nodiscard]] std::tuple getFinalOutputIds(SizeType batchIdx) const override; + [[nodiscard]] CudaEvent finalize(SizeType batchIdx) const; //! @brief Gather final beam search results for all requests. - //! @returns [batchSize, maxBeamWidth, maxInputLength + maxNewTokens], contains input token ids and generated token - //! ids without padding, on gpu - [[nodiscard]] TensorPtr getFinalOutputIds() const override; + void finalize() const override; //! @returns [batchSize, maxBeamWidth, maxInputLength + maxNewTokens], contains parent ids collected during beam //! search without padding, on gpu @@ -119,6 +117,28 @@ class GptDecoderBatch : public IGptDecoderBatch return ITensor::slice(mJointDecodingOutput->cumLogProbs, 0, mActualBatchSize); } + //! @returns [maxBeamWidth], cumulative log probabilities (per beam), on gpu + [[nodiscard]] TensorPtr getCumLogProbs(SizeType batchIdx) const + { + auto tensor = ITensor::slice(mJointDecodingOutput->cumLogProbs, batchIdx, 1); + tensor->squeeze(0); + return tensor; + } + + //! @returns [batchSize, maxBeamWidth, maxSequenceLength], log probabilities (per beam), on gpu + [[nodiscard]] TensorPtr getLogProbs() const override + { + return ITensor::slice(mJointDecodingOutput->logProbs, 0, mActualBatchSize); + } + + //! @returns [maxBeamWidth, maxSequenceLength], log probabilities (per beam), on gpu + [[nodiscard]] TensorPtr getLogProbs(SizeType batchIdx) const + { + auto tensor = ITensor::slice(mJointDecodingOutput->logProbs, batchIdx, 1); + tensor->squeeze(0); + return tensor; + } + //! @returns [batchSize, maxBeamWidth], tokens generated in last forward pass, on gpu [[nodiscard]] TensorPtr getNewTokens() const override { diff --git a/cpp/include/tensorrt_llm/runtime/gptModelConfig.h b/cpp/include/tensorrt_llm/runtime/gptModelConfig.h index 325a7246629d..c0dba9ba3a73 100644 --- a/cpp/include/tensorrt_llm/runtime/gptModelConfig.h +++ b/cpp/include/tensorrt_llm/runtime/gptModelConfig.h @@ -50,6 +50,7 @@ class GptModelConfig , mMaxOutputLen(0) , mMaxNumTokens(std::nullopt) , mComputeContextLogits(false) + , mComputeGenerationLogits(false) , mModelVariant(ModelVariant::kGpt) , mUseCustomAllReduce(false) , mMaxPromptEmbeddingTableSize(0) @@ -222,6 +223,16 @@ class GptModelConfig mComputeContextLogits = computeContextLogits; } + [[nodiscard]] bool constexpr computeGenerationLogits() const noexcept + { + return mComputeGenerationLogits; + } + + void constexpr computeGenerationLogits(bool computeGenerationLogits) noexcept + { + mComputeGenerationLogits = computeGenerationLogits; + } + [[nodiscard]] ModelVariant getModelVariant() const { return mModelVariant; @@ -260,6 +271,7 @@ class GptModelConfig std::optional mMaxNumTokens; bool mComputeContextLogits; + bool mComputeGenerationLogits; ModelVariant mModelVariant; bool mUseCustomAllReduce; diff --git a/cpp/include/tensorrt_llm/runtime/gptSession.h b/cpp/include/tensorrt_llm/runtime/gptSession.h index f6ef791452d1..50fda8ef5d76 100644 --- a/cpp/include/tensorrt_llm/runtime/gptSession.h +++ b/cpp/include/tensorrt_llm/runtime/gptSession.h @@ -63,6 +63,8 @@ class GptSession { using KvCacheManager = batch_manager::kv_cache_manager::KVCacheManager; using KvCacheConfig = batch_manager::kv_cache_manager::KvCacheConfig; + using TensorPtr = runtime::ITensor::SharedPtr; + using TokenGeneratedCallback = std::function; public: using LoggerPtr = std::shared_ptr; @@ -108,7 +110,7 @@ class GptSession [[nodiscard]] nvinfer1::ILogger& getLogger() const; - [[nodiscard]] BufferManager& getBufferManager() const; + [[nodiscard]] BufferManager const& getBufferManager() const; [[nodiscard]] GptModelConfig const& getModelConfig() const { @@ -133,8 +135,9 @@ class GptSession return !mCudaGraphInstances.empty(); } - void generateBatched(GenerationOutput& outputs, std::vector const& microBatches, - SamplingConfig const& samplingConfig); + void generateBatched(std::vector& microBatchesOutputs, + std::vector const& microBatchesInputs, SamplingConfig const& samplingConfig, + TokenGeneratedCallback const& onTokenGenerated); void setup(Config const& sessionConfig); @@ -148,9 +151,9 @@ class GptSession void executeContextStep(std::vector const& microBatches, std::vector const& microBatchOffsets, KvCacheManager const* kvCacheManager); - SizeType executeGenerationStep(SizeType step, std::vector const& microBatches, - std::vector const& microBatchOffsets, KvCacheManager* kvCacheManager, - std::vector& microBatchesFinished); + SizeType executeGenerationStep(SizeType step, std::vector const& microBatchesInputs, + std::vector& microBatchesOutputs, std::vector const& microBatchOffsets, + KvCacheManager* kvCacheManager, std::vector& microBatchesFinished); //! @brief Execute decoder on last PP rank, receive decoder output on other PP ranks. void decoderStepAsync(SizeType decoderStep, SizeType microBatchId); @@ -158,17 +161,17 @@ class GptSession //! @brief Synchronize with the decoder and return the `shouldStop` flag. bool shouldStopSync(SizeType batchSize, SizeType beamWidth, SizeType microBatchId); - //! @brief Collect final output ids on last PP rank and send them to first PP rank. + //! @brief Collect final output ids and log probs on last PP rank and send them to first PP rank. //! @details Receives are asynchronous on host, so synchronization is required before access. - void finalizeOutputIds(SizeType microBatchId); + void finalize(SizeType microBatchId); void kvCacheAddSequences(SizeType beamWidth, SizeType microBatchId, SizeType firstBatchIdx); //! @brief Populate outputIds and return reference to newTokens tensor - ITensor::SharedPtr initDecoder(ITensor& outputIds, GenerationInput const& inputs, + ITensor::SharedPtr initDecoder(ITensor& outputIds, GenerationInput const& inputs, GenerationOutput const& outputs, SamplingConfig const& samplingConfig, SizeType microBatchId) const; - std::function createOnTokenGeneratedCallback(GenerationOutput& outputs); + TokenGeneratedCallback createOnTokenGeneratedCallback(GenerationOutput& outputs); class CudaGraphExecutor { diff --git a/cpp/include/tensorrt_llm/runtime/iBuffer.h b/cpp/include/tensorrt_llm/runtime/iBuffer.h index a71c587821b5..27892bbd8bde 100644 --- a/cpp/include/tensorrt_llm/runtime/iBuffer.h +++ b/cpp/include/tensorrt_llm/runtime/iBuffer.h @@ -68,84 +68,108 @@ struct MemoryTypeString //! \brief For converting a TensorRT data type to a C++ data type. template -struct CppDataType +struct DataTypeTraits { }; template <> -struct CppDataType +struct DataTypeTraits { using type = float; + static char constexpr name[] = "float"; + static auto constexpr size = sizeof(type); }; template <> -struct CppDataType +struct DataTypeTraits { using type = half; + static char constexpr name[] = "half"; + static auto constexpr size = sizeof(type); }; template <> -struct CppDataType +struct DataTypeTraits { using type = std::int8_t; + static char constexpr name[] = "int8"; + static auto constexpr size = sizeof(type); }; template <> -struct CppDataType +struct DataTypeTraits { using type = std::int32_t; + static char constexpr name[] = "int32"; + static auto constexpr size = sizeof(type); }; template <> -struct CppDataType +struct DataTypeTraits { using type = std::int64_t; + static char constexpr name[] = "int64"; + static auto constexpr size = sizeof(type); }; template <> -struct CppDataType +struct DataTypeTraits { using type = std::uint32_t; + static char constexpr name[] = "uint32"; + static auto constexpr size = sizeof(type); }; template <> -struct CppDataType +struct DataTypeTraits { using type = std::uint64_t; + static char constexpr name[] = "uint64"; + static auto constexpr size = sizeof(type); }; template -struct CppDataType +struct DataTypeTraits { using type = bool; + static char constexpr name[] = "bool"; + static auto constexpr size = sizeof(type); }; template -struct CppDataType +struct DataTypeTraits { using type = std::uint8_t; + static char constexpr name[] = "uint8"; + static auto constexpr size = sizeof(type); }; #ifdef ENABLE_BF16 template <> -struct CppDataType +struct DataTypeTraits { using type = __nv_bfloat16; + static char constexpr name[] = "bfloat16"; + static auto constexpr size = sizeof(type); }; #endif #ifdef ENABLE_FP8 template <> -struct CppDataType +struct DataTypeTraits { using type = __nv_fp8_e4m3; + static char constexpr name[] = "fp8"; + static auto constexpr size = sizeof(type); }; #endif template -struct CppDataType +struct DataTypeTraits { - using type = typename CppDataType::type*; + using type = typename DataTypeTraits::type*; + static char constexpr name[] = "*"; + static auto constexpr size = sizeof(type); }; //! \brief A wrapper around `nvinfer1::DataType` that provides a support for pointer types. @@ -377,11 +401,15 @@ class IBuffer //! [[nodiscard]] virtual DataType getDataType() const = 0; + virtual char const* getDataTypeName() const; + //! //! \brief Returns the memory type of the buffer. //! [[nodiscard]] virtual MemoryType getMemoryType() const = 0; + virtual char const* getMemoryTypeName() const; + //! //! \brief Resizes the buffer. This is a no-op if the new size is smaller than or equal to the current capacity. //! diff --git a/cpp/include/tensorrt_llm/runtime/iGptDecoderBatch.h b/cpp/include/tensorrt_llm/runtime/iGptDecoderBatch.h index 667c1e58f109..d41bf351b170 100644 --- a/cpp/include/tensorrt_llm/runtime/iGptDecoderBatch.h +++ b/cpp/include/tensorrt_llm/runtime/iGptDecoderBatch.h @@ -39,10 +39,12 @@ class Request using TensorPtr = std::shared_ptr; explicit Request(ConstTensorPtr ids, std::optional maxNewTokens = std::nullopt, - std::optional endId = std::nullopt, std::optional padId = std::nullopt) + std::optional endId = std::nullopt) : ids{std::move(ids)} , maxNewTokens{maxNewTokens} , endId{endId} + , computeCumLogProbs(false) + , computeLogProbs(false) { } @@ -55,6 +57,9 @@ class Request TensorPtr embeddingBias; // [vocabSizePadded], on gpu TensorPtr badWordsList; // [2, badWordsLength], on gpu TensorPtr stopWordsList; // [2, stopWordsLength], on gpu + + bool computeCumLogProbs; // boolean that controls if cumLogProbs should be computed for that request + bool computeLogProbs; // boolean that controls if cumLogProbs should be computed for that request }; class Input : public decoder::Input @@ -128,9 +133,7 @@ class IGptDecoderBatch : public virtual IStatefulGptDecoder //! @brief Gather final beam search results for request `batchIdx`. //! Result will only be available after event returned - //! @returns [maxBeamWidth, maxInputLength + maxNewTokens], contains input token ids and generated token ids without - //! padding for request `batchIdx`, on gpu - virtual std::tuple getFinalOutputIds(SizeType batchIdx) const = 0; + virtual CudaEvent finalize(SizeType batchIdx) const = 0; //! @returns [batchSize, beamWidth], marks finished requests (per beam), on gpu virtual TensorPtr getFinishedBeams() const = 0; @@ -144,6 +147,15 @@ class IGptDecoderBatch : public virtual IStatefulGptDecoder //! @returns [batchSize, beamWidth], cumulative log probabilities (per beam), on gpu virtual TensorPtr getCumLogProbs() const = 0; + //! @returns [beamWidth], cumulative log probabilities (per beam) for request batchIdx, on gpu + virtual TensorPtr getCumLogProbs(SizeType batchIdx) const = 0; + + //! @returns [batchSize, beamWidth, maxSeqLen], log probabilities (per beam), on gpu + virtual TensorPtr getLogProbs() const = 0; + + //! @returns [beamWidth, maxSeqLen], cumulative log probabilities (per beam) for request batchIdx, on gpu + virtual TensorPtr getLogProbs(SizeType batchIdx) const = 0; + virtual TensorPtr getParentIds() const = 0; virtual std::vector getNbSteps() const = 0; diff --git a/cpp/include/tensorrt_llm/runtime/iStatefulGptDecoder.h b/cpp/include/tensorrt_llm/runtime/iStatefulGptDecoder.h index 5f6697173ecf..a7632acebd60 100644 --- a/cpp/include/tensorrt_llm/runtime/iStatefulGptDecoder.h +++ b/cpp/include/tensorrt_llm/runtime/iStatefulGptDecoder.h @@ -18,6 +18,7 @@ #include "tensorrt_llm/runtime/cudaStream.h" #include "tensorrt_llm/runtime/generationInput.h" +#include "tensorrt_llm/runtime/generationOutput.h" #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/samplingConfig.h" @@ -78,7 +79,9 @@ class IStatefulGptDecoder = 0; //! @brief Initialize the decoder with new batch of inputs. - virtual void newBatch(GenerationInput const& inputs, SamplingConfig const& samplingConfig) = 0; + virtual void newBatch( + GenerationInput const& inputs, GenerationOutput const& outputs, SamplingConfig const& samplingConfig) + = 0; //! @brief Run one step for all requests without blocking the host thread. virtual void forwardAsync(decoder::Output& output, decoder::Input const& input) = 0; @@ -94,11 +97,17 @@ class IStatefulGptDecoder } //! @brief Gather final beam search results for all requests. - virtual TensorPtr getFinalOutputIds() const = 0; + virtual void finalize() const = 0; //! @returns [batchSize, beamWidth, maxSequenceLength], all token ids, on gpu virtual TensorPtr getOutputIds() const = 0; + //! @returns [batchSize, maxBeamWidth], cumulative log probabilities (per beam), on gpu + virtual TensorPtr getCumLogProbs() const = 0; + + //! @returns [batchSize, maxBeamWidth, maxSequenceLength], log probabilities (per beam), on gpu + virtual TensorPtr getLogProbs() const = 0; + //! @returns [batchSize, beamWidth], latests generated tokens (per beam), on gpu virtual TensorPtr getNewTokens() const = 0; diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index 0720649d2c58..ac94fd73f9e4 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/cpp/tensorrt_llm/CMakeLists.txt @@ -38,22 +38,22 @@ add_subdirectory(runtime) set(BATCH_MANAGER_TARGET tensorrt_llm_batch_manager_static) set(BATCH_MANAGER_TARGET_ARCH "unknown") -execute_process( - COMMAND grep -oP "(?<=^ID=).+" /etc/os-release - COMMAND tr -d "\"" - COMMAND tr -d "\n" - RESULT_VARIABLE _OS_ID_SUCCESS - OUTPUT_VARIABLE OS_ID) -execute_process( - COMMAND grep -oP "(?<=^VERSION_ID=).+" /etc/os-release - COMMAND tr -d "\"" - COMMAND tr -d "\n" - RESULT_VARIABLE _OS_VERSION_ID_SUCCESS - OUTPUT_VARIABLE OS_VERSION_ID) -message(STATUS "Operating System: ${OS_ID}, ${OS_VERSION_ID}") - message(STATUS "CMAKE_SYSTEM_PROCESSOR: ${CMAKE_SYSTEM_PROCESSOR}") if(NOT WIN32) # Linux + execute_process( + COMMAND grep -oP "(?<=^ID=).+" /etc/os-release + COMMAND tr -d "\"" + COMMAND tr -d "\n" + RESULT_VARIABLE _OS_ID_SUCCESS + OUTPUT_VARIABLE OS_ID) + execute_process( + COMMAND grep -oP "(?<=^VERSION_ID=).+" /etc/os-release + COMMAND tr -d "\"" + COMMAND tr -d "\n" + RESULT_VARIABLE _OS_VERSION_ID_SUCCESS + OUTPUT_VARIABLE OS_VERSION_ID) + message(STATUS "Operating System: ${OS_ID}, ${OS_VERSION_ID}") + if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64") set(BATCH_MANAGER_TARGET_ARCH "x86_64-linux-gnu") elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64") diff --git a/cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/libtensorrt_llm_batch_manager_static.a b/cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/libtensorrt_llm_batch_manager_static.a deleted file mode 100644 index d62103f67441..000000000000 --- a/cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/libtensorrt_llm_batch_manager_static.a +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7a3ec9a8760d7b8ace53e420572aeb1b3607effc92fd56e13351fa4cbddbbb37 -size 1646420 diff --git a/cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/libtensorrt_llm_batch_manager_static.pre_cxx11.a b/cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/libtensorrt_llm_batch_manager_static.pre_cxx11.a deleted file mode 100644 index fd17c5bc6fb7..000000000000 --- a/cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/libtensorrt_llm_batch_manager_static.pre_cxx11.a +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:114348de9f6d1b3fa147f4fbccede10b7dbe13da6c5c86e968bb56bf05f9ec5a -size 1657852 diff --git a/cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/version.txt b/cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/version.txt deleted file mode 100644 index f381e540149f..000000000000 --- a/cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/version.txt +++ /dev/null @@ -1,3 +0,0 @@ -0776a4d41c06192c4ca0409ad8b837de libtensorrt_llm_batch_manager_static.a -c901725d5d278fd8d41f524f81fe5170 libtensorrt_llm_batch_manager_static.pre_cxx11.a -b3330c65d9b23d4f20c2b8d5a7c24cd45c910cd4 commit diff --git a/cpp/tensorrt_llm/batch_manager/x86_64-linux-gnu/libtensorrt_llm_batch_manager_static.a b/cpp/tensorrt_llm/batch_manager/x86_64-linux-gnu/libtensorrt_llm_batch_manager_static.a index e5ed0eca9832..c6e6af3c54ef 100644 --- a/cpp/tensorrt_llm/batch_manager/x86_64-linux-gnu/libtensorrt_llm_batch_manager_static.a +++ b/cpp/tensorrt_llm/batch_manager/x86_64-linux-gnu/libtensorrt_llm_batch_manager_static.a @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:abdce9bc64cecddb39ed14809eefc8bcf7164524a6dd20ec7c8167229f3c22a3 -size 1557782 +oid sha256:681917aea11f45d83ba1429ded44ced97cb8ce5f54eb1c3fb3055bc342f0ffbf +size 1600734 diff --git a/cpp/tensorrt_llm/batch_manager/x86_64-linux-gnu/libtensorrt_llm_batch_manager_static.pre_cxx11.a b/cpp/tensorrt_llm/batch_manager/x86_64-linux-gnu/libtensorrt_llm_batch_manager_static.pre_cxx11.a index 6ca65959a2c9..bd661ab46cd5 100644 --- a/cpp/tensorrt_llm/batch_manager/x86_64-linux-gnu/libtensorrt_llm_batch_manager_static.pre_cxx11.a +++ b/cpp/tensorrt_llm/batch_manager/x86_64-linux-gnu/libtensorrt_llm_batch_manager_static.pre_cxx11.a @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:a9109b506e993a041ea238f992bec2a5064dffd9c0a7af10cca0d4d96c5047a9 -size 1557482 +oid sha256:d59b04e3229358ec2d9476b07f0361aa4a8539e543312c8952b690173040663d +size 1598666 diff --git a/cpp/tensorrt_llm/batch_manager/x86_64-linux-gnu/version.txt b/cpp/tensorrt_llm/batch_manager/x86_64-linux-gnu/version.txt index bc433dbfe7ff..7aa1380c27d2 100644 --- a/cpp/tensorrt_llm/batch_manager/x86_64-linux-gnu/version.txt +++ b/cpp/tensorrt_llm/batch_manager/x86_64-linux-gnu/version.txt @@ -1,2 +1,2 @@ -25d1ebdd5977208c25023329c621e970 libtensorrt_llm_batch_manager_static.a -5cb1a7a13db34fcaee6b89fcdc1212ce libtensorrt_llm_batch_manager_static.pre_cxx11.a +c9d5678a2ec347188457ad4a3a59d483 libtensorrt_llm_batch_manager_static.a +53261d576d540ab330f2f2e1f8d99677 libtensorrt_llm_batch_manager_static.pre_cxx11.a diff --git a/cpp/tensorrt_llm/common/mpiUtils.h b/cpp/tensorrt_llm/common/mpiUtils.h index f5be0ca88466..e1ffee538fd3 100644 --- a/cpp/tensorrt_llm/common/mpiUtils.h +++ b/cpp/tensorrt_llm/common/mpiUtils.h @@ -16,6 +16,8 @@ #pragma once +#include "tensorrt_llm/runtime/utils/multiDeviceUtils.h" + #include #include #include @@ -24,16 +26,7 @@ #include #define COMM_WORLD MpiComm(MPI_COMM_WORLD) -#define MPICHECK(cmd) \ - do \ - { \ - int e = cmd; \ - if (e != MPI_SUCCESS) \ - { \ - printf("Failed: MPI error %s:%d '%d'\n", __FILE__, __LINE__, e); \ - exit(EXIT_FAILURE); \ - } \ - } while (0) +#define MPICHECK(cmd) TLLM_MPI_CHECK(cmd) // A wrapper module of the MPI library. namespace tensorrt_llm::mpi diff --git a/cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/threadblock/epilogue_tensor_op_int32.h b/cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/threadblock/epilogue_tensor_op_int32.h index 8f44c49bba43..6f26d7901703 100644 --- a/cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/threadblock/epilogue_tensor_op_int32.h +++ b/cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/epilogue/threadblock/epilogue_tensor_op_int32.h @@ -92,31 +92,18 @@ namespace threadblock namespace detail { -/// Partial specialization for half <= int32_t x 8 epilogues avoids shared memory bank conflicts. -template -struct DefaultIteratorsTensorOp -{ - - using WarpTileIterator - = cutlass::epilogue::warp::TileIteratorTensorOp; - - using SharedLoadIterator = cutlass::epilogue::threadblock::SharedLoadIterator; - - static int const kFragmentsPerIteration = 1; -}; - /// Partial specialization for bfloat16_t <= int32_t x 8 epilogues avoids shared memory bank conflicts. template struct DefaultIteratorsTensorOp { - using WarpTileIterator - = cutlass::epilogue::warp::TileIteratorTensorOp; + = cutlass::epilogue::warp::TileIteratorTensorOpMixed; - using SharedLoadIterator = cutlass::epilogue::threadblock::SharedLoadIterator; + using SharedLoadIterator + = cutlass::epilogue::threadblock::SharedLoadIteratorMixed; - static int const kFragmentsPerIteration = 1; + static int const kFragmentsPerIteration = 2; }; ///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/device/gemm_universal_base_compat.h b/cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/device/gemm_universal_base_compat.h new file mode 100644 index 000000000000..2edd5a228b47 --- /dev/null +++ b/cpp/tensorrt_llm/cutlass_extensions/include/cutlass_extensions/gemm/device/gemm_universal_base_compat.h @@ -0,0 +1,438 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief The universal GEMM accommodates serial reductions, parallel reductions, batched strided, and + batched array variants. +*/ + +#pragma once + +// #include + +#include "cutlass/arch/arch.h" +#include "cutlass/cutlass.h" +#include "cutlass/device_kernel.h" +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/kernel/gemm_universal.h" +#include "cutlass/gemm/threadblock/threadblock_swizzle.h" + +#include "cutlass/gemm/device/default_gemm_configuration.h" +#include "cutlass/gemm/kernel/default_gemm_universal.h" + +#include "cutlass/trace.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass +{ +namespace gemm +{ +namespace device +{ + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/* + This is the device layer from CUTLASS 2.10 (SHA - cc85b64cf676c45f98a17e3a47c0aafcf817f088) + It is replicated here since we needed to duplicate kernel level APIs for mixed dtype GEMMs + and SmoothQuant. The newer device layer is not compatible with these older kernel level APIs. + + Note: While CUTLASS 3.x supports stream-k, none of the kernels in the extensions folder support + that feature at the moment. + */ + +template +class GemmUniversalBaseCompat +{ +public: + using GemmKernel = GemmKernel_; + using ThreadblockShape = typename GemmKernel::Mma::Shape; + + using ElementA = typename GemmKernel::ElementA; + using LayoutA = typename GemmKernel::LayoutA; + using TensorRefA = TensorRef; + static ComplexTransform const kTransformA = GemmKernel::kTransformA; + + using ElementB = typename GemmKernel::ElementB; + using LayoutB = typename GemmKernel::LayoutB; + using TensorRefB = TensorRef; + static ComplexTransform const kTransformB = GemmKernel::kTransformB; + + using ElementC = typename GemmKernel::ElementC; + using LayoutC = typename GemmKernel::LayoutC; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + + using ElementAccumulator = typename GemmKernel::Mma::Policy::Operator::ElementC; + + using EpilogueOutputOp = typename GemmKernel::EpilogueOutputOp; + using ThreadblockSwizzle = typename GemmKernel::ThreadblockSwizzle; + using Operator = typename GemmKernel::Operator; + + /// Argument structure + using Arguments = typename GemmKernel::Arguments; + +protected: + /// Kernel parameters object + typename GemmKernel::Params params_; + +protected: + /// Private helper to obtain the grid dimensions with fix-up for split-K + static void get_grid_shape_(gemm::GemmCoord& grid_tiled_shape, int& gemm_k_size, Arguments const& args) + { + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + grid_tiled_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, args.batch_count); + + gemm_k_size = args.problem_size.k(); + + if (args.mode == GemmUniversalMode::kGemm || args.mode == GemmUniversalMode::kGemmSplitKParallel) + { + + int const kAlignK + = const_max(const_max(128 / sizeof_bits::value, 128 / sizeof_bits::value), 1); + + gemm_k_size = round_up(ceil_div(args.problem_size.k(), args.batch_count), kAlignK); + + if (gemm_k_size) + { + grid_tiled_shape.k() = ceil_div(args.problem_size.k(), gemm_k_size); + } + } + } + +public: + /// Constructs the GEMM. + GemmUniversalBaseCompat() {} + + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const& args) + { + + // Determine grid shape + cutlass::gemm::GemmCoord grid_tiled_shape; + int gemm_k_size = 0; + + get_grid_shape_(grid_tiled_shape, gemm_k_size, args); + + ThreadblockSwizzle threadblock_swizzle; + dim3 grid = threadblock_swizzle.get_grid_shape(grid_tiled_shape); + + uint32_t const kGridYZMax = ((1 << (sizeof(uint16_t) * 8)) - 1); + + if (!(grid.y <= kGridYZMax && grid.z <= kGridYZMax)) + { + + return Status::kErrorInvalidProblem; + } + + return GemmKernel::can_implement(args); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const& args) + { + + CUTLASS_TRACE_HOST("GemmUniversalBaseCompat::get_workspace_size()"); + + size_t workspace_bytes = 0; + + // Determine grid shape + cutlass::gemm::GemmCoord grid_tiled_shape; + int gemm_k_size = 0; + + get_grid_shape_(grid_tiled_shape, gemm_k_size, args); + + if (args.mode == GemmUniversalMode::kGemmSplitKParallel) + { + + // Split-K parallel always requires a temporary workspace + workspace_bytes = sizeof(ElementC) * size_t(args.batch_stride_D) * size_t(grid_tiled_shape.k()); + } + else if (args.mode == GemmUniversalMode::kGemm && grid_tiled_shape.k() > 1) + { + + // Serial split-K only requires a temporary workspace if the number of partitions along the + // GEMM K dimension is greater than one. + workspace_bytes = sizeof(int) * size_t(grid_tiled_shape.m()) * size_t(grid_tiled_shape.n()); + } + + CUTLASS_TRACE_HOST(" workspace_bytes: " << workspace_bytes); + + workspace_bytes += GemmKernel::get_extra_workspace_size(args, grid_tiled_shape); + + return workspace_bytes; + } + + /// Computes the grid shape + static dim3 get_grid_shape(Arguments const& args) + { + + CUTLASS_TRACE_HOST("GemmUniversalBaseCompat::get_grid_shape()"); + + ThreadblockSwizzle threadblock_swizzle; + + cutlass::gemm::GemmCoord grid_tiled_shape; + int gemm_k_size = 0; + + get_grid_shape_(grid_tiled_shape, gemm_k_size, args); + dim3 result = threadblock_swizzle.get_grid_shape(grid_tiled_shape); + + CUTLASS_TRACE_HOST(" grid_tiled_shape: " << grid_tiled_shape << "\n" + << " result = {" << result << "}"); + + return result; + } + + /// Computes the maximum number of active blocks per multiprocessor + static int maximum_active_blocks(int smem_capacity = -1) + { + + CUTLASS_TRACE_HOST("GemmUniversalBaseCompat::maximum_active_blocks()"); + + int max_active_blocks = -1; + int smem_size = int(sizeof(typename GemmKernel::SharedStorage)); + + CUTLASS_TRACE_HOST(" smem_size: " << smem_size << " bytes"); + + if (smem_size <= (48 << 10)) + { + + cudaError_t result = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_active_blocks, Kernel, GemmKernel::kThreadCount, smem_size); + + if (result == cudaSuccess) + { + CUTLASS_TRACE_HOST(" max_active_blocks: " << max_active_blocks); + return max_active_blocks; + } + } + else + { + + // Query assuming zero shared memory then compute occupancy limit based on SMEM + cudaError_t result = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_active_blocks, Kernel, GemmKernel::kThreadCount, 0); + + if (result != cudaSuccess) + { + + CUTLASS_TRACE_HOST( + " cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error " << cudaGetErrorString(result)); + + return -1; + } + + if (smem_capacity < 0) + { + int device_idx = 0; + result = cudaGetDevice(&device_idx); + + if (result != cudaSuccess) + { + return -1; + } + + cudaDeviceProp properties; + result = cudaGetDeviceProperties(&properties, device_idx); + + if (result != cudaSuccess) + { + return -1; + } + + smem_capacity = static_cast(properties.sharedMemPerMultiprocessor); + } + + int occupancy = std::min(max_active_blocks, smem_capacity / smem_size); + + CUTLASS_TRACE_HOST(" occupancy: " << occupancy); + + return occupancy; + } + + CUTLASS_TRACE_HOST(" returning internal error"); + + return -1; + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) + { + + CUTLASS_TRACE_HOST("GemmUniversalBaseCompat::initialize() - workspace " + << workspace << ", stream: " << (stream ? "non-null" : "null")); + + size_t workspace_bytes = get_workspace_size(args); + + CUTLASS_TRACE_HOST(" workspace_bytes: " << workspace_bytes); + + if (workspace_bytes) + { + + if (!workspace) + { + CUTLASS_TRACE_HOST(" error: device workspace must not be null"); + + return Status::kErrorWorkspaceNull; + } + + if (args.mode == GemmUniversalMode::kGemm) + { + CUTLASS_TRACE_HOST(" clearing device workspace"); + cudaError_t result = cudaMemsetAsync(workspace, 0, workspace_bytes, stream); + + if (result != cudaSuccess) + { + CUTLASS_TRACE_HOST(" cudaMemsetAsync() returned error " << cudaGetErrorString(result)); + + return Status::kErrorInternal; + } + } + } + + // Get CUDA grid shape + cutlass::gemm::GemmCoord grid_tiled_shape; + int gemm_k_size = 0; + + get_grid_shape_(grid_tiled_shape, gemm_k_size, args); + + // Initialize the Params structure + params_ = typename GemmKernel::Params(args, grid_tiled_shape, gemm_k_size, static_cast(workspace)); + + // Specify shared memory capacity for kernel. + int smem_size = int(sizeof(typename GemmKernel::SharedStorage)); + + if (smem_size >= (48 << 10)) + { + cudaError_t result + = cudaFuncSetAttribute(Kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + + if (result != cudaSuccess) + { + return Status::kErrorInternal; + } + } + + return Status::kSuccess; + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const& args, void* workspace = nullptr) + { + + CUTLASS_TRACE_HOST("GemmUniversalBaseCompat()::update() - workspace: " << workspace); + + size_t workspace_bytes = get_workspace_size(args); + + if (workspace_bytes && !workspace) + { + return Status::kErrorWorkspaceNull; + } + + params_.update(args, workspace); + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(cudaStream_t stream = nullptr) + { + CUTLASS_TRACE_HOST("GemmUniversalBaseCompat::run()"); + + // + // Configure grid and block dimensions + // + + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape); + dim3 block(GemmKernel::kThreadCount, 1, 1); + + int smem_size = int(sizeof(typename GemmKernel::SharedStorage)); + + // + // Launch kernel + // + + CUTLASS_TRACE_HOST(" grid: (" << grid << "), block: (" << block << "), SMEM: " << smem_size << " bytes"); + + // Launch + cutlass::Kernel<<>>(params_); + + // + // Query for errors + // + cudaError_t result = cudaGetLastError(); + + if (result != cudaSuccess) + { + CUTLASS_TRACE_HOST(" grid launch failed with error " << cudaGetErrorString(result)); + return Status::kErrorInternal; + } + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status operator()(cudaStream_t stream = nullptr) + { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) + { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) + { + status = run(stream); + } + + return status; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt index a7dab1378d44..8d543231176c 100644 --- a/cpp/tensorrt_llm/kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt @@ -18,6 +18,12 @@ file(GLOB_RECURSE SRC_CPP *.cpp) file(GLOB_RECURSE SRC_CU *.cu) +# skip mmha 48, 80, 96, 112, 144, 160, 192 and 224 for fast build +if(FAST_BUILD) + list(FILTER SRC_CU EXCLUDE REGEX + "decoderMaskedMultiheadAttention(48|80|96|112|144|160|192|224).*cu$") +endif() + add_library(kernels_src OBJECT ${SRC_CPP} ${SRC_CU}) set_property(TARGET kernels_src PROPERTY POSITION_INDEPENDENT_CODE ON) set_property(TARGET kernels_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) diff --git a/cpp/tensorrt_llm/kernels/beamSearchPenaltyKernels.cu b/cpp/tensorrt_llm/kernels/beamSearchPenaltyKernels.cu index 5ca0f47ac121..56b305f69582 100644 --- a/cpp/tensorrt_llm/kernels/beamSearchPenaltyKernels.cu +++ b/cpp/tensorrt_llm/kernels/beamSearchPenaltyKernels.cu @@ -14,6 +14,7 @@ * limitations under the License. */ +#include // all_of #include #include "tensorrt_llm/common/assert.h" diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm_template.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm_template.h index 2d60cbc02b40..4858fef00d4e 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm_template.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm_template.h @@ -19,9 +19,9 @@ #pragma GCC diagnostic ignored "-Wstrict-aliasing" #endif // #ifndef _WIN32 -#include "cutlass/gemm/device/gemm_universal_base.h" #include "cutlass/gemm/kernel/default_gemm.h" #include "cutlass_extensions/compute_occupancy.h" +#include "cutlass_extensions/gemm/device/gemm_universal_base_compat.h" #include "cutlass_extensions/epilogue_helpers.h" #include "cutlass_extensions/gemm/kernel/default_fpA_intB_traits.h" @@ -124,7 +124,7 @@ void generic_mixed_gemm_kernelLauncher(const T* A, const WeightType* B, const T* return; } - using Gemm = cutlass::gemm::device::GemmUniversalBase; + using Gemm = cutlass::gemm::device::GemmUniversalBaseCompat; const int ldb = cutlass::platform::is_same::value ? n diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/int8_gemm/int8_gemm_template.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/int8_gemm/int8_gemm_template.h index ea0e2b004e39..6f783f64598e 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/int8_gemm/int8_gemm_template.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/int8_gemm/int8_gemm_template.h @@ -22,12 +22,13 @@ // clang-format off #include #include -#include +#include #include #include // clang-format on #include "cutlass_extensions/compute_occupancy.h" +#include "cutlass_extensions/epilogue/threadblock/epilogue_per_row_per_col_scale.h" #include "cutlass_extensions/epilogue/threadblock/epilogue_tensor_op_int32.h" #include "cutlass_extensions/epilogue_helpers.h" #include "cutlass_extensions/gemm_configs.h" @@ -123,7 +124,7 @@ void genericInt8GemmKernelLauncher(const int8_t* A, const int8_t* B, tk::QuantMo return; } - using Gemm = cutlass::gemm::device::GemmUniversalBase; + using Gemm = cutlass::gemm::device::GemmUniversalBaseCompat; typename EpilogueOp::Params linearScalingParams; // TODO: right now it's unused (scaling is done in // visitor, no activation needed) diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionLaunch.h b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionLaunch.h index abb0f5d83e88..b602a1242917 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionLaunch.h +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionLaunch.h @@ -226,7 +226,7 @@ void mmha_launch_kernel_ex( } // If blocks with larger block size already fill all SMs, then disable the multi blocks mode. - mmha::multi_block_grid_setup(grid, params, dynamic_block_size, available_blocks, tlength, DO_MULTI_BLOCK); + mmha::multi_block_grid_setup(grid, params, available_blocks, dynamic_block_size, tlength, DO_MULTI_BLOCK); // Launch kernels based on the valid block size. switch (dynamic_block_size) diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionTemplate.h b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionTemplate.h index df47baa7531e..49e7d61bcb66 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionTemplate.h +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderMaskedMultiheadAttentionTemplate.h @@ -1411,7 +1411,8 @@ __global__ void masked_multihead_attention_kernel( bool has_relative_attention_bias = params.relative_attention_bias != nullptr; // Compute relative attention bias on the fly, with relative attention table [head_num/TP, num_buckets] passed in. // num_buckets passed as relative_attention_bias_stride, max_distance passed as params.max_distance - const bool implicit_rel_attn_bias = DO_CROSS_ATTENTION && params.max_distance != 0 && has_relative_attention_bias; + // this is a common optimization for both self attention and cross attention + const bool implicit_rel_attn_bias = params.max_distance != 0 && has_relative_attention_bias; int relative_attention_bias_stride = params.relative_attention_bias_stride; // num_buckets might be modified below, save it beforehand int max_distance = params.max_distance; @@ -1693,12 +1694,15 @@ __global__ void masked_multihead_attention_kernel( // Pre-compute the pointer for the relative attention bias. const T* relative_attention_bias_ptr = nullptr; + const T* relative_attention_bias_ptr_fixed = nullptr; // record the base for offset if (has_relative_attention_bias) { + // "hi" is unsigned, subtracting int from unsigned int causes underflow. Cast to int int64_t offset = implicit_rel_attn_bias - ? (hi * relative_attention_bias_stride - tlength) - : (hi * relative_attention_bias_stride + tlength) * relative_attention_bias_stride; + ? ((int64_t) hi * relative_attention_bias_stride - tlength) + : ((int64_t) hi * relative_attention_bias_stride + tlength) * relative_attention_bias_stride; relative_attention_bias_ptr = ¶ms.relative_attention_bias[offset]; + relative_attention_bias_ptr_fixed = ¶ms.relative_attention_bias[offset]; } // Load the value. @@ -1706,7 +1710,7 @@ __global__ void masked_multihead_attention_kernel( if (has_relative_attention_bias && tidx == 0) { // TODO: Use a better way to convert from T to float. - add(relative_attention_bias, relative_attention_bias_ptr[tlength]); + relative_attention_bias = add(relative_attention_bias, relative_attention_bias_ptr[tlength]); } // Store that value in shared memory. Keep the Q*K^T value in register for softmax. @@ -1769,7 +1773,19 @@ __global__ void masked_multihead_attention_kernel( // Pick a number of keys to make sure all the threads of a warp enter (due to shfl_sync). // Take all previous cache as context when we have no beam searching in order to batch as many LDGs as possible. - const int context_length = HAS_BEAMS ? beam0_context_length : kv_loop_length; + const int context_length + = DO_CROSS_ATTENTION ? kv_loop_length : (HAS_BEAMS ? beam0_context_length : kv_loop_length); + // Clarifications: + // - in self attn, input_length is input text length, tlength is current timestep + // - in cross attn, input_length is *decoder* input length (usually 1), tlength is *encoder* input context length + // - in beam search, since the cache during generation is organized differently, the following KV compute needs + // split into context cache compute and generation cache compute + // - for self attn, no-beam search: entire cache can be treated as context cache --> context_length = tlength + // - for self attn, beam search: cache of input text length is context cache, other are generation cache --> + // context_length = input_length + // - for cross attn, no-beam/beam search: cache length is fixed, not differ context/generation cache --> + // context_length = tlength Suggestion: we could have a flag HANDLE_GEN_CACHE + const auto context_ti_end = MULTI_BLOCK_FLAG ? divUp(timesteps_per_block, UNROLLED_K_PER_WARP) * UNROLLED_K_PER_WARP : divUp(static_cast(context_length), UNROLLED_K_PER_WARP) * UNROLLED_K_PER_WARP; @@ -1872,7 +1888,7 @@ __global__ void masked_multihead_attention_kernel( relative_position_if_large = min(relative_position_if_large, num_buckets - 1); relative_buckets += is_small ? relative_position : relative_position_if_large; relative_attention_bias_ptr - = relative_attention_bias_ptr + (tlength - local_time_now) + relative_buckets; + = relative_attention_bias_ptr_fixed + (tlength - local_time_now) + relative_buckets; } // Prefetch the relative attention bias. @@ -1880,7 +1896,7 @@ __global__ void masked_multihead_attention_kernel( if (is_active && has_relative_attention_bias) { // TODO: Use a better way to convert from T to float. - add(relative_attention_bias, relative_attention_bias_ptr[local_time_now]); + relative_attention_bias = add(relative_attention_bias, relative_attention_bias_ptr[local_time_now]); } // Compute the dot product between Q and K. @@ -1937,7 +1953,9 @@ __global__ void masked_multihead_attention_kernel( // Handle generation key cache with beam searching. // Note that it may be overlapped with the context key loop, but it won't impact the corretness. - if (HAS_BEAMS && (!MULTI_BLOCK_FLAG || (c_tile + 1) * timesteps_per_block > beam0_context_length)) + // Can skip in cross attention mode. + if (HAS_BEAMS && !DO_CROSS_ATTENTION + && (!MULTI_BLOCK_FLAG || (c_tile + 1) * timesteps_per_block > beam0_context_length)) { // The input length; const int input_length_ = MULTI_BLOCK_FLAG ? beam0_context_length % timesteps_per_block : beam0_context_length; @@ -1987,7 +2005,8 @@ __global__ void masked_multihead_attention_kernel( * (num_buckets - max_exact)); relative_position_if_large = min(relative_position_if_large, num_buckets - 1); relative_buckets += is_small ? relative_position : relative_position_if_large; - relative_attention_bias_ptr = relative_attention_bias_ptr + (tlength - time_now) + relative_buckets; + relative_attention_bias_ptr + = relative_attention_bias_ptr_fixed + (tlength - time_now) + relative_buckets; } // Prefetch the relative attention bias. @@ -1995,7 +2014,7 @@ __global__ void masked_multihead_attention_kernel( if (is_active && has_relative_attention_bias) { // TODO: Use a better way to convert from T to float. - add(relative_attention_bias, relative_attention_bias_ptr[time_now]); + relative_attention_bias = add(relative_attention_bias, relative_attention_bias_ptr[time_now]); } // Perform the dot product and normalize qk. @@ -2260,7 +2279,8 @@ __global__ void masked_multihead_attention_kernel( // Handle both context and generation value cache without beam searching. // Explicit batching of LDGs (by V_LOOP_UNROLL) as it doesn't depend on indirection tables. // Take all previous cache as context when we have no beam searching in order to batch as many LDGs as possible. - const int context_length = HAS_BEAMS ? beam0_context_length : kv_loop_length; + const int context_length + = DO_CROSS_ATTENTION ? kv_loop_length : (HAS_BEAMS ? beam0_context_length : kv_loop_length); int context_v_loop_end = MULTI_BLOCK_FLAG ? timesteps_per_block : context_length; int generation_v_loop_end = MULTI_BLOCK_FLAG ? timesteps_per_block : kv_loop_length; for (int ti = vo; ti < context_v_loop_end; ti += UNROLLED_V_PER_ITER) @@ -2300,7 +2320,7 @@ __global__ void masked_multihead_attention_kernel( } // Handle generation value cache with beam searching. - if (HAS_BEAMS) + if (HAS_BEAMS && !DO_CROSS_ATTENTION) { const auto generation_start_ti = MULTI_BLOCK_FLAG ? vo : (vo + (beam0_context_length / V_PER_ITER) * V_PER_ITER); diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttentionUtils.h b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttentionUtils.h index 1c9799d1f883..54083e79ff46 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttentionUtils.h +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttentionUtils.h @@ -2226,6 +2226,8 @@ inline __device__ Float8_ mul(uint4 a, int64_t b) fc.y = mul(a.y, make_float2(int8[2], int8[3])); fc.z = mul(a.z, make_float2(int8[4], int8[5])); fc.w = mul(a.w, make_float2(int8[6], int8[7])); + + return fc; } //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -2247,6 +2249,8 @@ inline __device__ Float8_ mul(Float8_ fa, int64_t b) fc.y = mul(fa.y, make_float2(int8[2], int8[3])); fc.z = mul(fa.z, make_float2(int8[4], int8[5])); fc.w = mul(fa.w, make_float2(int8[6], int8[7])); + + return fc; } //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -2323,6 +2327,8 @@ inline __device__ Float8_ mul(bf16_8_t a, int64_t b) fc.y = mul(a.y, make_float2(int8[2], int8[3])); fc.z = mul(a.z, make_float2(int8[4], int8[5])); fc.w = mul(a.w, make_float2(int8[6], int8[7])); + + return fc; } #endif // ENABLE_BF16 diff --git a/cpp/tensorrt_llm/kernels/decodingKernels.cu b/cpp/tensorrt_llm/kernels/decodingKernels.cu index db571cb07118..a6d300c13416 100644 --- a/cpp/tensorrt_llm/kernels/decodingKernels.cu +++ b/cpp/tensorrt_llm/kernels/decodingKernels.cu @@ -416,8 +416,13 @@ __global__ void finalize(int* output_ids, int* sequence_lengths, float* cum_log_ = topk_output_ids[blockIdx.x * (beam_width * 2) * max_seq_len + s_rank[beam_idx] * max_seq_len + i]; if (output_log_probs != nullptr) { - output_log_probs[blockIdx.x * beam_width * max_seq_len + beam_idx * max_seq_len + i] - = topk_log_probs[blockIdx.x * (beam_width * 2) * max_seq_len + s_rank[beam_idx] * max_seq_len + i]; + int input_len = input_lengths[blockIdx.x * beam_width + beam_idx]; + if (i >= input_len) + { + output_log_probs[blockIdx.x * beam_width * max_seq_len + beam_idx * max_seq_len + i - input_len] + = topk_log_probs[blockIdx.x * (beam_width * 2) * max_seq_len + s_rank[beam_idx] * max_seq_len + + i]; + } } } } @@ -471,5 +476,32 @@ void invokeCopyNextStepIds(int* next_step_ids, int** output_ids_ptr, const int* next_step_ids, output_ids_ptr, sequence_lengths, batch_size, beam_width, max_seq_len); } +__global__ void transposeLogProbs(float* output_log_probs, float* output_log_probs_tiled, const int* sequence_lengths, + int batch_size, int beam_width, int max_seq_len) +{ + int index = blockIdx.x * blockDim.x + threadIdx.x; + + const int batch_idx = index / (beam_width * max_seq_len); + const int tmp_idx = index % (beam_width * max_seq_len); + const int beam_idx = tmp_idx / max_seq_len; + const int pos = tmp_idx % max_seq_len; + + if (batch_idx < batch_size && pos < sequence_lengths[batch_idx]) + { + + output_log_probs[index] + = output_log_probs_tiled[pos * batch_size * beam_width + batch_idx * beam_width + beam_idx]; + } +} + +void invokeTransposeLogProbs(float* output_log_probs, float* output_log_probs_tiled, const int* sequence_lengths, + int batch_size, int beam_width, int max_seq_len, cudaStream_t stream) +{ + dim3 block(256); + dim3 grid(divUp(batch_size * beam_width * max_seq_len, block.x)); + transposeLogProbs<<>>( + output_log_probs, output_log_probs_tiled, sequence_lengths, batch_size, beam_width, max_seq_len); +} + } // namespace kernels } // namespace tensorrt_llm diff --git a/cpp/tensorrt_llm/kernels/decodingKernels.h b/cpp/tensorrt_llm/kernels/decodingKernels.h index d4e05a3fa1d3..21a13102c7a9 100644 --- a/cpp/tensorrt_llm/kernels/decodingKernels.h +++ b/cpp/tensorrt_llm/kernels/decodingKernels.h @@ -62,5 +62,8 @@ void invokeInitializeOutput(int* output_ids, const int* end_ids, int batch_beam, void invokeCopyNextStepIds(int* next_step_ids, int** output_ids_ptr, const int* sequence_lengths, int batch_size, int beam_width, int max_seq_len, cudaStream_t stream); +void invokeTransposeLogProbs(float* output_log_probs, float* output_log_probs_tiled, const int* sequence_lengths, + int batch_size, int beam_width, int max_seq_len, cudaStream_t stream); + } // namespace kernels } // namespace tensorrt_llm diff --git a/cpp/tensorrt_llm/kernels/gptKernels.cu b/cpp/tensorrt_llm/kernels/gptKernels.cu index e341a03de524..60729a72e805 100644 --- a/cpp/tensorrt_llm/kernels/gptKernels.cu +++ b/cpp/tensorrt_llm/kernels/gptKernels.cu @@ -195,12 +195,23 @@ __global__ void computeAttentionMask(AttentionMaskDataType* attentionMask, const isValid = (rowIdx < seqLength - 1 && colIdx < seqLength - 1) || (rowIdx == seqLength - 1 && colIdx < seqLength); // clang-format on - // seq_length==4, max_seq_len==5, only use in context phase + // seq_length==4, max_seq_len==5 // 1 1 1 0 0 // 1 1 1 0 0 // 1 1 1 0 0 // 1 1 1 1 0 // 0 0 0 0 0 + case AttentionMaskType::BIDIRECTIONALGLM: + // clang-format off + isValid = (colIdx < seqLength - 1) || + (rowIdx == maxSeqLength - 1 && colIdx == maxSeqLength - 1); + // clang-format on + // seq_length==4, max_seq_len==5 + // 1 1 1 1 0 + // 1 1 1 1 0 + // 1 1 1 1 0 + // 1 1 1 1 0 + // 1 1 1 1 1 break; } diff --git a/cpp/tensorrt_llm/kernels/gptKernels.h b/cpp/tensorrt_llm/kernels/gptKernels.h index b30edfa807c6..e486dade6c09 100644 --- a/cpp/tensorrt_llm/kernels/gptKernels.h +++ b/cpp/tensorrt_llm/kernels/gptKernels.h @@ -31,7 +31,10 @@ enum class AttentionMaskType // Mask the padded tokens and all the tokens that come after in a sequence. CAUSAL = 1, // See ChatGLM-6B mask. - BIDIRECTIONAL = 2 + BIDIRECTIONAL = 2, + // See GLM-10B mask. + // TODO: merge this mask into BIDIRECTIONAL + BIDIRECTIONALGLM = 3 }; enum class PositionEmbeddingType : int8_t @@ -58,7 +61,7 @@ struct BuildDecoderInfoParams { // The offsets to the 1st token in each sequence. Shape: [batchSize+1]. int* seqOffsets; - // The number of padded tokens in the corresponding padded tensor. Shape: [numTokens]. + // The number of padded tokens in the corresponding padded tensor before the current token. Shape: [numTokens]. int* paddingOffsets; // The mask to mark invalid tokens in Attention - that's not used by the plugins as it can be diff --git a/cpp/tensorrt_llm/kernels/kvCacheUtils.h b/cpp/tensorrt_llm/kernels/kvCacheUtils.h index e65a2cca0a7c..d342644ab2ff 100644 --- a/cpp/tensorrt_llm/kernels/kvCacheUtils.h +++ b/cpp/tensorrt_llm/kernels/kvCacheUtils.h @@ -18,6 +18,7 @@ #include "tensorrt_llm/common/assert.h" #include #include +#include namespace tensorrt_llm { @@ -64,6 +65,11 @@ struct KVBlockArray const float tokensPerBlockSeqLog2 = log2(mTokensPerBlock); TLLM_CHECK_WITH_INFO( ceil(tokensPerBlockSeqLog2) == floor(tokensPerBlockSeqLog2), "tokensPerBlock must be power of 2"); + // NOTE: pointer offset arithmetic offset is performed on int32_t (see this.getRowPtr). + // If needed, we could do it on uint32_t or even uint64_t, but that might have performance implications + TLLM_CHECK_WITH_INFO(static_cast(mMaxSeqs - 1) * mMaxBlocksPerSeq * 2 + maxBlocksPerSeq + <= std::numeric_limits::max(), + "kv cache is too large for gpt_attention_plugin"); mTokensPerBlockLog2 = static_cast(tokensPerBlockSeqLog2); } @@ -140,6 +146,11 @@ struct KVLinearBuffer , mMaxSeqLen(tokensPerBlock) , mBytesPerSeq(tokensPerBlock * sizePerToken) { + // NOTE: pointer offset arithmetic offset is performed on int32_t (see this.getRowPtr). + // If needed, we could do it on uint32_t or even uint64_t, but that might have performance implications + TLLM_CHECK_WITH_INFO( + static_cast(mMaxSeqs - 1) * mBytesPerSeq * 2 + mBytesPerSeq <= std::numeric_limits::max(), + "kv cache is too large for gpt_attention_plugin"); } __host__ __device__ inline void** getRowPtr(KVIdxType kvIdx, int32_t seqIdx) diff --git a/cpp/tensorrt_llm/kernels/layernormKernels.cu b/cpp/tensorrt_llm/kernels/layernormKernels.cu index c4ba76092aef..8145f507351d 100644 --- a/cpp/tensorrt_llm/kernels/layernormKernels.cu +++ b/cpp/tensorrt_llm/kernels/layernormKernels.cu @@ -198,12 +198,10 @@ void dispatch_layernorm_type_square_method(const T* input, const T* gamma, const float* scale_orig_quant_per_token, int8_t* normed_output_quant, const dim3 grid, const dim3 block, const size_t shmem_size, cudaStream_t stream) { - bool use_shmem = true; if (shmem_size >= (48 << 10)) { cudaError_t ret = cudaFuncSetAttribute( generalLayerNorm, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_size); - use_shmem = ret == cudaSuccess; } generalLayerNorm<<>>(input, gamma, beta, normed_output, eps, tokens, hidden_dim, scale_orig_quant_per_tensor, scale_orig_quant_per_token, normed_output_quant, true); diff --git a/cpp/tensorrt_llm/kernels/rmsnormKernels.cu b/cpp/tensorrt_llm/kernels/rmsnormKernels.cu index 6042667b5399..326733605967 100644 --- a/cpp/tensorrt_llm/kernels/rmsnormKernels.cu +++ b/cpp/tensorrt_llm/kernels/rmsnormKernels.cu @@ -155,12 +155,10 @@ void dispatch_rmsnorm_type_square_method(const T* input, const T* gamma, const T float* scale_orig_quant_per_token, int8_t* normed_output_quant, const dim3 grid, const dim3 block, const size_t shmem_size, cudaStream_t stream) { - bool use_shmem = true; if (shmem_size >= (48 << 10)) { cudaError_t ret = cudaFuncSetAttribute(generalRmsNorm, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_size); - use_shmem = ret == cudaSuccess; } generalRmsNorm<<>>(input, gamma, beta, normed_output, eps, tokens, hidden_dim, scale_orig_quant_per_tensor, scale_orig_quant_per_token, normed_output_quant, true); diff --git a/cpp/tensorrt_llm/kernels/samplingTopPKernels.cu b/cpp/tensorrt_llm/kernels/samplingTopPKernels.cu index ce04c6cc06ad..5815bbaebfa2 100644 --- a/cpp/tensorrt_llm/kernels/samplingTopPKernels.cu +++ b/cpp/tensorrt_llm/kernels/samplingTopPKernels.cu @@ -203,7 +203,6 @@ __global__ void topPSsampling(T* sortedLogProbs, int* sortedIdVals, int** ids, i * output. */ - __shared__ int stopShared; __shared__ float randNumS; const int tid = threadIdx.x; @@ -233,7 +232,6 @@ __global__ void topPSsampling(T* sortedLogProbs, int* sortedIdVals, int** ids, i // will choose the token which probability makes cumulative probability sum to exceed P' if (threadIdx.x == 0) { - stopShared = 0; randNumS = curand_uniform(curandstate + blockIdx.x) * probThreshold; } diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.cu b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.cu index 2b86b605b340..4d0db9c249d5 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.cu +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.cu @@ -1361,7 +1361,6 @@ __global__ void add_fusedQKV_bias_transpose_kernel(T* q_buf, T* k_buf, T* v_buf, switch (position_embedding_type) { - case PositionEmbeddingType::kRELATIVE: case PositionEmbeddingType::kROPE_GPTJ: { mmha::apply_rotary_embedding( diff --git a/cpp/tensorrt_llm/layers/baseBeamSearchLayer.cu b/cpp/tensorrt_llm/layers/baseBeamSearchLayer.cu index d8c0ba0bf8d1..946d0214bf73 100644 --- a/cpp/tensorrt_llm/layers/baseBeamSearchLayer.cu +++ b/cpp/tensorrt_llm/layers/baseBeamSearchLayer.cu @@ -18,6 +18,9 @@ #include "tensorrt_llm/common/memoryUtils.h" #include "tensorrt_llm/kernels/beamSearchPenaltyKernels.h" #include "tensorrt_llm/layers/baseBeamSearchLayer.h" +#include "tensorrt_llm/layers/fillBuffers.h" + +#include using namespace tensorrt_llm::common; using namespace tensorrt_llm::kernels; @@ -32,7 +35,7 @@ __global__ void update_indir_cache_kernel(int* tgt_indir_cache, const int* src_i int beam_width, int max_kv_cache_length, int max_seq_len) { int time_step = threadIdx.x + blockIdx.x * blockDim.x; - int bb_id = threadIdx.y + blockIdx.y * blockDim.y; + int bb_id = threadIdx.y + blockIdx.y * blockDim.y; // should be just blockIdx.y? const int current_step{sequence_lengths[bb_id] - 1}; // the sequence_lengths is updated, need to minus 1 const int batch_id = bb_id / beam_width; const int beam_id = bb_id % beam_width; @@ -46,9 +49,9 @@ __global__ void update_indir_cache_kernel(int* tgt_indir_cache, const int* src_i const int src_beam = parent_ids[batch_id][beam_id * max_seq_len + current_step]; // for the indir tables, we have the cyclic kv cache. - const uint tgt_offset + const uint32_t tgt_offset = batch_id * beam_width * max_kv_cache_length + beam_id * max_kv_cache_length + time_step_circ; - const uint src_offset + const uint32_t src_offset = batch_id * beam_width * max_kv_cache_length + src_beam * max_kv_cache_length + time_step_circ; tgt_indir_cache[tgt_offset] = (time_step == current_step) ? beam_id : src_indir_cache[src_offset]; @@ -113,7 +116,7 @@ void BaseBeamSearchLayer::allocateBuffer(size_t batch_size) repetition_penalty_buf_ = allocator_->reMalloc(repetition_penalty_buf_, sizeof(float) * batch_size, false); is_allocate_buffer_ = true; - TLLM_LOG_DEBUG("% stop", __PRETTY_FUNCTION__); + TLLM_LOG_DEBUG("%s stop", __PRETTY_FUNCTION__); } template @@ -122,25 +125,7 @@ void BaseBeamSearchLayer::setupBase(size_t batch_size, SetupParams const& set allocateBuffer(batch_size); TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); // Setup penalties. - auto fillBuffers - = [this, &batch_size](auto const& optParam, auto const defaultValue, auto& hostBuffer, auto& deviceBuffer) - { - hostBuffer.resize(batch_size); - if (!optParam) - { - std::fill(std::begin(hostBuffer), std::end(hostBuffer), defaultValue); - } - else if (optParam->size() == 1) - { - std::fill(std::begin(hostBuffer), std::end(hostBuffer), optParam->front()); - } - else - { - TLLM_CHECK_WITH_INFO(optParam->size() == batch_size, "Argument vector size mismatch."); - std::copy(optParam->begin(), optParam->end(), std::begin(hostBuffer)); - } - cudaAutoCpy(deviceBuffer, hostBuffer.data(), batch_size, stream_); - }; + FillBuffers const fillBuffers{batch_size, stream_}; fillBuffers(setupParams.temperature, 1.0f, mTemperature, temperature_buf_); fillBuffers(setupParams.min_length, 1, mMinLength, min_lengths_buf_); diff --git a/cpp/tensorrt_llm/layers/dynamicDecodeLayer.cpp b/cpp/tensorrt_llm/layers/dynamicDecodeLayer.cpp index fbc0de7cda0c..f7458eb8aaae 100644 --- a/cpp/tensorrt_llm/layers/dynamicDecodeLayer.cpp +++ b/cpp/tensorrt_llm/layers/dynamicDecodeLayer.cpp @@ -321,7 +321,7 @@ void DynamicDecodeLayer::forward(OutputParams& outputs, ForwardParams const& = outputs.cum_log_probs->slice({dynamic_decode_batch_size * beam_width}, dynamic_id_offset); dynamic_decode_outputs.beamHypotheses = outputs.beamHypotheses; - dynamic_decode_outputs.output_log_probs = outputs.output_log_probs; + dynamic_decode_outputs.output_log_probs = outputs.output_log_probs_tiled; // only OnlineBeamSearchLayer support beam_search_diversity_rate // when beamHypotheses is used @@ -365,11 +365,11 @@ void DynamicDecodeLayer::forward(OutputParams& outputs, ForwardParams const& decode_outputs.cum_log_probs = outputs.cum_log_probs->slice({local_batch_size * beam_width}, local_batch_offset); } - if (outputs.output_log_probs) + if (outputs.output_log_probs_tiled) { auto const generationStep = step - params.max_input_length; TLLM_CHECK(generationStep >= 0); - Tensor& output_log_probs = outputs.output_log_probs.value(); + Tensor& output_log_probs = outputs.output_log_probs_tiled.value(); size_t step_offset = generationStep * batch_size * beam_width; decode_outputs.output_log_probs = output_log_probs.slice({output_log_probs.shape[0] - generationStep, local_batch_size * beam_width}, @@ -411,6 +411,17 @@ void DynamicDecodeLayer::forward(OutputParams& outputs, ForwardParams const& invokeCopyNextStepIds(outputs.newTokens.template getPtr(), idsPtrHost, outputs.sequence_length->template getPtr(), batch_size, beam_width, max_seq_len, stream_); + + // Transpose the output log probs from [max_seq_len, bs, beam_width] to [batch_size, beam_width, max_seq_len] + if (outputs.output_log_probs_tiled) + { + auto logProbsMaxSeqLen = outputs.output_log_probs_tiled.value().shape[0]; + + invokeTransposeLogProbs(outputs.output_log_probs.value().template getPtr(), + outputs.output_log_probs_tiled.value().template getPtr(), + outputs.sequence_length->template getPtr(), batch_size, beam_width, logProbsMaxSeqLen, stream_); + } + sync_check_cuda_error(); } diff --git a/cpp/tensorrt_llm/layers/dynamicDecodeLayer.h b/cpp/tensorrt_llm/layers/dynamicDecodeLayer.h index ae7ac8cd22c8..d20a96e50f67 100644 --- a/cpp/tensorrt_llm/layers/dynamicDecodeLayer.h +++ b/cpp/tensorrt_llm/layers/dynamicDecodeLayer.h @@ -129,14 +129,16 @@ class DynamicDecodeLayer : public BaseLayer std::optional parent_ids; // [max_seq_len, batch_size * beam_width], necessary in beam search std::optional sequence_length; // [batch_size * beam_width], optional std::optional - output_log_probs; // [request_ouptut_length, batch_size * beam_width], must be float*, optional + output_log_probs_tiled; // [request_output_length, batch_size, beam_width], must be float*, optional std::optional - tgt_cache_indirection; // [local_batch_size, beam_width, max_seq_len], the k/v cache index for beam search + output_log_probs; // [batchSize, beam_width, request_ouptut_length], must be float*, optional + std::optional + tgt_cache_indirection; // [local_batch_size, beam_width, max_seq_len], the k/v cache index for beam search std::shared_ptr - beamHypotheses; // a special structure which maintains some pointers of beam search + beamHypotheses; // a special structure which maintains some pointers of beam search - tc::Tensor output_ids_ptr; // [batch_size] int* (2-d array), each int* has [beam_width, max_seq_len] - tc::Tensor parent_ids_ptr; // [batch_size] int* (2-d array), each int* has [beam_width, max_seq_len] + tc::Tensor output_ids_ptr; // [batch_size] int* (2-d array), each int* has [beam_width, max_seq_len] + tc::Tensor parent_ids_ptr; // [batch_size] int* (2-d array), each int* has [beam_width, max_seq_len] }; void forward(OutputParams& outputs, ForwardParams const& params); diff --git a/cpp/tensorrt_llm/layers/fillBuffers.h b/cpp/tensorrt_llm/layers/fillBuffers.h new file mode 100644 index 000000000000..a5e377e64b7b --- /dev/null +++ b/cpp/tensorrt_llm/layers/fillBuffers.h @@ -0,0 +1,68 @@ +/* + * Copyright (c) 2019-2023, NVIDIA CORPORATION. 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. + */ + +#pragma once + +#include +#include +#include +#include + +#include + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/memoryUtils.h" + +namespace tensorrt_llm +{ +namespace layers +{ + +// Using a local lambda in beam search layers to fill buffers causes an internal compiler error on nvcc windows. +// As a workaround and to promote DRY, the fill logic is refactored into FillBuffers below. +struct FillBuffers +{ + + template + void operator()(std::optional> const& optParam, T const defaultValue, std::vector& hostBuffer, + T*& deviceBuffer) const + { + using tensorrt_llm::common::cudaAutoCpy; + + hostBuffer.resize(batch_size); + if (!optParam) + { + std::fill(std::begin(hostBuffer), std::end(hostBuffer), defaultValue); + } + else if (optParam->size() == 1) + { + std::fill(std::begin(hostBuffer), std::end(hostBuffer), optParam->front()); + } + else + { + TLLM_CHECK_WITH_INFO(optParam->size() == batch_size, "Argument vector size mismatch."); + std::copy(optParam->begin(), optParam->end(), std::begin(hostBuffer)); + } + cudaAutoCpy(deviceBuffer, hostBuffer.data(), batch_size, stream); + } + + size_t batch_size; + cudaStream_t stream; +}; + +} // namespace layers + +} // namespace tensorrt_llm diff --git a/cpp/tensorrt_llm/layers/onlineBeamSearchLayer.cu b/cpp/tensorrt_llm/layers/onlineBeamSearchLayer.cu index c0152c1f9b70..e9d5fff7899a 100644 --- a/cpp/tensorrt_llm/layers/onlineBeamSearchLayer.cu +++ b/cpp/tensorrt_llm/layers/onlineBeamSearchLayer.cu @@ -16,6 +16,7 @@ #include "tensorrt_llm/common/memoryUtils.h" #include "tensorrt_llm/kernels/beamSearchTopkKernels.h" +#include "tensorrt_llm/layers/fillBuffers.h" #include "tensorrt_llm/layers/onlineBeamSearchLayer.h" using namespace tensorrt_llm::common; @@ -49,7 +50,7 @@ __global__ void update_kernel(bool* finished, int** parent_ids_ptr, int* sequenc // Increase the seq_len even if the request has finished. // On the following iteration we check if the sequence has finished before - if (!finished[beam_idx]) + if (!finished[blockIdx.x * beam_width + beam_idx]) { s_sequence_lengths[beam_idx]++; } @@ -96,25 +97,7 @@ void OnlineBeamSearchLayer::setup(size_t batch_size, SetupParams const& setup mDiversityRate = setupParams.beam_search_diversity_rate.value_or(std::vector(0.0f)); mLengthPenalty = setupParams.length_penalty.value_or(std::vector(0.0f)); - auto fillBuffers - = [this, &batch_size](auto const& optParam, auto const defaultValue, auto& hostBuffer, auto& deviceBuffer) - { - hostBuffer.resize(batch_size); - if (!optParam) - { - std::fill(std::begin(hostBuffer), std::end(hostBuffer), defaultValue); - } - else if (optParam->size() == 1) - { - std::fill(std::begin(hostBuffer), std::end(hostBuffer), optParam->front()); - } - else - { - TLLM_CHECK_WITH_INFO(optParam->size() == batch_size, "Argument vector size mismatch."); - std::copy(optParam->begin(), optParam->end(), std::begin(hostBuffer)); - } - cudaAutoCpy(deviceBuffer, hostBuffer.data(), batch_size, stream_); - }; + FillBuffers const fillBuffers{batch_size, stream_}; fillBuffers(setupParams.beam_search_diversity_rate, 0.0f, mDiversityRate, diversity_rates_buf_); fillBuffers(setupParams.length_penalty, 0.0f, mLengthPenalty, length_penalties_buf_); @@ -153,7 +136,6 @@ void OnlineBeamSearchLayer::invokeSoftMax(BeamSearchOutputParams& outputs, So beamHypotheses.end_ids = end_ids; } - output_log_probs = (outputs.output_log_probs) ? outputs.output_log_probs->template getPtr() : nullptr; invokeTopkSoftMax(logits.template getPtr(), (const T*) (nullptr), finished, sequence_lengths, outputs.cum_log_probs->template getPtr(), output_log_probs, output_ids_ptr.getPtr(), topk_softmax_workspace_, topk_softmax_workspace_size_, &beamHypotheses, local_batch_size, beam_width, diff --git a/cpp/tensorrt_llm/plugins/CMakeLists.txt b/cpp/tensorrt_llm/plugins/CMakeLists.txt index 9bc273c412f8..0996a2de8362 100755 --- a/cpp/tensorrt_llm/plugins/CMakeLists.txt +++ b/cpp/tensorrt_llm/plugins/CMakeLists.txt @@ -44,7 +44,8 @@ set(PLUGIN_LISTS rmsnormQuantizationPlugin weightOnlyGroupwiseQuantMatmulPlugin weightOnlyQuantMatmulPlugin - lookupPlugin) + lookupPlugin + loraPlugin) foreach(PLUGIN_ITER ${PLUGIN_LISTS}) include_directories(${PLUGIN_ITER}) diff --git a/cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp b/cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp index aa0fde60211d..bbf724618d21 100644 --- a/cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp +++ b/cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp @@ -26,6 +26,7 @@ #include "tensorrt_llm/plugins/layernormPlugin/layernormPlugin.h" #include "tensorrt_llm/plugins/layernormQuantizationPlugin/layernormQuantizationPlugin.h" #include "tensorrt_llm/plugins/lookupPlugin/lookupPlugin.h" +#include "tensorrt_llm/plugins/loraPlugin/loraPlugin.h" #if ENABLE_MULTI_DEVICE #include "tensorrt_llm/plugins/ncclPlugin/allgatherPlugin.h" #include "tensorrt_llm/plugins/ncclPlugin/allreducePlugin.h" @@ -151,6 +152,7 @@ extern "C" weightOnlyGroupwiseQuantMatmulPluginCreator; static tensorrt_llm::plugins::WeightOnlyQuantMatmulPluginCreator weightOnlyQuantMatmulPluginCreator; static tensorrt_llm::plugins::LookupPluginCreator lookupPluginCreator; + static tensorrt_llm::plugins::LoraPluginCreator loraPluginCreator; static std::array pluginCreators = { creatorPtr(identityPluginCreator), @@ -173,6 +175,7 @@ extern "C" creatorPtr(weightOnlyGroupwiseQuantMatmulPluginCreator), creatorPtr(weightOnlyQuantMatmulPluginCreator), creatorPtr(lookupPluginCreator), + creatorPtr(loraPluginCreator), }; nbCreators = pluginCreators.size(); return pluginCreators.data(); diff --git a/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.cpp b/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.cpp index 09ede9f42388..8eb90b343bbb 100644 --- a/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.cpp +++ b/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.cpp @@ -18,6 +18,7 @@ #include "tensorrt_llm/kernels/decoderMaskedMultiheadAttention.h" #include "tensorrt_llm/kernels/gptKernels.h" #include "tensorrt_llm/kernels/unfusedAttentionKernels.h" +#include "tensorrt_llm/runtime/iBuffer.h" using namespace nvinfer1; using namespace tensorrt_llm::kernels; @@ -32,7 +33,8 @@ PluginFieldCollection BertAttentionPluginCreator::mFC{}; std::vector BertAttentionPluginCreator::mPluginAttributes; BertAttentionPlugin::BertAttentionPlugin(int num_heads, int head_size, float q_scaling, bool qk_half_accum, - ContextFMHAType context_fmha_type, nvinfer1::DataType type, bool do_relative_attention, int max_distance) + ContextFMHAType context_fmha_type, nvinfer1::DataType type, bool do_relative_attention, int max_distance, + bool remove_padding) : mNumHeads(num_heads) , mHeadSize(head_size) , mQScaling(q_scaling) @@ -42,6 +44,7 @@ BertAttentionPlugin::BertAttentionPlugin(int num_heads, int head_size, float q_s , mType(type) , mRelativeAttention(do_relative_attention) , mMaxDistance(max_distance) + , mRemovePadding(remove_padding) { // pre-check whether FMHA is supported in order to save memory allocation mEnableContextFMHA = mEnableContextFMHA && (mType == DataType::kHALF) && MHARunner::fmha_supported(mHeadSize, mSM); @@ -60,6 +63,7 @@ BertAttentionPlugin::BertAttentionPlugin(const void* data, size_t length) read(d, mType); read(d, mRelativeAttention); read(d, mMaxDistance); + read(d, mRemovePadding); TLLM_CHECK(d == a + length); } @@ -84,13 +88,29 @@ nvinfer1::DimsExprs BertAttentionPlugin::getOutputDimensions( bool BertAttentionPlugin::supportsFormatCombination( int pos, const nvinfer1::PluginTensorDesc* inOut, int nbInputs, int nbOutputs) noexcept { - if (pos == 1) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; + // inputs: [0] qkv, [1] input_lengths, [2] max_input_length (optional), [3] relative_attention_bias (optional) + // outputs: [X] hidden_states + if (nbInputs == 2) + { // BERT + if (pos == 1) + { + return inOut[pos].type == nvinfer1::DataType::kINT32; + } + else + { + return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); + } } - else - { - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); + else if (nbInputs > 2) + { // Encoder in encoder-decoder + if (pos == 1 || pos == 2) + { + return inOut[pos].type == nvinfer1::DataType::kINT32; + } + else + { + return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); + } } } @@ -102,25 +122,17 @@ void BertAttentionPlugin::configurePlugin(const nvinfer1::DynamicPluginTensorDes size_t BertAttentionPlugin::getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, int nbInputs, const nvinfer1::PluginTensorDesc* outputs, int nbOutputs) const noexcept { - const int batch_size = inputs[0].dims.d[0]; - const int input_seq_len = inputs[0].dims.d[1]; + // if remove padding, inputs[0] "qkv_hidden_states" dim is [1, num_tokens, 3*hidden_dim] which doesn't have shape + // info should get max_batch_size and max_input_length from inputs[1] "input_lengths" and input[2] + // "max_input_length" + const int batch_size = mRemovePadding ? inputs[1].dims.d[0] : inputs[0].dims.d[0]; + const int input_seq_len = mRemovePadding ? inputs[2].dims.d[0] : inputs[0].dims.d[1]; const int local_hidden_units_ = inputs[0].dims.d[2] / 3; - const int beam_width = 1; - const int max_input_length = input_seq_len; - size_t size{0U}; - if (inputs[0].type == DataType::kHALF) - { - size = sizeof(half); - } - else if (inputs[0].type == DataType::kFLOAT) - { - size = sizeof(float); - } + auto const size = tensorrt_llm::runtime::BufferDataType(inputs[0].type).getSize(); - const size_t attention_mask_size - = mEnableContextFMHA ? 0 : size * batch_size * beam_width * max_input_length * max_input_length; - const size_t cu_seqlens_size = sizeof(int) * (batch_size * beam_width + 1); + const size_t attention_mask_size = mEnableContextFMHA ? 0 : size * batch_size * input_seq_len * input_seq_len; + const size_t cu_seqlens_size = sizeof(int) * (batch_size + 1); const size_t q_buf_2_size = size * batch_size * input_seq_len * local_hidden_units_; const size_t k_buf_2_size = size * batch_size * input_seq_len * local_hidden_units_; const size_t v_buf_2_size = size * batch_size * input_seq_len * local_hidden_units_; @@ -153,25 +165,27 @@ int BertAttentionPlugin::enqueueImpl(const nvinfer1::PluginTensorDesc* inputDesc { // inputs - // input_tensor [batch_size, seq_len, local_hidden_size * 3] + // input_tensor [batch_size, seq_len, local_hidden_size*3] or [1, num_tokens, local_hidden_size*3] // input_lengths [batch_size] - // relative_attention_bias [num_heads, num_buckets] (optional) + // max_input_length [max_input_length] -- use shape dim to represent max value. If remove padding, this records + // the max input length among sequences; otherwise same as input_tensor's padded dim[1] relative_attention_bias + // [num_heads, num_buckets] (optional) // outputs - // output_tensor [batch_size, seq_len, local_hidden_size] + // output_tensor [batch_size, seq_len, local_hidden_size] or [1, num_tokens, local_hidden_size] - const int batch_size = inputDesc[0].dims.d[0]; + // if remove padding, inputs[0] dim is [1, num_tokens] which doesn't have workspace info + // should get max_batch_size from inputs[1] and max_input_length from plugin attribute + const int batch_size = mRemovePadding ? inputDesc[1].dims.d[0] : inputDesc[0].dims.d[0]; + const int input_seq_len = mRemovePadding ? inputDesc[2].dims.d[0] : inputDesc[0].dims.d[1]; + const int num_tokens = mRemovePadding ? inputDesc[0].dims.d[1] : batch_size * input_seq_len; const int request_batch_size = batch_size; - const int input_seq_len = inputDesc[0].dims.d[1]; const int request_seq_len = input_seq_len; - const int beam_width = 1; const int local_hidden_units_ = inputDesc[0].dims.d[2] / 3; const float q_scaling = mQScaling; const T* attention_input = reinterpret_cast(inputs[0]); - const int* input_lengths = reinterpret_cast(inputs[1]); - const T* relative_attn_table = mRelativeAttention ? reinterpret_cast(inputs[2]) : nullptr; - + const T* relative_attn_table = mRelativeAttention ? reinterpret_cast(inputs[3]) : nullptr; T* context_buf_ = (T*) (outputs[0]); auto cublasHandle = mCublasWrapper->getCublasHandle(); @@ -186,10 +200,15 @@ int BertAttentionPlugin::enqueueImpl(const nvinfer1::PluginTensorDesc* inputDesc { mCublasWrapper->setFP32GemmConfig(); } +#ifdef ENABLE_BF16 + else if constexpr (std::is_same_v) + { + mCublasWrapper->setBF16GemmConfig(); + } +#endif - const size_t attention_mask_size - = mEnableContextFMHA ? 0 : sizeof(T) * batch_size * beam_width * input_seq_len * input_seq_len; - const size_t cu_seqlens_size = sizeof(int) * (batch_size * beam_width + 1); + const size_t attention_mask_size = mEnableContextFMHA ? 0 : sizeof(T) * batch_size * input_seq_len * input_seq_len; + const size_t cu_seqlens_size = sizeof(int) * (batch_size + 1); const size_t q_buf_2_size = sizeof(T) * batch_size * input_seq_len * local_hidden_units_; const size_t k_buf_2_size = sizeof(T) * batch_size * input_seq_len * local_hidden_units_; const size_t v_buf_2_size = sizeof(T) * batch_size * input_seq_len * local_hidden_units_; @@ -200,8 +219,6 @@ int BertAttentionPlugin::enqueueImpl(const nvinfer1::PluginTensorDesc* inputDesc = mEnableContextFMHA ? 0 : sizeof(float) * batch_size * mNumHeads * input_seq_len * input_seq_len; const size_t padding_offset_size = sizeof(int) * batch_size * input_seq_len; - mMaxInputLength = input_seq_len; - // Workspace pointer shift int8_t* workspace_byte_ptr = reinterpret_cast(workspace); size_t offset = CUBLAS_WORKSPACE_SIZE; @@ -223,17 +240,17 @@ int BertAttentionPlugin::enqueueImpl(const nvinfer1::PluginTensorDesc* inputDesc params.paddingOffsets = padding_offset; params.attentionMask = attention_mask; params.seqLengths = input_lengths; - params.batchSize = batch_size * beam_width; - params.maxSeqLength = mMaxInputLength; - params.numTokens = batch_size * beam_width * mMaxInputLength; + params.batchSize = batch_size; + params.maxSeqLength = input_seq_len; + params.numTokens = num_tokens; params.attentionMaskType = AttentionMaskType::PADDING; invokeBuildDecoderInfo(params, stream); + sync_check_cuda_error(); - // Padding offset = nullptr here (remove padding is not supported). invokeAddFusedQKVBiasTranspose(q_buf_2_, k_buf_2_, v_buf_2_, const_cast(attention_input), input_lengths, - nullptr, request_batch_size, request_seq_len, batch_size * input_seq_len, mNumHeads, mNumHeads, mHeadSize, - mEnableContextFMHA, 0, 0.0f, RotaryScalingType::kNONE, 0.0f, 0, PositionEmbeddingType::kLEARNED_ABSOLUTE, - (float*) nullptr, 0, stream); + mRemovePadding ? padding_offset : nullptr, batch_size, input_seq_len, num_tokens, mNumHeads, mNumHeads, + mHeadSize, mEnableContextFMHA, 0, 0.0f, RotaryScalingType::kNONE, 0.0f, 0, + PositionEmbeddingType::kLEARNED_ABSOLUTE, (float*) nullptr, 0, stream); const auto gemm_data_type = tc::CudaDataType::value; const int attention_seq_len_1 = request_seq_len; // q length @@ -285,7 +302,7 @@ int BertAttentionPlugin::enqueueImpl(const nvinfer1::PluginTensorDesc* inputDesc // [num_heads, num_buckets], with necessary params (max_distance, num_buckets) passed at the end invokeAddRelativeAttentionBiasUnaligned(qk_buf_float_, relative_attn_table, request_batch_size, mNumHeads, attention_seq_len_1, attention_seq_len_2, stream, mMaxDistance > 0, - inputDesc[2].dims.d[1], mMaxDistance, true /* bidirectional */); + inputDesc[3].dims.d[1], mMaxDistance, true /* bidirectional */); } MaskedSoftmaxParam param; @@ -317,7 +334,7 @@ int BertAttentionPlugin::enqueueImpl(const nvinfer1::PluginTensorDesc* inputDesc // max_output_len + 1. In implicit mode, relative_attention_bias is rel attn table // [num_heads, num_buckets], with necessary params (max_distance, num_buckets) passed at the end invokeAddRelativeAttentionBiasUnaligned(qk_buf_, relative_attn_table, request_batch_size, mNumHeads, - attention_seq_len_1, attention_seq_len_2, stream, mMaxDistance > 0, inputDesc[2].dims.d[1], + attention_seq_len_1, attention_seq_len_2, stream, mMaxDistance > 0, inputDesc[3].dims.d[1], mMaxDistance, true /* bidirectional */); } @@ -339,16 +356,15 @@ int BertAttentionPlugin::enqueueImpl(const nvinfer1::PluginTensorDesc* inputDesc attention_seq_len_1 * attention_seq_len_2, qkv_buf_2_, mHeadSize, attention_seq_len_1 * mHeadSize, request_batch_size * mNumHeads); - if (padding_offset == nullptr) + if (!mRemovePadding) { invokeTransposeQKV(context_buf_, qkv_buf_2_, request_batch_size, attention_seq_len_1, mNumHeads, mHeadSize, (float*) nullptr, 0, stream); } else { - invokeTransposeAttentionOutRemovePadding(qkv_buf_2_, context_buf_, batch_size * input_seq_len, - request_batch_size, attention_seq_len_1, mNumHeads, mHeadSize, padding_offset, (float*) nullptr, 0, - stream); + invokeTransposeAttentionOutRemovePadding(qkv_buf_2_, context_buf_, num_tokens, request_batch_size, + request_seq_len, mNumHeads, mHeadSize, padding_offset, (float*) nullptr, 0, stream); } } return 0; @@ -362,6 +378,12 @@ template int BertAttentionPlugin::enqueueImpl(const nvinfer1::PluginTenso const nvinfer1::PluginTensorDesc* outputDesc, const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); +#ifdef ENABLE_BF16 +template int BertAttentionPlugin::enqueueImpl<__nv_bfloat16>(const nvinfer1::PluginTensorDesc* inputDesc, + const nvinfer1::PluginTensorDesc* outputDesc, const void* const* inputs, void* const* outputs, void* workspace, + cudaStream_t stream); +#endif + int BertAttentionPlugin::enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc, const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept @@ -374,6 +396,12 @@ int BertAttentionPlugin::enqueue(const nvinfer1::PluginTensorDesc* inputDesc, { return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); } +#ifdef ENABLE_BF16 + else if (mType == DataType::kBF16) + { + return enqueueImpl<__nv_bfloat16>(inputDesc, outputDesc, inputs, outputs, workspace, stream); + } +#endif return 0; } @@ -425,7 +453,8 @@ void BertAttentionPlugin::destroy() noexcept size_t BertAttentionPlugin::getSerializationSize() const noexcept { return sizeof(mNumHeads) + sizeof(mHeadSize) + sizeof(mQScaling) + sizeof(mQKHalfAccum) + sizeof(mEnableContextFMHA) - + sizeof(mFMHAForceFP32Acc) + sizeof(mType) + sizeof(mRelativeAttention) + sizeof(mMaxDistance); + + sizeof(mFMHAForceFP32Acc) + sizeof(mType) + sizeof(mRelativeAttention) + sizeof(mMaxDistance) + + sizeof(mRemovePadding); } void BertAttentionPlugin::serialize(void* buffer) const noexcept @@ -440,6 +469,7 @@ void BertAttentionPlugin::serialize(void* buffer) const noexcept write(d, mType); write(d, mRelativeAttention); write(d, mMaxDistance); + write(d, mRemovePadding); assert(d == a + getSerializationSize()); } @@ -459,6 +489,7 @@ BertAttentionPluginCreator::BertAttentionPluginCreator() mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32, 1)); mPluginAttributes.emplace_back(PluginField("do_relative_attention", nullptr, PluginFieldType::kINT8, 0)); mPluginAttributes.emplace_back(PluginField("max_distance", nullptr, PluginFieldType::kINT32, 0)); + mPluginAttributes.emplace_back(PluginField("remove_padding", nullptr, PluginFieldType::kINT8, 0)); mFC.nbFields = mPluginAttributes.size(); mFC.fields = mPluginAttributes.data(); } @@ -488,6 +519,7 @@ IPluginV2* BertAttentionPluginCreator::createPlugin(const char* name, const Plug nvinfer1::DataType type; bool do_relative_attention; int max_distance; + bool remove_padding; // Read configurations from each fields for (int i = 0; i < fc->nbFields; ++i) { @@ -532,11 +564,16 @@ IPluginV2* BertAttentionPluginCreator::createPlugin(const char* name, const Plug TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); max_distance = static_cast(*(static_cast(fields[i].data))); } + else if (!strcmp(attrName, "remove_padding")) + { + TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); + remove_padding = static_cast(*(static_cast(fields[i].data))); + } } try { auto* obj = new BertAttentionPlugin(num_heads, head_size, q_scaling, qk_half_accum, context_fmha_type, type, - do_relative_attention, max_distance); + do_relative_attention, max_distance, remove_padding); obj->setPluginNamespace(mNamespace.c_str()); return obj; } diff --git a/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.h b/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.h index c475e3b3c0be..151022d953f9 100644 --- a/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.h +++ b/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.h @@ -36,7 +36,7 @@ class BertAttentionPlugin : public BasePlugin BertAttentionPlugin(int num_heads, int head_size, float q_scaling, bool qk_half_accum, tensorrt_llm::kernels::ContextFMHAType context_fmha_type, nvinfer1::DataType type, - bool do_relative_attention = false, int max_distance = 0); + bool do_relative_attention = false, int max_distance = 0, bool remove_padding = false); BertAttentionPlugin(const void* data, size_t length); @@ -78,11 +78,11 @@ class BertAttentionPlugin : public BasePlugin int mNumHeads; int mHeadSize; - int mMaxInputLength; float mQScaling; nvinfer1::DataType mType; bool mRelativeAttention = false; int mMaxDistance = 0; + bool mRemovePadding = false; // unfused mha bool mQKHalfAccum = false; diff --git a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp b/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp index 26df67deb4c0..88b2551c4043 100644 --- a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp +++ b/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp @@ -217,7 +217,7 @@ void fusedQKV_masked_attention_dispatch(Multihead_attention_params&, \ const FusedQKVMaskedAttentionDispatchParams&, cudaStream_t stream); \ template void fusedQKV_masked_attention_dispatch(Multihead_attention_params&, \ @@ -226,12 +226,12 @@ void fusedQKV_masked_attention_dispatch(Multihead_attention_params&, cudaStream_t stream); \ template void fusedQKV_masked_attention_dispatch(Multihead_attention_params&, \ const FusedQKVMaskedAttentionDispatchParams&, cudaStream_t stream); -INSTANTIATE_MMHA_DISPATH(float, float) -INSTANTIATE_MMHA_DISPATH(uint16_t, half) +INSTANTIATE_MMHA_DISPATCH(float, float) +INSTANTIATE_MMHA_DISPATCH(uint16_t, half) #ifdef ENABLE_BF16 -INSTANTIATE_MMHA_DISPATH(__nv_bfloat16, __nv_bfloat16) +INSTANTIATE_MMHA_DISPATCH(__nv_bfloat16, __nv_bfloat16) #endif -#undef INSTANTIATE_MMHA_DISPATH +#undef INSTANTIATE_MMHA_DISPATCH GPTAttentionPluginCommon::GPTAttentionPluginCommon(int num_heads, int num_kv_heads, int head_size, int unidirectional, float q_scaling, tensorrt_llm::kernels::PositionEmbeddingType position_embedding_type, @@ -498,7 +498,8 @@ int GPTAttentionPluginCommon::enqueueContext(const EnqueueContextParams(nextWorkspacePtr(workspace_byte_ptr, offset, padding_offset_size)); // build attention_mask, cu_seqlens, and padding_offset tensors + // Note: self attn and cross attn should use different params + // cross attn's seqlen info is from encoder input lengths, not decoder input lengths! + // moreover, attn mask for cross attn should be set separately (see below) BuildDecoderInfoParams decoder_params; memset(&decoder_params, 0, sizeof(decoder_params)); decoder_params.seqOffsets = cu_seqlens; decoder_params.paddingOffsets = padding_offset; - decoder_params.attentionMask = attention_mask; - decoder_params.seqLengths = params.context_lengths; + decoder_params.attentionMask = isCrossAttention() ? nullptr : attention_mask; // manually set for cross attn + decoder_params.seqLengths = isCrossAttention() ? params.encoder_input_lengths : params.context_lengths; decoder_params.batchSize = params.batch_size; - decoder_params.maxSeqLength = params.input_seq_length; + decoder_params.maxSeqLength = isCrossAttention() ? params.cross_qkv_length : params.input_seq_length; decoder_params.maxKvCacheLength = params.cyclic_kv_cache_length; decoder_params.numTokens = params.num_tokens; decoder_params.attentionMaskType = mMaskType; @@ -533,9 +537,27 @@ int GPTAttentionPluginCommon::enqueueContext(const EnqueueContextParams h_attention_mask(params.batch_size * params.cross_qkv_length * params.input_seq_length, 1.); + std::vector h_attention_mask(params.batch_size * params.input_seq_length * params.cross_qkv_length, 1.); + std::vector h_encoder_input_lengths(params.batch_size); + cudaMemcpyAsync(h_encoder_input_lengths.data(), params.encoder_input_lengths, + sizeof(int32_t) * params.batch_size, cudaMemcpyDeviceToHost, stream); + for (int bi = 0; bi < params.batch_size; bi++) + { + int b_offset = bi * params.input_seq_length * params.cross_qkv_length; + for (int qi = 0; qi < params.input_seq_length; qi++) + { + int q_offset = b_offset + qi * params.cross_qkv_length; + if (h_encoder_input_lengths[bi] < params.cross_qkv_length) + { + std::fill(h_attention_mask.begin() + q_offset + h_encoder_input_lengths[bi], + h_attention_mask.begin() + q_offset + params.cross_qkv_length, 0.f); + } + } + } cudaMemcpyAsync(attention_mask, h_attention_mask.data(), sizeof(T) * params.batch_size * params.cross_qkv_length * params.input_seq_length, cudaMemcpyHostToDevice, stream); @@ -580,6 +602,8 @@ int GPTAttentionPluginCommon::enqueueContext(const EnqueueContextParams(qk_buf_) - reinterpret_cast(k_buf_2_), + // stream); cudaMemsetAsync(k_buf_2_, 0, reinterpret_cast(v_buf_2_) - reinterpret_cast(k_buf_2_) + v_buf_2_size, stream); @@ -698,23 +722,22 @@ int GPTAttentionPluginCommon::enqueueContext(const EnqueueContextParams 0, - relative_attention_bias_stride, max_distance, true /* bidirectional */); - } - if (is_qk_buf_float_ == true) { + // add relative position bias + if (isRelativePosition()) + { + // Add relative_attention_bias + // QK is (batch_size, local_head_num, q_length, k_length), relative_attention_bias is (1, + // local_head_num, max_output_len + 1, max_output_len + 1). broadcast along 1st dim. max_seq_len is + // already max_output_len + 1. In implicit mode, relative_attention_bias is relative_attention_table + // [num_heads, num_buckets], with necessary params (max_distance, num_buckets) passed at the end + invokeAddRelativeAttentionBiasUnaligned(qk_buf_float_, relative_attention_bias, params.batch_size, + mNumHeads, attention_seq_len_1, + isCrossAttention() ? params.cross_qkv_length : params.cyclic_kv_cache_length, stream, + max_distance > 0, relative_attention_bias_stride, max_distance, true /* bidirectional */); + } + MaskedSoftmaxParam param; param.attention_score = qk_buf_; // (batch_size, head_num, q_length, k_length) param.qk = qk_buf_float_; // (batch_size, head_num, q_length, k_length) @@ -729,6 +752,19 @@ int GPTAttentionPluginCommon::enqueueContext(const EnqueueContextParams 0, relative_attention_bias_stride, max_distance, true /* bidirectional */); + } + MaskedSoftmaxParam param; param.attention_score = qk_buf_; // (batch_size, head_num, q_length, k_length) param.qk = qk_buf_; // (batch_size, head_num, q_length, k_length) @@ -935,7 +971,6 @@ int GPTAttentionPluginCommon::enqueueGeneration( dispatch_params.input_lengths = params.context_lengths; dispatch_params.step = step; dispatch_params.q_scaling = q_scaling; - dispatch_params.relative_attention_bias_stride = relative_attention_bias_stride; dispatch_params.linear_bias_slopes = isALiBi() ? params.alibi_slopes : nullptr; dispatch_params.ia3_tasks = ia3_tasks; dispatch_params.ia3_key_weights = ia3_key_weights; diff --git a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp b/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp index 0633762c1e52..6f1d9ccec4ad 100644 --- a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp +++ b/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp @@ -250,7 +250,7 @@ int GPTAttentionPlugin::enqueueSome(int32_t seqIdxBeg, int32_t localNbSeq, int32 // such model has an encoder context (for cross attn) and an decoder context (for self attn) // clarify 3 lens: // -- max_context_len: len of decoder input. No "max" concept, it's what it is given. - // Also called (decoder_)input_seq_length + // Also called (decoder_)input_seq_length, normally 1 for encoder-decoder start token // -- max_seq_len: max allowed len of decoder output, i.e. final results // -- max_encoder_context_len: len of encoder input (in cross attn). Also called encoder_input_seq_length diff --git a/cpp/tensorrt_llm/plugins/loraPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/loraPlugin/CMakeLists.txt new file mode 100644 index 000000000000..86876224fccd --- /dev/null +++ b/cpp/tensorrt_llm/plugins/loraPlugin/CMakeLists.txt @@ -0,0 +1,21 @@ +# +# 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. +# +file(GLOB SRCS *.cpp) +set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) +set(PLUGIN_SOURCES + ${PLUGIN_SOURCES} + PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.cpp b/cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.cpp new file mode 100644 index 000000000000..aced835f94f8 --- /dev/null +++ b/cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.cpp @@ -0,0 +1,573 @@ +/* + * 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 "loraPlugin.h" +#include "tensorrt_llm/common/memoryUtils.h" +#include "tensorrt_llm/runtime/iBuffer.h" + +using namespace nvinfer1; +using namespace tensorrt_llm::common; +using tensorrt_llm::plugins::LoraPluginCreator; +using tensorrt_llm::plugins::LoraPlugin; +using tensorrt_llm::plugins::CublasGemmWrapperPtr; +using tensorrt_llm::plugins::read; +using tensorrt_llm::plugins::write; + +static const char* LORA_PLUGIN_VERSION{"1"}; +static const char* LORA_PLUGIN_NAME{"Lora"}; +PluginFieldCollection LoraPluginCreator::mFC{}; +std::vector LoraPluginCreator::mPluginAttributes; + +// TODO should reuse the function in gemmPlugin +void _getProblemParams(cublasOperation_t& transa, cublasOperation_t& transb, int& m, int& n, int& k, int& lda, int& ldb, + int& ldc, bool transA, bool transB, int M, int N, int K) +{ + transa = transB ? CUBLAS_OP_T : CUBLAS_OP_N; + transb = transA ? CUBLAS_OP_T : CUBLAS_OP_N; + m = N; + n = M; + k = K; + lda = transB ? K : N; + ldb = transA ? M : K; + ldc = N; +} + +// TODO should reuse the function in gemmPlugin +void _runGemm(const int M, const int N, const int K, const bool transA, const bool transB, + const nvinfer1::DataType type, const CublasGemmWrapperPtr& cublasWrapperPtr, const void* act, const void* weight, + void* output, const std::optional& heuristic, void* workspace, cudaStream_t stream) +{ + cublasWrapperPtr->setStream(stream); + cublasWrapperPtr->setWorkspace(workspace); + + cublasOperation_t transa, transb; + int m, n, k; + int lda, ldb, ldc; + _getProblemParams(transa, transb, m, n, k, lda, ldb, ldc, transA, transB, M, N, K); + + cublasWrapperPtr->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); + cublasWrapperPtr->Gemm(transa, transb, m, n, k, weight, lda, act, ldb, output, ldc, heuristic); + cublasWrapperPtr->destroyDescriptors(); +} + +LoraPlugin::LoraPlugin(int in_hidden_size, int out_hidden_size, int transA, int transB, int lora_module_number, + nvinfer1::DataType type, const LoraPlugin::PluginProfilerPtr& pluginProfiler, bool remove_input_padding, + int max_context_length, int max_low_rank) + : mInHiddenSize(in_hidden_size) + , mOutHiddenSize(out_hidden_size) + , mTransA(transA) + , mTransB(transB) + , mType(type) + , mPluginProfiler(pluginProfiler) + , mRemoveInputPadding(remove_input_padding) + , mMaxContextLength(max_context_length) + , mMaxLowRank(max_low_rank) +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + init(); +} + +// Parameterized constructor +LoraPlugin::LoraPlugin(const void* data, size_t length, const LoraPlugin::PluginProfilerPtr& pluginProfiler) + : mPluginProfiler(pluginProfiler) +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + const char *d = reinterpret_cast(data), *a = d; + read(d, mInHiddenSize); + read(d, mOutHiddenSize); + read(d, mTransA); + read(d, mTransB); + read(d, mType); + read(d, mRemoveInputPadding); + read(d, mMaxContextLength); + read(d, mMaxLowRank); + + init(); + + mPluginProfiler->deserialize(d, mDims, mGemmId); + + TLLM_CHECK(d == a + length); +} + +void LoraPlugin::init() +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + auto cublasHandle = getCublasHandle(); + auto cublasLtHandle = getCublasLtHandle(); + mCublasWrapper = std::make_shared(cublasHandle, cublasLtHandle, nullptr, nullptr); + + mPluginProfiler->setTranspose(mTransA, mTransB); + + mGemmId = GemmIdCublas(mDims.n, mDims.k, mType, mTransA, mTransB); +} + +void LoraPlugin::setGemmConfig() +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + if (mType == DataType::kHALF) + { + mCublasWrapper->setFP16GemmConfig(); + } + else if (mType == DataType::kFLOAT) + { + mCublasWrapper->setFP32GemmConfig(); + } +#ifdef ENABLE_BF16 + else if (mType == DataType::kBF16) + { + mCublasWrapper->setBF16GemmConfig(); + } +#endif +} + +void LoraPlugin::configGemm() +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + if (!mDims.isInitialized()) + { + return; + } + + setGemmConfig(); + + mPluginProfiler->profileTactics(mCublasWrapper, mType, mDims, mGemmId); +} + +// IPluginV2DynamicExt Methods +nvinfer1::IPluginV2DynamicExt* LoraPlugin::clone() const noexcept +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + auto* plugin = new LoraPlugin(*this); + return plugin; +} + +nvinfer1::DimsExprs LoraPlugin::getOutputDimensions( + int outputIndex, const nvinfer1::DimsExprs* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + try + { + TLLM_CHECK(outputIndex == 0); + const int nbDimsA = inputs[getInputTensorIdx()].nbDims; + DimsExprs ret; + ret.nbDims = nbDimsA; + + for (int i = 0; i < ret.nbDims; ++i) + { + ret.d[0] = 0; + } + + if (mTransA) + { + for (int i = 1; i < nbDimsA; ++i) + { + ret.d[i - 1] = inputs[getInputTensorIdx()].d[i]; + } + } + else + { + for (int i = 0; i < nbDimsA - 1; ++i) + { + ret.d[i] = inputs[getInputTensorIdx()].d[i]; + } + } + + auto const* outHiddenSize = exprBuilder.constant(mOutHiddenSize); + TLLM_CHECK(outHiddenSize != nullptr); + ret.d[ret.nbDims - 1] = outHiddenSize; + return ret; + } + catch (const std::exception& e) + { + caughtError(e); + } + return DimsExprs{}; +} + +bool LoraPlugin::supportsFormatCombination( + int pos, const nvinfer1::PluginTensorDesc* inOut, int nbInputs, int nbOutputs) noexcept +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + if (pos == getHostRequestTypesIdx()) + { + return inOut[pos].type == nvinfer1::DataType::kINT32; + } + else if (pos == getLoraRanksIdx()) + { + return inOut[pos].type == nvinfer1::DataType::kINT32; + } + else if (pos == getLoraWeightsPtrsIdx()) + { + return inOut[pos].type == nvinfer1::DataType::kINT64; + } + else if (mRemoveInputPadding && pos == getHostContextLengthsIdx()) + { + return inOut[pos].type == nvinfer1::DataType::kINT32; + } + else + { + return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); + } +} + +int32_t _computeMDimension(bool transA, const int32_t nbDims, const int32_t* dims) +{ + int32_t M = 1; + if (transA) + { + for (int i = nbDims - 1; i > 0; --i) + { + M *= dims[i]; + } + } + else + { + for (int i = 0; i < nbDims - 1; ++i) + { + M *= dims[i]; + } + } + return M; +} + +int32_t _computeNDimension(bool transB, const int32_t nbDims, const int32_t* dims) +{ + int32_t N = 1; + if (transB) + { + for (int i = 0; i < nbDims - 1; ++i) + { + N *= dims[i]; + } + } + else + { + for (int i = nbDims - 1; i > 0; --i) + { + N *= dims[i]; + } + } + return N; +} + +void LoraPlugin::configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in, int nbInputs, + const nvinfer1::DynamicPluginTensorDesc* out, int nbOutputs) noexcept +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + const int nbDimsA = in[0].max.nbDims; + const int nbDimsB = in[1].max.nbDims; + + const auto minM = _computeMDimension(mTransA, nbDimsA, in[0].min.d); + const auto maxM = _computeMDimension(mTransA, nbDimsA, in[0].max.d); + const auto N = _computeNDimension(mTransB, nbDimsB, in[1].max.d); + const auto K = mTransA ? in[0].max.d[0] : in[0].max.d[nbDimsA - 1]; + + if (!mDims.isInitialized()) + { + mDims = {minM, maxM, N, K}; + } + mGemmId.n = N; + mGemmId.k = K; +} + +size_t LoraPlugin::getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, int nbInputs, + const nvinfer1::PluginTensorDesc* outputs, int nbOutputs) const noexcept +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + const int nbReq = inputs[getLoraRanksIdx()].dims.d[0]; + auto const type = inputs[getInputTensorIdx()].type; + auto const typeSize = tensorrt_llm::runtime::BufferDataType(type).getSize(); + + size_t const lowRankWorkSpaceSize = nbReq * mMaxContextLength * mMaxLowRank * typeSize; + + return CUBLAS_WORKSPACE_SIZE + lowRankWorkSpaceSize; +} + +int LoraPlugin::enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc, + const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + // inputs + // input [-1, K] (view as 2D) + // host_request_type [batch_size] on cpu + // lora_ranks [batch_size] on cpu + // lora_weights_ptr [batch_size, 2] on cpu + // host_context_lengths [batch_size] on cpu + // outputs + // output [-1, N] (view as 2D) + + auto const typeSize = tensorrt_llm::runtime::BufferDataType(mType).getSize(); + void* cublasWorkSpace = workspace; + void* lowRankWorkSpace = static_cast(cublasWorkSpace) + CUBLAS_WORKSPACE_SIZE; + + setGemmConfig(); + auto const batch_size = inputDesc[getLoraRanksIdx()].dims.d[0]; + auto const lora_ranks = static_cast(inputs[getLoraRanksIdx()]); + auto const lora_weights_ptr = static_cast(inputs[getLoraWeightsPtrsIdx()]); + auto const host_context_lengths + = mRemoveInputPadding ? static_cast(inputs[getHostContextLengthsIdx()]) : nullptr; + RequestType const* reqTypes = static_cast(inputs[getHostRequestTypesIdx()]); + + size_t handled_token_num = 0; + for (int batchIdx = 0; batchIdx < batch_size; batchIdx++) + { + const RequestType reqType = reqTypes[batchIdx]; + const auto M = (reqType != RequestType::kCONTEXT) + ? 1 + : (mRemoveInputPadding ? host_context_lengths[batchIdx] : inputDesc[0].dims.d[1]); + const auto lora_rank = lora_ranks[batchIdx]; + + if (lora_rank <= 0) + { + const auto N = outputDesc[0].dims.d[outputDesc[0].dims.nbDims - 1]; + void* output = static_cast(static_cast(outputs[0]) + handled_token_num * N * typeSize); + if (typeSize == 2) + { + deviceFill((half*) output, M * N, (half) 0.0f, stream); + } + else + { + deviceFill((float*) output, M * N, 0.0f, stream); + } + } + else + { + // the input shape should be [1, token_num, K] under remove_input_padding, + // [batch, seqlen, K] under non-remove_input_padding + auto bestTactic = mPluginProfiler->getBestConfig(M, mGemmId); + + const int nbDimsA = inputDesc[0].dims.nbDims; + const auto N = lora_rank; + + TLLM_CHECK_WITH_INFO(N <= mMaxLowRank, + fmtstr("Invalid low_rank (%d). low_rank must be smaller than mMaxLowRank (%d)", N, mMaxLowRank)); + const auto K = mTransA ? inputDesc[0].dims.d[0] : inputDesc[0].dims.d[nbDimsA - 1]; // input hidden size + const auto N2 = outputDesc[0].dims.d[nbDimsA - 1]; + // [M, K] -> [M, N] -> [M, N2] + + void* lora_in_weight = reinterpret_cast(lora_weights_ptr[batchIdx * 2 + 0]); + void* lora_out_weight = reinterpret_cast(lora_weights_ptr[batchIdx * 2 + 1]); + const void* input + = static_cast(static_cast(inputs[0]) + handled_token_num * K * typeSize); + void* output = static_cast(static_cast(outputs[0]) + handled_token_num * N2 * typeSize); + _runGemm(M, N, K, mTransA, mTransB, mType, mCublasWrapper, input, lora_in_weight, lowRankWorkSpace, + bestTactic, cublasWorkSpace, stream); + + _runGemm(M, N2, N, mTransA, mTransB, mType, mCublasWrapper, lowRankWorkSpace, lora_out_weight, output, + bestTactic, cublasWorkSpace, stream); + } + handled_token_num += M; + } + return 0; +} + +// IPluginV2Ext Methods +nvinfer1::DataType LoraPlugin::getOutputDataType( + int index, const nvinfer1::DataType* inputTypes, int nbInputs) const noexcept +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + TLLM_CHECK(index == 0); + return inputTypes[0]; +} + +// IPluginV2 Methods + +const char* LoraPlugin::getPluginType() const noexcept +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + return LORA_PLUGIN_NAME; +} + +const char* LoraPlugin::getPluginVersion() const noexcept +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + return LORA_PLUGIN_VERSION; +} + +int LoraPlugin::getNbOutputs() const noexcept +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + return 1; +} + +int LoraPlugin::initialize() noexcept +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + configGemm(); + return 0; +} + +void LoraPlugin::destroy() noexcept +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + delete this; +} + +size_t LoraPlugin::getSerializationSize() const noexcept +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + return sizeof(mInHiddenSize) + sizeof(mOutHiddenSize) + sizeof(mTransA) + sizeof(mTransB) + sizeof(mType) + + mPluginProfiler->getSerializationSize(mGemmId) + sizeof(mRemoveInputPadding) + sizeof(mMaxContextLength) + + sizeof(mMaxLowRank); // selected tactics container size +} + +void LoraPlugin::serialize(void* buffer) const noexcept +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + char *d = static_cast(buffer), *a = d; + write(d, mInHiddenSize); + write(d, mOutHiddenSize); + write(d, mTransA); + write(d, mTransB); + write(d, mType); + write(d, mRemoveInputPadding); + write(d, mMaxContextLength); + write(d, mMaxLowRank); + mPluginProfiler->serialize(d, mGemmId); + + assert(d == a + getSerializationSize()); +} + +void LoraPlugin::terminate() noexcept {} + +/////////////// + +LoraPluginCreator::LoraPluginCreator() +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + // Fill PluginFieldCollection with PluginField arguments metadata + mPluginAttributes.clear(); + mPluginAttributes.emplace_back(PluginField("transA", nullptr, PluginFieldType::kINT32, 0)); + mPluginAttributes.emplace_back(PluginField("transB", nullptr, PluginFieldType::kINT32, 0)); + mPluginAttributes.emplace_back(PluginField("lora_module_number", nullptr, PluginFieldType::kINT32, 0)); + mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32, 1)); + mFC.nbFields = mPluginAttributes.size(); + mFC.fields = mPluginAttributes.data(); +} + +const char* LoraPluginCreator::getPluginName() const noexcept +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + return LORA_PLUGIN_NAME; +} + +const char* LoraPluginCreator::getPluginVersion() const noexcept +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + return LORA_PLUGIN_VERSION; +} + +const PluginFieldCollection* LoraPluginCreator::getFieldNames() noexcept +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + return &mFC; +} + +IPluginV2* LoraPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) noexcept +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + + const PluginField* fields = fc->fields; + nvinfer1::DataType type; + int lora_module_number; + int in_hidden_size, out_hidden_size, transA, transB; + bool remove_input_padding; + int max_context_length; + int max_low_rank; + // Read configurations from each fields + for (int i = 0; i < fc->nbFields; ++i) + { + const char* attrName = fields[i].name; + if (!strcmp(attrName, "in_hidden_size")) + { + TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); + in_hidden_size = static_cast(*(static_cast(fields[i].data))); + } + else if (!strcmp(attrName, "out_hidden_size")) + { + TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); + out_hidden_size = static_cast(*(static_cast(fields[i].data))); + } + else if (!strcmp(attrName, "transa")) + { + TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); + transA = static_cast(*(static_cast(fields[i].data))); + } + else if (!strcmp(attrName, "transb")) + { + TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); + transB = static_cast(*(static_cast(fields[i].data))); + } + else if (!strcmp(attrName, "type_id")) + { + TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); + type = static_cast(*(static_cast(fields[i].data))); + } + else if (!strcmp(attrName, "remove_input_padding")) + { + TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); + remove_input_padding = static_cast(*(static_cast(fields[i].data))); + } + else if (!strcmp(attrName, "max_context_length")) + { + TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); + max_context_length = static_cast(*(static_cast(fields[i].data))); + } + else if (!strcmp(attrName, "max_low_rank")) + { + TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); + max_low_rank = static_cast(*(static_cast(fields[i].data))); + } + } + try + { + // LoraPluginCreator is unique and shared for an engine generation + // Create plugin profiler with shared tactics map + // FIXME enable tactic profiler + auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ false, /* skip */ true); + auto* obj = new LoraPlugin(in_hidden_size, out_hidden_size, transA, transB, lora_module_number, type, + pluginProfiler, remove_input_padding, max_context_length, max_low_rank); + obj->setPluginNamespace(mNamespace.c_str()); + return obj; + } + catch (const std::exception& e) + { + caughtError(e); + } + return nullptr; +} + +IPluginV2* LoraPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) noexcept +{ + TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); + // This object will be deleted when the network is destroyed, which will + // call LoraPlugin::destroy() + try + { + // LoraPluginCreator is unique and shared for an engine generation + // Create plugin profiler with shared tactics map + // FIXME enable tactic profiler + auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ true, /* skip */ true); + auto* obj = new LoraPlugin(serialData, serialLength, pluginProfiler); + obj->setPluginNamespace(mNamespace.c_str()); + return obj; + } + catch (const std::exception& e) + { + caughtError(e); + } + return nullptr; +} diff --git a/cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.h b/cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.h new file mode 100644 index 000000000000..9a40ef856e93 --- /dev/null +++ b/cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.h @@ -0,0 +1,161 @@ +/* + * 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. + */ +#ifndef TRT_LORA_PLUGIN_H +#define TRT_LORA_PLUGIN_H +#include "tensorrt_llm/common/cublasMMWrapper.h" +#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" +#include "tensorrt_llm/plugins/common/plugin.h" +#include "tensorrt_llm/plugins/gemmPlugin/gemmPlugin.h" +#include +#include +#include +#include + +namespace tensorrt_llm::plugins +{ + +using CublasGemmWrapper = tensorrt_llm::common::CublasMMWrapper; +using CublasGemmWrapperPtr = std::shared_ptr; + +class LoraPlugin : public BasePlugin +{ +public: + using PluginProfilerPtr = std::shared_ptr; + + LoraPlugin() = delete; + + LoraPlugin(int in_hidden_size, int out_hidden_size, int transA, int transB, int lora_module_number, + nvinfer1::DataType type, const PluginProfilerPtr& profiler, bool remove_input_padding, int max_context_length, + int max_low_rank); + + LoraPlugin(const void* data, size_t length, const PluginProfilerPtr& profiler); + + ~LoraPlugin() override = default; + + // IPluginV2DynamicExt Methods + nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; + nvinfer1::DimsExprs getOutputDimensions(int outputIndex, const nvinfer1::DimsExprs* inputs, int nbInputs, + nvinfer1::IExprBuilder& exprBuilder) noexcept override; + bool supportsFormatCombination( + int pos, const nvinfer1::PluginTensorDesc* inOut, int nbInputs, int nbOutputs) noexcept override; + void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in, int nbInputs, + const nvinfer1::DynamicPluginTensorDesc* out, int nbOutputs) noexcept override; + size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, int nbInputs, + const nvinfer1::PluginTensorDesc* outputs, int nbOutputs) const noexcept override; + int enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc, + const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; + + // IPluginV2Ext Methods + nvinfer1::DataType getOutputDataType( + int index, const nvinfer1::DataType* inputTypes, int nbInputs) const noexcept override; + + // IPluginV2 Methods + const char* getPluginType() const noexcept override; + const char* 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; + +private: + void init(); + void configGemm(); + void setGemmConfig(); + + using IndexType = std::int32_t; + + IndexType getInputTensorIdx() const + { + return 0; + } + + IndexType getHostRequestTypesIdx() const + { + return 1; + } + + IndexType getLoraRanksIdx() const + { + return 2; + } + + IndexType getLoraWeightsPtrsIdx() const + { + return 3; + } + + IndexType getHostContextLengthsIdx() const + { + TLLM_CHECK(mRemoveInputPadding); + return 4; + } + + enum class RequestType : int32_t + { + kCONTEXT = 0, + kGENERATION = 1 + }; + +private: + const std::string mLayerName; + + int mInHiddenSize; + int mOutHiddenSize; + int mTransA; + int mTransB; + nvinfer1::DataType mType; + bool mRemoveInputPadding; + int mMaxContextLength; + int mMaxLowRank; + + // @fixme: seems this is shared across multiple clones. + // If we deep copy the wrapper inside clone(), then we may avoid the mutex inside the wrapper? + CublasGemmWrapperPtr mCublasWrapper; + + GemmDims mDims{}; + GemmIdCublas mGemmId{}; + + PluginProfilerPtr mPluginProfiler; +}; + +class LoraPluginCreator : public BaseCreator +{ +public: + LoraPluginCreator(); + + const char* getPluginName() const noexcept override; + + const char* getPluginVersion() const noexcept override; + + const nvinfer1::PluginFieldCollection* getFieldNames() noexcept override; + + nvinfer1::IPluginV2* createPlugin(const char* name, const nvinfer1::PluginFieldCollection* fc) noexcept override; + + nvinfer1::IPluginV2* deserializePlugin( + const char* name, const void* serialData, size_t serialLength) noexcept override; + +private: + GemmPluginProfilerManager gemmPluginProfileManager; + static nvinfer1::PluginFieldCollection mFC; + static std::vector mPluginAttributes; +}; + +} // namespace tensorrt_llm::plugins + +#endif // TRT_LORA_PLUGIN_H diff --git a/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.cpp b/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.cpp index fcd3315646bc..87835ba2bbd7 100644 --- a/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.cpp +++ b/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.cpp @@ -333,12 +333,13 @@ int WeightOnlyGroupwiseQuantMatmulPlugin::enqueue(const nvinfer1::PluginTensorDe const half* biases_ptr = (mQuantAlgo & BIAS) ? reinterpret_cast(inputs[mBiasesInputIdx]) : nullptr; const half* act_ptr = reinterpret_cast((mQuantAlgo & PRE_QUANT_SCALE) ? workspace : inputs[0]); - TLLM_CHECK_WITH_INFO(mType == nvinfer1::DataType::kHALF #if defined(ENABLE_BF16) - || mType == nvinfer1::DataType::kBF16 -#endif - , + TLLM_CHECK_WITH_INFO(mType == nvinfer1::DataType::kHALF || mType == nvinfer1::DataType::kBF16, "No valid weightOnlyGropwiseQuantMatmul configuration"); +#else + TLLM_CHECK_WITH_INFO(mType == nvinfer1::DataType::kHALF, "No valid weightOnlyGropwiseQuantMatmul configuration"); +#endif + tensorrt_llm::kernels::WeightOnlyActivationType weight_only_act_type; int real_n = n * INT8_INT4_RATIO; if (mType == nvinfer1::DataType::kHALF) diff --git a/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.cpp b/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.cpp index 122868347197..0bc7b59d85aa 100644 --- a/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.cpp +++ b/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.cpp @@ -288,12 +288,12 @@ int WeightOnlyQuantMatmulPlugin::enqueue(const nvinfer1::PluginTensorDesc* input const int ws_size = m_weightOnlyGemmRunner->getWorkspaceSize(m, n, k); const auto& bestTactic = mPluginProfiler->getBestConfig(m, mGemmId); TLLM_CHECK_WITH_INFO(bestTactic, "No valid weight only groupwise GEMM tactic"); - TLLM_CHECK_WITH_INFO(mType == nvinfer1::DataType::kHALF || #if defined(ENABLE_BF16) - mType == nvinfer1::DataType::kBF16 -#endif - , + TLLM_CHECK_WITH_INFO(mType == nvinfer1::DataType::kHALF || mType == nvinfer1::DataType::kBF16, "No valid weightOnlyQuantMatmul configuration"); +#else + TLLM_CHECK_WITH_INFO(mType == nvinfer1::DataType::kHALF, "No valid weightOnlyQuantMatmul configuration"); +#endif tensorrt_llm::kernels::WeightOnlyQuantType weight_only_quant_type; tensorrt_llm::kernels::WeightOnlyActivationType weight_only_act_type; diff --git a/cpp/tensorrt_llm/pybind/bindings.cpp b/cpp/tensorrt_llm/pybind/bindings.cpp index 6485b473275b..df8a2a37ffb8 100644 --- a/cpp/tensorrt_llm/pybind/bindings.cpp +++ b/cpp/tensorrt_llm/pybind/bindings.cpp @@ -15,6 +15,7 @@ * limitations under the License. */ +#include #include #include @@ -72,7 +73,8 @@ PYBIND11_MODULE(TRTLLM_PYBIND_MODULE, m) .def_readwrite("ids", &tpr::GenerationOutput::ids) .def_readwrite("lengths", &tpr::GenerationOutput::lengths) .def_readwrite("log_probs", &tpr::GenerationOutput::logProbs) - .def_readwrite("context_logits", &tpr::GenerationOutput::contextLogits); + .def_readwrite("context_logits", &tpr::GenerationOutput::contextLogits) + .def_readwrite("on_token_generated", &tpr::GenerationOutput::onTokenGenerated); py::class_(m, "KvCacheConfig") .def(py::init, std::optional, std::optional>(), @@ -175,6 +177,9 @@ PYBIND11_MODULE(TRTLLM_PYBIND_MODULE, m) .def_property("compute_context_logits", py::overload_cast<>(&tr::GptModelConfig::computeContextLogits, py::const_), py::overload_cast(&tr::GptModelConfig::computeContextLogits)) + .def_property("compute_generation_logits", + py::overload_cast<>(&tr::GptModelConfig::computeGenerationLogits, py::const_), + py::overload_cast(&tr::GptModelConfig::computeGenerationLogits)) .def_property("model_variant", &tr::GptModelConfig::getModelVariant, &tr::GptModelConfig::setModelVariant) .def_property("use_custom_all_reduce", py::overload_cast<>(&tr::GptModelConfig::useCustomAllReduce, py::const_), py::overload_cast(&tr::GptModelConfig::useCustomAllReduce)); diff --git a/cpp/tensorrt_llm/pybind/runtime/generationOutput.cpp b/cpp/tensorrt_llm/pybind/runtime/generationOutput.cpp index e6d97b983312..0f163ecede79 100644 --- a/cpp/tensorrt_llm/pybind/runtime/generationOutput.cpp +++ b/cpp/tensorrt_llm/pybind/runtime/generationOutput.cpp @@ -16,6 +16,7 @@ */ #include "generationOutput.h" +#include "tensorrt_llm/runtime/torch.h" #include "tensorrt_llm/runtime/torchView.h" namespace tr = tensorrt_llm::runtime; @@ -34,6 +35,12 @@ std::shared_ptr GenerationOutput::toTrtLlm() const { output->contextLogits = tr::TorchView::of(contextLogits.value()); } - // TODO(mseznec): add support for onTokenGenerated + + if (onTokenGenerated) + { + output->onTokenGenerated = [delegate = onTokenGenerated]( + tr::GenerationOutput::TensorPtr const& ids, tr::SizeType step, bool finished) + { delegate(tr::Torch::tensor(ids), step, finished); }; + } return output; } diff --git a/cpp/tensorrt_llm/runtime/bufferManager.cpp b/cpp/tensorrt_llm/runtime/bufferManager.cpp index b0189d28b5e0..810d5daf6bdb 100644 --- a/cpp/tensorrt_llm/runtime/bufferManager.cpp +++ b/cpp/tensorrt_llm/runtime/bufferManager.cpp @@ -114,7 +114,8 @@ void BufferManager::copy(IBuffer const& src, void* dst, MemoryType dstType) cons void BufferManager::copy(IBuffer const& src, IBuffer& dst) const { - TLLM_CHECK_WITH_INFO(src.getDataType() == dst.getDataType(), "Incompatible data types"); + TLLM_CHECK_WITH_INFO(src.getDataType() == dst.getDataType(), + tc::fmtstr("Incompatible data types: %s != %s", src.getDataTypeName(), dst.getDataTypeName())); TLLM_CHECK_WITH_INFO(src.getSizeInBytes() == dst.getSizeInBytes(), tc::fmtstr("Incompatible buffer sizes: %lu != %lu", src.getSizeInBytes(), dst.getSizeInBytes())); copy(src, dst.data(), dst.getMemoryType()); @@ -192,3 +193,49 @@ void BufferManager::initMemoryPool(int device) auto maxThreshold = std::numeric_limits::max(); TLLM_CUDA_CHECK(cudaMemPoolSetAttribute(memPool, cudaMemPoolAttrReleaseThreshold, &maxThreshold)); } + +std::size_t BufferManager::memoryPoolReserved(int device) +{ + ::cudaMemPool_t memPool; + TLLM_CUDA_CHECK(cudaDeviceGetDefaultMemPool(&memPool, device)); + std::size_t reserved = 0; + TLLM_CUDA_CHECK(cudaMemPoolGetAttribute(memPool, cudaMemPoolAttrReservedMemCurrent, &reserved)); + return reserved; +} + +std::size_t BufferManager::memoryPoolUsed(int device) +{ + ::cudaMemPool_t memPool; + TLLM_CUDA_CHECK(cudaDeviceGetDefaultMemPool(&memPool, device)); + std::size_t used = 0; + TLLM_CUDA_CHECK(cudaMemPoolGetAttribute(memPool, cudaMemPoolAttrUsedMemCurrent, &used)); + return used; +} + +void BufferManager::memoryPoolTrimTo(int device, std::size_t size) +{ + ::cudaMemPool_t memPool; + TLLM_CUDA_CHECK(cudaDeviceGetDefaultMemPool(&memPool, device)); + TLLM_CUDA_CHECK(cudaMemPoolTrimTo(memPool, size)); +} + +std::size_t BufferManager::memoryPoolReserved() const +{ + return memoryPoolReserved(mStream->getDevice()); +} + +std::size_t BufferManager::memoryPoolUsed() const +{ + return memoryPoolUsed(mStream->getDevice()); +} + +std::size_t BufferManager::memoryPoolFree() const +{ + return memoryPoolFree(mStream->getDevice()); +} + +void BufferManager::memoryPoolTrimTo(std::size_t size) +{ + mStream->synchronize(); + memoryPoolTrimTo(mStream->getDevice(), size); +} diff --git a/cpp/tensorrt_llm/runtime/gptDecoder.cpp b/cpp/tensorrt_llm/runtime/gptDecoder.cpp index 5d8209c27328..e434b9d7e1d5 100644 --- a/cpp/tensorrt_llm/runtime/gptDecoder.cpp +++ b/cpp/tensorrt_llm/runtime/gptDecoder.cpp @@ -41,10 +41,13 @@ GptDecoder::GptDecoder(size_t vocabSize, size_t vocabSizePadded, CudaStreamPt mDynamicDecodeLayer = std::make_shared>( vocabSize, vocabSizePadded, stream->get(), &mAllocator, isFreeBufferAfterForward, &prop); + + auto constexpr nvFloatType = TRTDataType::value; + mLogProbsTiled = mManager.emptyTensor(MemoryType::kGPU, nvFloatType); } template -void GptDecoder::setup(SamplingConfig const& samplingConfig, size_t batchSize) +void GptDecoder::setup(SamplingConfig const& samplingConfig, size_t batchSize, SizeType maxSequenceLength) { typename layers::DynamicDecodeLayer::SetupParams setupParams; @@ -71,6 +74,10 @@ void GptDecoder::setup(SamplingConfig const& samplingConfig, size_t batchSize setupParams.length_penalty = samplingConfig.lengthPenalty; mDynamicDecodeLayer->setup(batchSize, samplingConfig.beamWidth, setupParams); + + mLogProbsTiled->reshape( + ITensor::makeShape({maxSequenceLength, static_cast(batchSize), samplingConfig.beamWidth})); + mManager.setZero(*mLogProbsTiled); } namespace @@ -128,7 +135,7 @@ typename tl::DynamicDecodeLayer::ForwardParams prepareInputs(DecodingInput co template typename tl::DynamicDecodeLayer::OutputParams prepareOutputs( - DecodingOutput& output, DecodingInput::TensorPtr const& inputLengths) + DecodingOutput& output, DecodingInput::TensorPtr const& inputLengths, DecodingOutput::TensorPtr& logProbsTiled) { TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); typename tl::DynamicDecodeLayer::OutputParams outputParams(tcc::toTllmTensor(*output.ids)); @@ -168,6 +175,7 @@ typename tl::DynamicDecodeLayer::OutputParams prepareOutputs( if (output.logProbs) { outputParams.output_log_probs = tcc::toTllmTensor(*output.logProbs); + outputParams.output_log_probs_tiled = tcc::toTllmTensor(*logProbsTiled); } outputParams.beamHypotheses = std::make_shared(); @@ -218,7 +226,7 @@ bool GptDecoder::forward(DecodingOutput& output, DecodingInput const& input) { TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); auto forwardParams = prepareInputs(input); - auto outputParams = prepareOutputs(output, input.lengths); + auto outputParams = prepareOutputs(output, input.lengths, mLogProbsTiled); BufferManager::ITensorPtr finishedSum; std::int32_t* finishedSumHost = nullptr; @@ -256,19 +264,14 @@ void GptDecoder::forwardAsync(DecodingOutput& output, DecodingInput const& in { TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); auto forwardParams = prepareInputs(input); - auto outputParams = prepareOutputs(output, input.lengths); + auto outputParams = prepareOutputs(output, input.lengths, mLogProbsTiled); mDynamicDecodeLayer->forward(outputParams, forwardParams); } -namespace tensorrt_llm::runtime -{ -template class GptDecoder; -template class GptDecoder; -} // namespace tensorrt_llm::runtime - // this should be similar to gatherTree in cpp/tensorrt_llm/thop/gatherTreeOp.cpp -void IGptDecoder::gatherTree(ITensor& finalOutputIds, DecodingOutput const& decodingOutput, +template +void GptDecoder::gatherTree(ITensor& finalOutputIds, DecodingOutput const& decodingOutput, DecodingInput const& decodingInput, BufferManager const& manager) { TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); @@ -300,7 +303,7 @@ void IGptDecoder::gatherTree(ITensor& finalOutputIds, DecodingOutput const& deco beamHypotheses.sequence_lengths_src = bufferCast(*decodingOutput.lengths); beamHypotheses.parent_ids_src = bufferCast(*decodingOutput.parentIds); beamHypotheses.output_ids_src = bufferCast(*decodingOutput.ids); - beamHypotheses.log_probs_src = nullptr; + beamHypotheses.log_probs_src = bufferCast(*mLogProbsTiled); beamHypotheses.max_seq_len = maxSeqLength; beamHypotheses.length_penalties = nullptr; // TODO (bhsueh) should set length penalties, this should be a gpu tensor When it is set as @@ -316,17 +319,24 @@ void IGptDecoder::gatherTree(ITensor& finalOutputIds, DecodingOutput const& deco beamHypotheses.is_done = bufferCast(*decodingOutput.beamHypotheses.isDone); beamHypotheses.input_lengths = bufferCast(*decodingInput.lengths); + // This is where transpose is done tensorrt_llm::kernels::invokeInsertUnfinishedPath(beamHypotheses, bufferCast(*decodingOutput.finished), bufferCast(*decodingOutput.cumLogProbs), batchSize, beamWidth, stream.get()); sync_check_cuda_error(); tensorrt_llm::kernels::invokeFinalize(bufferCast(finalOutputIds), bufferCast(*decodingOutput.lengths), bufferCast(*decodingOutput.cumLogProbs), - nullptr, // output_logs - beamHypotheses.output_ids_tgt, beamHypotheses.sequence_lengths_tgt, beamHypotheses.normed_scores, - beamHypotheses.cum_log_probs, beamHypotheses.log_probs, beamHypotheses.num_beams, beamHypotheses.input_lengths, - beamWidth, maxSeqLength, batchSize, stream.get()); + decodingOutput.logProbs ? bufferCast(*decodingOutput.logProbs) : nullptr, beamHypotheses.output_ids_tgt, + beamHypotheses.sequence_lengths_tgt, beamHypotheses.normed_scores, beamHypotheses.cum_log_probs, + beamHypotheses.log_probs, beamHypotheses.num_beams, beamHypotheses.input_lengths, beamWidth, maxSeqLength, + batchSize, stream.get()); sync_check_cuda_error(); TLLM_LOG_DEBUG("%s stop", __PRETTY_FUNCTION__); } + +namespace tensorrt_llm::runtime +{ +template class GptDecoder; +template class GptDecoder; +} // namespace tensorrt_llm::runtime diff --git a/cpp/tensorrt_llm/runtime/gptDecoderBatch.cpp b/cpp/tensorrt_llm/runtime/gptDecoderBatch.cpp index 158bf05b217d..ccac65f0f9d5 100644 --- a/cpp/tensorrt_llm/runtime/gptDecoderBatch.cpp +++ b/cpp/tensorrt_llm/runtime/gptDecoderBatch.cpp @@ -100,6 +100,7 @@ GptDecoderBatch::GptDecoderBatch( mFinishedSum = mBufferManager.pinned(ITensor::makeShape({1}), nvSizeType); dOutput->lengths = mBufferManager.emptyTensor(MemoryType::kGPU, nvSizeType); dOutput->cumLogProbs = mBufferManager.emptyTensor(MemoryType::kGPU, nvFloatType); + dOutput->logProbs = mBufferManager.emptyTensor(MemoryType::kGPU, nvFloatType); dOutput->beamHypotheses.empty(mBufferManager); TLLM_LOG_DEBUG("%s stop", __PRETTY_FUNCTION__); } @@ -144,10 +145,14 @@ void GptDecoderBatch::setup(SizeType maxBatchSize, SizeType maxBeamWidth, SizeTy dOutput.finishedSum->reshape(maxBatchSizeShape); mBufferManager.setZero(*dOutput.finishedSum); + dOutput.cumLogProbs->reshape(maxBatchSizeXmaxBeamWidth); + mBufferManager.setZero(*dOutput.cumLogProbs); + + dOutput.logProbs->reshape(ITensor::makeShape({maxBatchSize, maxBeamWidth, mMaxSequenceLength})); + mBufferManager.setZero(*dOutput.logProbs); + if (maxBeamWidth > 1) { - dOutput.cumLogProbs->reshape(maxBatchSizeXmaxBeamWidth); - mBufferManager.setZero(*dOutput.cumLogProbs); dOutput.beamHypotheses.reshape(maxBatchSize, maxBeamWidth, mMaxSequenceLength); } else @@ -262,13 +267,25 @@ void GptDecoderBatch::newRequest( dOutput->newTokens = ITensor::slice(dJointOutput.newTokens, batchIdx, localBatchSize); manager.setZero(*dOutput->newTokens); - if (beamWidth > 1) + // cumLogProb is mandatory for beamWidth > 1 + dOutput->cumLogProbs = nullptr; + if (request.computeCumLogProbs || beamWidth > 1) { dOutput->cumLogProbs = ITensor::slice(dJointOutput.cumLogProbs, batchIdx, localBatchSize); - manager.setZero(*IBuffer::slice(dOutput->cumLogProbs, 0, 1)); + manager.setZero(*dOutput->cumLogProbs); + } + + dOutput->logProbs = nullptr; + if (request.computeLogProbs) + { + dOutput->logProbs = ITensor::slice(dJointOutput.logProbs, batchIdx, localBatchSize); + manager.setZero(*dOutput->logProbs); + } + + if (beamWidth > 1) + { kernels::invokeFill( *IBuffer::slice(dOutput->cumLogProbs, 1, beamWidth - 1), DecodingOutput::kNegativeInfinity, *stream); - dOutput->parentIds = ITensor::slice(dJointOutput.parentIds, batchIdx, localBatchSize); dOutput->parentIds->reshape(outputIdsShape); manager.setZero(*dOutput->parentIds); @@ -277,7 +294,7 @@ void GptDecoderBatch::newRequest( } // remaining - mDecoders[batchIdx]->setup(samplingConfig, localBatchSize); + mDecoders[batchIdx]->setup(samplingConfig, localBatchSize, mMaxSequenceLength); mBeamWidths[batchIdx] = beamWidth; mNbSteps[batchIdx] = 0; mFinished[batchIdx] = false; @@ -407,6 +424,7 @@ CudaEvent GptDecoderBatch::postProcessRequest(SizeType batchIdx) const TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); auto& stream = mStreams[batchIdx]; auto manager = BufferManager{stream}; + auto& decoder = *mDecoders[batchIdx]; auto& dInput = *mDecodingInputs[batchIdx]; auto& dOutput = *mDecodingOutputs[batchIdx]; @@ -414,7 +432,7 @@ CudaEvent GptDecoderBatch::postProcessRequest(SizeType batchIdx) const // TODO can we do this inplace? auto& outputIds = dOutput.ids; auto finalOutputIds = manager.gpu(outputIds->getShape(), outputIds->getDataType()); - IGptDecoder::gatherTree(*finalOutputIds, dOutput, dInput, manager); + decoder.gatherTree(*finalOutputIds, dOutput, dInput, manager); manager.copy(*finalOutputIds, *outputIds); CudaEvent event{}; @@ -424,7 +442,8 @@ CudaEvent GptDecoderBatch::postProcessRequest(SizeType batchIdx) const return event; } -void GptDecoderBatch::newBatch(GenerationInput const& inputs, SamplingConfig const& samplingConfig) +void GptDecoderBatch::newBatch( + GenerationInput const& inputs, GenerationOutput const& outputs, SamplingConfig const& samplingConfig) { TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); // split batch into single requests @@ -458,7 +477,10 @@ void GptDecoderBatch::newBatch(GenerationInput const& inputs, SamplingConfig con inputView = ITensor::slice(inputs.ids, batchIdx, 1); inputView->reshape(inputShape); } - auto request = decoder_batch::Request{inputView, inputs.maxNewTokens, inputs.endId, inputs.padId}; + + auto request = decoder_batch::Request{inputView, inputs.maxNewTokens, inputs.endId}; + request.computeCumLogProbs = (outputs.cumLogProbs != nullptr); + request.computeLogProbs = (outputs.logProbs != nullptr); if (inputs.embeddingBias) { @@ -517,7 +539,7 @@ void GptDecoderBatch::forwardSync() TLLM_LOG_DEBUG("%s stop", __PRETTY_FUNCTION__); } -IStatefulGptDecoder::TensorPtr GptDecoderBatch::getFinalOutputIds() const +void GptDecoderBatch::finalize() const { TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); for (SizeType batchIdx = 0; batchIdx < mActualBatchSize; ++batchIdx) @@ -525,13 +547,12 @@ IStatefulGptDecoder::TensorPtr GptDecoderBatch::getFinalOutputIds() const postProcessRequest(batchIdx); } TLLM_LOG_DEBUG("%s stop", __PRETTY_FUNCTION__); - return getOutputIds(); } -std::tuple GptDecoderBatch::getFinalOutputIds(SizeType batchIdx) const +CudaEvent GptDecoderBatch::finalize(SizeType batchIdx) const { TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); auto event = postProcessRequest(batchIdx); TLLM_LOG_DEBUG("%s stop", __PRETTY_FUNCTION__); - return {std::move(event), getOutputIds(batchIdx)}; + return event; } diff --git a/cpp/tensorrt_llm/runtime/gptJsonConfig.cpp b/cpp/tensorrt_llm/runtime/gptJsonConfig.cpp index 3860e4b611ac..f499464ee2e6 100644 --- a/cpp/tensorrt_llm/runtime/gptJsonConfig.cpp +++ b/cpp/tensorrt_llm/runtime/gptJsonConfig.cpp @@ -107,6 +107,7 @@ GptJsonConfig parseJson(InputType&& i) = parseJsonFieldOr(builderConfig, "max_prompt_embedding_table_size", 0); auto const computeContextLogits = parseJsonFieldOr(builderConfig, "gather_all_token_logits", false); + auto const computeGenerationLogits = parseJsonFieldOr(builderConfig, "gather_all_token_logits", false); auto const& pluginConfig = json.at("plugin_config"); auto const pagedKvCache = pluginConfig.at("paged_kv_cache"); @@ -125,6 +126,7 @@ GptJsonConfig parseJson(InputType&& i) modelConfig.setQuantMode(quantMode); modelConfig.setNbKvHeads(numKvHeads); modelConfig.computeContextLogits(computeContextLogits); + modelConfig.computeGenerationLogits(computeGenerationLogits); modelConfig.setMaxBatchSize(maxBatchSize); modelConfig.setMaxInputLen(maxInputLen); @@ -132,10 +134,10 @@ GptJsonConfig parseJson(InputType&& i) modelConfig.setMaxNumTokens(maxNumTokens); modelConfig.setMaxPromptEmbeddingTableSize(maxPromptEmbeddingTableSize); - if (name == std::string("chatglm-6b")) + if (name == std::string("chatglm_6b") || name == std::string("glm_10b")) { modelConfig.setModelVariant(GptModelConfig::ModelVariant::kGlm); - // kGlm is only for ChatGLM-6B and Glm-10B + // kGlm is only for ChatGLM-6B and GLM-10B } return GptJsonConfig{name, precision, tensorParallelism, pipelineParallelism, modelConfig}; diff --git a/cpp/tensorrt_llm/runtime/gptSession.cpp b/cpp/tensorrt_llm/runtime/gptSession.cpp index e79cbad52ebf..8ab57bb60108 100644 --- a/cpp/tensorrt_llm/runtime/gptSession.cpp +++ b/cpp/tensorrt_llm/runtime/gptSession.cpp @@ -55,7 +55,7 @@ GptSession::GptSession(Config const& sessionConfig, GptModelConfig const& modelC { if (mWorldConfig.isPipelineParallel()) { - mPipelineComm = NcclCommunicator::createPipelineComm(mWorldConfig, *mLogger); + mPipelineComm = NcclCommunicator::createPipelineComm(mWorldConfig); mCommStream = std::make_shared(); } @@ -72,7 +72,7 @@ nvinfer1::ILogger& GptSession::getLogger() const return *mLogger; } -BufferManager& GptSession::getBufferManager() const +BufferManager const& GptSession::getBufferManager() const { return mRuntime->getBufferManager(); } @@ -163,7 +163,8 @@ void GptSession::createKvCacheManager(SizeType batchSize, SizeType beamWidth, Si kvDtype = mModelConfig.getDataType(); } - auto const maxNumTokens = bmkv::KVCacheManager::getMaxNumTokens(config, kvDtype, mModelConfig, mWorldConfig); + auto const maxNumTokens + = bmkv::KVCacheManager::getMaxNumTokens(config, kvDtype, mModelConfig, mWorldConfig, getBufferManager()); TLLM_LOG_INFO("Using %d tokens in paged KV cache.", maxNumTokens); auto const maxNumBlocks = tc::ceilDiv(maxNumTokens, tokensPerBlock); auto const maxBlocksPerSeq = tc::ceilDiv(maxSequenceLength, tokensPerBlock); @@ -302,12 +303,12 @@ void GptSession::kvCacheAddSequences(SizeType beamWidth, SizeType microBatchId, } ITensor::SharedPtr GptSession::initDecoder(ITensor& outputIds, GenerationInput const& inputs, - SamplingConfig const& samplingConfig, SizeType microBatchId) const + GenerationOutput const& outputs, SamplingConfig const& samplingConfig, SizeType microBatchId) const { if (mWorldConfig.isLastPipelineParallelRank()) { auto& decoder = mDecoders.at(microBatchId); - decoder->newBatch(inputs, samplingConfig); + decoder->newBatch(inputs, outputs, samplingConfig); return decoder->getNewTokens(); } else if (mWorldConfig.isFirstPipelineParallelRank()) @@ -444,6 +445,39 @@ std::vector splitInputs(GenerationInput const& inputs, SizeType return inputBatches; } +std::vector splitOutputs(GenerationOutput& outputs, SizeType microBatchSize, BufferManager& manager) +{ + auto const numRequests = outputs.ids->getShape().d[0]; + + std::vector outputBatches; + for (auto batchOffset = 0; batchOffset < numRequests; batchOffset += microBatchSize) + { + auto const batchSize = std::min(microBatchSize, numRequests - batchOffset); + + outputBatches.emplace_back(ITensor::slice(outputs.ids, batchOffset, batchSize), + ITensor::slice(outputs.lengths, batchOffset, batchSize)); + + if (outputs.cumLogProbs) + { + outputBatches.back().cumLogProbs = ITensor::slice(outputs.cumLogProbs, batchOffset, batchSize); + } + if (outputs.logProbs) + { + outputBatches.back().logProbs = ITensor::slice(outputs.logProbs, batchOffset, batchSize); + } + if (outputs.contextLogits) + { + outputBatches.back().contextLogits = ITensor::slice(outputs.contextLogits, batchOffset, batchSize); + } + if (outputs.generationLogits) + { + outputBatches.back().generationLogits = ITensor::slice(outputs.generationLogits, batchOffset, batchSize); + } + } + + return outputBatches; +} + void updateOutputIds(ITensor::SharedPtr const& outputIds, ITensor::SharedPtr const& newTokens, SizeType decoderStep, CudaStream const& stream) { // assemble outputIds of all micro batches @@ -473,32 +507,86 @@ void GptSession::generate( auto const beamWidth = samplingConfig.beamWidth; outputs.ids->reshape(ITensor::makeShape({batchSize, beamWidth, mDecoderMaxSequenceLength})); outputs.lengths->reshape(ITensor::makeShape({batchSize, beamWidth})); - if (mWorldConfig.isLastPipelineParallelRank() && mModelConfig.computeContextLogits()) + if (mWorldConfig.isLastPipelineParallelRank()) { - TLLM_CHECK_WITH_INFO(outputs.contextLogits, - "outputs.contextLogits is nullptr. It must be allocated when computeContextLogits() is enabled."); - auto const vocabSizePadded = mModelConfig.getVocabSizePadded(mWorldConfig.getSize()); - auto const inputLengthsHost = manager.copyFrom(*inputLengths, MemoryType::kCPU); - auto const inputLengthsRange = BufferRange(*inputLengthsHost); - auto const maxInputLength = *std::max_element(inputLengthsRange.begin(), inputLengthsRange.end()); - outputs.contextLogits->reshape(ITensor::makeShape({batchSize, maxInputLength, vocabSizePadded})); + if (outputs.cumLogProbs) + { + TLLM_CHECK_WITH_INFO(outputs.cumLogProbs, + "outputs.cumLogProbs is nullptr. It must be allocated when computeLogProbs is true"); + outputs.cumLogProbs->reshape(ITensor::makeShape({batchSize, beamWidth})); + } + if (outputs.logProbs) + { + TLLM_CHECK_WITH_INFO( + outputs.logProbs, "outputs.logProbs is nullptr. It must be allocated when computeLogProbs is true"); + outputs.logProbs->reshape(ITensor::makeShape({batchSize, beamWidth, mDecoderMaxSequenceLength})); + } + if (mModelConfig.computeContextLogits() || mModelConfig.computeGenerationLogits()) + { + TLLM_CHECK_WITH_INFO(outputs.contextLogits, + "outputs.contextLogits is nullptr. It must be allocated when computeContextLogits() is enabled."); + auto const vocabSizePadded = mModelConfig.getVocabSizePadded(mWorldConfig.getSize()); + auto const inputLengthsHost = manager.copyFrom(*inputLengths, MemoryType::kCPU); + auto const inputLengthsRange = BufferRange(*inputLengthsHost); + auto const maxInputLength = *std::max_element(inputLengthsRange.begin(), inputLengthsRange.end()); + + if (mModelConfig.computeContextLogits()) + { + outputs.contextLogits->reshape(ITensor::makeShape({batchSize, maxInputLength, vocabSizePadded})); + } + + // Initialize the output generation logits buffer + if (mModelConfig.computeGenerationLogits()) + { + SizeType maxNewTokens = 0; + if (inputs.maxNewTokens) + { + maxNewTokens = inputs.maxNewTokens.value(); + } + else + { + for (auto iter = inputLengthsRange.begin(); iter != inputLengthsRange.end(); iter++) + { + maxNewTokens = std::max(maxNewTokens, mDecoderMaxSequenceLength - *iter); + } + } + + TLLM_CHECK_WITH_INFO(maxNewTokens, "maxNewTokens is null"); + + TLLM_CHECK_WITH_INFO(outputs.generationLogits, + "outputs.generationLogits is nullptr. It must be allocated when computeGenerationLogits() is " + "enabled."); + outputs.generationLogits->reshape( + ITensor::makeShape({batchSize, beamWidth, maxNewTokens - 1, vocabSizePadded})); + auto const generationLogitsShape = outputs.generationLogits->getShape(); + TLLM_CHECK_WITH_INFO(generationLogitsShape.d[0] == batchSize, "Invalid dim[0]"); + TLLM_CHECK_WITH_INFO(generationLogitsShape.d[1] == beamWidth, "Invalid dim[1]"); + TLLM_CHECK_WITH_INFO(generationLogitsShape.d[2] == maxNewTokens - 1, "Invalid dim[2]"); + TLLM_CHECK_WITH_INFO(generationLogitsShape.d[3] == vocabSizePadded, "Invalid dim[3]"); + }; + } } + // callbacks + auto const onTokenGenerated = createOnTokenGeneratedCallback(outputs); + if (batchSize <= mMicroBatchConfig.genBatchSize) { - std::vector microBatches{inputs}; - generateBatched(outputs, microBatches, samplingConfig); + std::vector microBatchesInputs{inputs}; + std::vector microBatchesOutputs{outputs}; + generateBatched(microBatchesOutputs, microBatchesInputs, samplingConfig, onTokenGenerated); } else { - auto const microBatches = splitInputs(inputs, mMicroBatchConfig.genBatchSize, manager); - generateBatched(outputs, microBatches, samplingConfig); + auto const microBatchesInputs = splitInputs(inputs, mMicroBatchConfig.genBatchSize, manager); + auto microBatchesOutputs = splitOutputs(outputs, mMicroBatchConfig.genBatchSize, manager); + generateBatched(microBatchesOutputs, microBatchesInputs, samplingConfig, onTokenGenerated); } TLLM_LOG_DEBUG("%s stop", __PRETTY_FUNCTION__); } -std::function GptSession::createOnTokenGeneratedCallback(GenerationOutput& outputs) +GptSession::TokenGeneratedCallback GptSession::createOnTokenGeneratedCallback(GenerationOutput& outputs) { if (outputs.onTokenGenerated && mWorldConfig.isFirstPipelineParallelRank()) { @@ -514,13 +602,15 @@ std::function GptSession::createOnTokenGener } } -void GptSession::generateBatched( - GenerationOutput& outputs, std::vector const& microBatches, SamplingConfig const& samplingConfig) +void GptSession::generateBatched(std::vector& microBatchesOutputs, + std::vector const& microBatchesInputs, SamplingConfig const& samplingConfig, + TokenGeneratedCallback const& onTokenGenerated) { TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); auto& manager = mRuntime->getBufferManager(); - auto const numMicroBatches = static_cast(microBatches.size()); + TLLM_CHECK(microBatchesInputs.size() == microBatchesOutputs.size()); + auto const numMicroBatches = static_cast(microBatchesInputs.size()); TLLM_CHECK(numMicroBatches > 0); TLLM_CHECK(numMicroBatches <= mMicroBatchConfig.numGenBatches); SizeType const beamWidth{samplingConfig.beamWidth}; @@ -528,7 +618,7 @@ void GptSession::generateBatched( // Initialize and reshape buffers for (auto microBatchId = 0; microBatchId < numMicroBatches; ++microBatchId) { - auto const& microBatchInputs = microBatches.at(microBatchId); + auto const& microBatchInputs = microBatchesInputs.at(microBatchId); auto& buffers = *mBuffers.at(microBatchId); buffers.initFromInput(*microBatchInputs.ids, microBatchInputs.lengths, microBatchInputs.packed, beamWidth, mDecoderMaxKvCacheLength, mDecoderMaxSequenceLength, manager); @@ -549,14 +639,29 @@ void GptSession::generateBatched( auto& buffers = *mBuffers.at(microBatchId); auto const batchOffset = microBatchOffsets.at(microBatchId); kvCacheAddSequences(beamWidth, microBatchId, batchOffset); - auto const& microBatchInputs = microBatches.at(microBatchId); - auto const microBatchSize = buffers.generationConfig.batchSize; - buffers.outputIds = ITensor::slice(outputs.ids, batchOffset, microBatchSize); - buffers.outputLengths = ITensor::slice(outputs.lengths, batchOffset, microBatchSize); - buffers.newTokens = initDecoder(*buffers.outputIds, microBatchInputs, samplingConfig, microBatchId); - if (mWorldConfig.isLastPipelineParallelRank() && mModelConfig.computeContextLogits()) + auto const& microBatchInputs = microBatchesInputs.at(microBatchId); + auto& microBatchOutputs = microBatchesOutputs.at(microBatchId); + buffers.outputIds = microBatchOutputs.ids; + buffers.outputLengths = microBatchOutputs.lengths; + buffers.newTokens + = initDecoder(*buffers.outputIds, microBatchInputs, microBatchOutputs, samplingConfig, microBatchId); + + if (mWorldConfig.isLastPipelineParallelRank()) { - buffers.logits = ITensor::slice(outputs.contextLogits, batchOffset, microBatchSize); + buffers.cumLogProbs = nullptr; + if (microBatchOutputs.cumLogProbs) + { + buffers.cumLogProbs = microBatchOutputs.cumLogProbs; + } + buffers.logProbs = nullptr; + if (microBatchOutputs.logProbs) + { + buffers.logProbs = microBatchOutputs.logProbs; + } + if (mModelConfig.computeContextLogits()) + { + buffers.logits = microBatchOutputs.contextLogits; + } } if (mModelConfig.usePromptTuning()) { @@ -564,9 +669,6 @@ void GptSession::generateBatched( } } - // Prepare the onTokenGenerated callback - auto const onTokenGenerated = createOnTokenGeneratedCallback(outputs); - if (useCudaGraphs()) { for (auto& instance : mCudaGraphInstances) @@ -577,7 +679,7 @@ void GptSession::generateBatched( auto kvCacheManager = mModelConfig.usePagedKvCache() ? mKvCacheManager.get() : nullptr; - executeContextStep(microBatches, microBatchOffsets, kvCacheManager); + executeContextStep(microBatchesInputs, microBatchOffsets, kvCacheManager); std::vector microBatchesFinished(numMicroBatches, false); SizeType numBatchesFinished{0}; @@ -585,8 +687,8 @@ void GptSession::generateBatched( while (numBatchesFinished < numMicroBatches) { ++step; - numBatchesFinished - += executeGenerationStep(step, microBatches, microBatchOffsets, kvCacheManager, microBatchesFinished); + numBatchesFinished += executeGenerationStep( + step, microBatchesInputs, microBatchesOutputs, microBatchOffsets, kvCacheManager, microBatchesFinished); onTokenGenerated(step - 1, numBatchesFinished == numMicroBatches); } @@ -608,9 +710,23 @@ void GptSession::generateBatched( // TODO(micro batching) use mCommStream? if (beamWidth > 1) - finalizeOutputIds(microBatchId); + { + finalize(microBatchId); + } else if (!mWorldConfig.isPipelineParallel()) - manager.copy(*mDecoders.at(microBatchId)->getOutputIds(), *mBuffers.at(microBatchId)->outputIds); + { + auto& buffers = *mBuffers.at(microBatchId); + auto& decoder = *mDecoders.at(microBatchId); + manager.copy(*decoder.getOutputIds(), *buffers.outputIds); + + auto& cumLogProbs = buffers.cumLogProbs; + if (cumLogProbs) + manager.copy(*decoder.getCumLogProbs(), *buffers.cumLogProbs); + + auto& logProbs = buffers.logProbs; + if (logProbs) + manager.copy(*decoder.getLogProbs(), *buffers.logProbs); + } } manager.getStream().synchronize(); @@ -668,14 +784,15 @@ void GptSession::executeContextStep(std::vector const& generati TLLM_LOG_DEBUG("%s stop", __PRETTY_FUNCTION__); } -SizeType GptSession::executeGenerationStep(SizeType step, std::vector const& microBatches, - std::vector const& microBatchOffsets, KvCacheManager* kvCacheManager, - std::vector& microBatchesFinished) +SizeType GptSession::executeGenerationStep(SizeType step, std::vector const& microBatchesInputs, + std::vector& microBatchesOutputs, std::vector const& microBatchOffsets, + KvCacheManager* kvCacheManager, std::vector& microBatchesFinished) { TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); + TLLM_CHECK(microBatchesInputs.size() == microBatchesOutputs.size()); auto& manager = mRuntime->getBufferManager(); - auto const numMicroBatches = static_cast(microBatches.size()); + auto const numMicroBatches = static_cast(microBatchesInputs.size()); SizeType numBatchesFinished{0}; auto const flipFlopId = step % 2; @@ -725,6 +842,18 @@ SizeType GptSession::executeGenerationStep(SizeType step, std::vectorsend(*decoder.getNbFinished(), pipelineGroup[peerIdx], *mCommStream, *mLogger); + mPipelineComm->send(*decoder.getNbFinished(), pipelineGroup[peerIdx], *mCommStream); if (beamWidth > 1) { - mPipelineComm->send(cacheIndirection, pipelineGroup[peerIdx], *mCommStream, *mLogger); + mPipelineComm->send(cacheIndirection, pipelineGroup[peerIdx], *mCommStream); } - mPipelineComm->send(sequenceLengths, pipelineGroup[peerIdx], *mCommStream, *mLogger); + mPipelineComm->send(sequenceLengths, pipelineGroup[peerIdx], *mCommStream); } - mPipelineComm->send(*decoder.getNewTokens(), pipelineGroup.front(), *mCommStream, *mLogger); + mPipelineComm->send(*decoder.getNewTokens(), pipelineGroup.front(), *mCommStream); } } else // pipeline parallel mode @@ -781,19 +910,19 @@ void GptSession::decoderStepAsync(SizeType decoderStep, SizeType microBatchId) mCommStream->wait(mCommEvent.get()); auto const pipelineGroup = mWorldConfig.getPipelineParallelGroup(); auto const peer = pipelineGroup.back(); - mPipelineComm->receive(*buffers.nbFinished, peer, *mCommStream, *mLogger); + mPipelineComm->receive(*buffers.nbFinished, peer, *mCommStream); auto& cacheIndirection = *buffers.cacheIndirectionDecoderOutput; auto& sequenceLengths = *buffers.sequenceLengths; auto const beamWidth = cacheIndirection.getShape().d[1]; if (beamWidth > 1) { - mPipelineComm->receive(cacheIndirection, peer, *mCommStream, *mLogger); + mPipelineComm->receive(cacheIndirection, peer, *mCommStream); } - mPipelineComm->receive(sequenceLengths, peer, *mCommStream, *mLogger); + mPipelineComm->receive(sequenceLengths, peer, *mCommStream); if (mWorldConfig.isFirstPipelineParallelRank()) { // receive newTokens from last rank on a separate stream - mPipelineComm->receive(*newTokens, peer, *mCommStream, *mLogger); + mPipelineComm->receive(*newTokens, peer, *mCommStream); updateOutputIds(outputIds, newTokens, decoderStep, *mCommStream); } mCommStream->record(mReceivedEvents.at(microBatchId).get()); @@ -837,12 +966,16 @@ bool GptSession::shouldStopSync(SizeType batchSize, SizeType beamWidth, SizeType return nbFinished == batchSize * beamWidth; } -void GptSession::finalizeOutputIds(SizeType microBatchId) +void GptSession::finalize(SizeType microBatchId) { TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); auto& manager = mRuntime->getBufferManager(); - auto& outputIds = *mBuffers.at(microBatchId)->outputIds; - auto& sequenceLengths = *mBuffers.at(microBatchId)->sequenceLengths; + auto& buffers = *mBuffers.at(microBatchId); + auto& decoder = mDecoders.at(microBatchId); + auto& outputIds = buffers.outputIds; + auto& cumLogProbs = buffers.cumLogProbs; + auto& logProbs = buffers.logProbs; + auto& sequenceLengths = buffers.sequenceLengths; if (mWorldConfig.isPipelineParallel()) { @@ -852,20 +985,56 @@ void GptSession::finalizeOutputIds(SizeType microBatchId) if (mWorldConfig.isLastPipelineParallelRank()) { // send ids from last to first auto const peer = pipelineGroup.front(); - auto const finalOutputIds = mDecoders.at(microBatchId)->getFinalOutputIds(); - mPipelineComm->send(*finalOutputIds, peer, stream, *mLogger); - mPipelineComm->send(sequenceLengths, peer, stream, *mLogger); + decoder->finalize(); + auto finalOutputIds = decoder->getOutputIds(); + + mPipelineComm->send(*finalOutputIds, peer, stream); + mPipelineComm->send(*sequenceLengths, peer, stream); + manager.copy(*finalOutputIds, *outputIds); + + if (cumLogProbs) + { + auto finalCumLogProbs = decoder->getCumLogProbs(); + mPipelineComm->send(*finalCumLogProbs, peer, stream); + manager.copy(*finalCumLogProbs, *cumLogProbs); + } + if (logProbs) + { + auto finalLogProbs = decoder->getLogProbs(); + mPipelineComm->send(*finalLogProbs, peer, stream); + manager.copy(*finalLogProbs, *logProbs); + } } else if (mWorldConfig.isFirstPipelineParallelRank()) { // receive ids from last on first auto const peer = pipelineGroup.back(); - mPipelineComm->receive(outputIds, peer, stream, *mLogger); - mPipelineComm->receive(sequenceLengths, peer, stream, *mLogger); + mPipelineComm->receive(*outputIds, peer, stream); + mPipelineComm->receive(*sequenceLengths, peer, stream); + if (cumLogProbs) + { + mPipelineComm->receive(*cumLogProbs, peer, stream); + } + if (logProbs) + { + mPipelineComm->receive(*logProbs, peer, stream); + } } } else { - manager.copy(*mDecoders.at(microBatchId)->getFinalOutputIds(), outputIds); + decoder->finalize(); + auto finalOutputIds = decoder->getOutputIds(); + manager.copy(*finalOutputIds, *outputIds); + if (cumLogProbs) + { + auto finalCumLogProbs = decoder->getCumLogProbs(); + manager.copy(*finalCumLogProbs, *cumLogProbs); + } + if (logProbs) + { + auto finalLogProbs = decoder->getLogProbs(); + manager.copy(*finalLogProbs, *logProbs); + } // sequenceLengths are already updated by decoder } diff --git a/cpp/tensorrt_llm/runtime/iBuffer.cpp b/cpp/tensorrt_llm/runtime/iBuffer.cpp index 89783e3317d9..543c5bde6ebf 100644 --- a/cpp/tensorrt_llm/runtime/iBuffer.cpp +++ b/cpp/tensorrt_llm/runtime/iBuffer.cpp @@ -17,6 +17,7 @@ #include "tensorrt_llm/runtime/iBuffer.h" #include "tensorrt_llm/runtime/iTensor.h" +#include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/runtime/bufferView.h" @@ -79,3 +80,31 @@ std::ostream& tensorrt_llm::runtime::operator<<(std::ostream& output, IBuffer co ITensor::makeShape({static_cast(buffer.getSize())}), buffer.getCapacity()); return output << *tensor; } + +char const* IBuffer::getDataTypeName() const +{ + switch (getDataType()) + { + case nvinfer1::DataType::kINT64: return DataTypeTraits::name; + case nvinfer1::DataType::kINT32: return DataTypeTraits::name; + case nvinfer1::DataType::kFLOAT: return DataTypeTraits::name; + case nvinfer1::DataType::kBF16: return DataTypeTraits::name; + case nvinfer1::DataType::kHALF: return DataTypeTraits::name; + case nvinfer1::DataType::kBOOL: return DataTypeTraits::name; + case nvinfer1::DataType::kUINT8: return DataTypeTraits::name; + case nvinfer1::DataType::kINT8: return DataTypeTraits::name; + case nvinfer1::DataType::kFP8: return DataTypeTraits::name; + } + TLLM_THROW("Unknown data type"); +} + +char const* IBuffer::getMemoryTypeName() const +{ + switch (getMemoryType()) + { + case MemoryType::kPINNED: return MemoryTypeString::value; + case MemoryType::kCPU: return MemoryTypeString::value; + case MemoryType::kGPU: return MemoryTypeString::value; + } + TLLM_THROW("Unknown memory type"); +} diff --git a/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp b/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp index c7bd2760b8ca..f18bdb07366c 100644 --- a/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp +++ b/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp @@ -64,39 +64,38 @@ struct NcclDataType } // namespace template -void NcclCommunicator::send( - T* sendbuff, size_t count, int peer, CudaStream const& stream, nvinfer1::ILogger& logger) const +void NcclCommunicator::send(T* sendbuff, size_t count, int peer, CudaStream const& stream) const { #if ENABLE_MULTI_DEVICE auto datatype = NcclDataType>::value; - TLLM_NCCL_CHECK(ncclSend(sendbuff, count, datatype, peer, mComm, stream.get()), logger); + TLLM_NCCL_CHECK(ncclSend(sendbuff, count, datatype, peer, mComm, stream.get())); #else TLLM_THROW("Multi device support is disabled."); #endif // ENABLE_MULTI_DEVICE } -template void NcclCommunicator::send(std::uint8_t*, size_t, int, CudaStream const&, nvinfer1::ILogger&) const; -template void NcclCommunicator::send(std::int32_t*, size_t, int, CudaStream const&, nvinfer1::ILogger&) const; -template void NcclCommunicator::send(std::uint8_t const*, size_t, int, CudaStream const&, nvinfer1::ILogger&) const; -template void NcclCommunicator::send(std::int32_t const*, size_t, int, CudaStream const&, nvinfer1::ILogger&) const; +template void NcclCommunicator::send(std::uint8_t*, size_t, int, CudaStream const&) const; +template void NcclCommunicator::send(std::int32_t*, size_t, int, CudaStream const&) const; +template void NcclCommunicator::send(std::uint8_t const*, size_t, int, CudaStream const&) const; +template void NcclCommunicator::send(std::int32_t const*, size_t, int, CudaStream const&) const; +template void NcclCommunicator::send(float const*, size_t, int, CudaStream const&) const; template -void NcclCommunicator::receive( - T* sendbuff, size_t count, int peer, CudaStream const& stream, nvinfer1::ILogger& logger) const +void NcclCommunicator::receive(T* sendbuff, size_t count, int peer, CudaStream const& stream) const { #if ENABLE_MULTI_DEVICE auto datatype = NcclDataType>::value; - TLLM_NCCL_CHECK(ncclRecv(sendbuff, count, datatype, peer, mComm, stream.get()), logger); + TLLM_NCCL_CHECK(ncclRecv(sendbuff, count, datatype, peer, mComm, stream.get())); #else TLLM_THROW("Multi device support is disabled."); #endif // ENABLE_MULTI_DEVICE } -template void NcclCommunicator::receive(std::uint8_t*, size_t, int, CudaStream const&, nvinfer1::ILogger&) const; -template void NcclCommunicator::receive(std::int32_t*, size_t, int, CudaStream const&, nvinfer1::ILogger&) const; +template void NcclCommunicator::receive(std::uint8_t*, size_t, int, CudaStream const&) const; +template void NcclCommunicator::receive(std::int32_t*, size_t, int, CudaStream const&) const; +template void NcclCommunicator::receive(float*, size_t, int, CudaStream const&) const; -std::shared_ptr NcclCommunicator::createPipelineComm( - WorldConfig const& worldConfig, nvinfer1::ILogger& logger) +std::shared_ptr NcclCommunicator::createPipelineComm(WorldConfig const& worldConfig) { #if ENABLE_MULTI_DEVICE int const myRank = worldConfig.getRank(); @@ -108,18 +107,18 @@ std::shared_ptr NcclCommunicator::createPipelineComm( ncclGetUniqueId(&id); for (auto peer = 1; peer < worldSize; ++peer) { - TLLM_MPI_CHECK(MPI_Send(&id, sizeof(id), MPI_BYTE, peer, 0, MPI_COMM_WORLD), logger); + TLLM_MPI_CHECK(MPI_Send(&id, sizeof(id), MPI_BYTE, peer, 0, MPI_COMM_WORLD)); } } else { auto constexpr peer = 0; MPI_Status status; - TLLM_MPI_CHECK(MPI_Recv(&id, sizeof(id), MPI_BYTE, peer, 0, MPI_COMM_WORLD, &status), logger); + TLLM_MPI_CHECK(MPI_Recv(&id, sizeof(id), MPI_BYTE, peer, 0, MPI_COMM_WORLD, &status)); } auto pipelineComm = std::make_shared(); - TLLM_NCCL_CHECK(ncclCommInitRank(&pipelineComm->mComm, worldSize, id, myRank), logger); + TLLM_NCCL_CHECK(ncclCommInitRank(&pipelineComm->mComm, worldSize, id, myRank)); return pipelineComm; #else diff --git a/cpp/tensorrt_llm/runtime/ncclCommunicator.h b/cpp/tensorrt_llm/runtime/ncclCommunicator.h index 44e0d67b7b85..f8066941b1bc 100644 --- a/cpp/tensorrt_llm/runtime/ncclCommunicator.h +++ b/cpp/tensorrt_llm/runtime/ncclCommunicator.h @@ -30,25 +30,24 @@ class NcclCommunicator { public: template - void send(T* sendbuff, size_t count, int peer, CudaStream const& stream, nvinfer1::ILogger& logger) const; + void send(T* sendbuff, size_t count, int peer, CudaStream const& stream) const; template - void send(IBuffer const& buf, int peer, CudaStream const& stream, nvinfer1::ILogger& logger) const + void send(IBuffer const& buf, int peer, CudaStream const& stream) const { - send(bufferCast(buf), buf.getSize(), peer, stream, logger); + send(bufferCast(buf), buf.getSize(), peer, stream); } template - void receive(T* sendbuff, size_t count, int peer, CudaStream const& stream, nvinfer1::ILogger& logger) const; + void receive(T* sendbuff, size_t count, int peer, CudaStream const& stream) const; template - void receive(IBuffer& buf, int peer, CudaStream const& stream, nvinfer1::ILogger& logger) const + void receive(IBuffer& buf, int peer, CudaStream const& stream) const { - receive(bufferCast(buf), buf.getSize(), peer, stream, logger); + receive(bufferCast(buf), buf.getSize(), peer, stream); } - static std::shared_ptr createPipelineComm( - WorldConfig const& worldConfig, nvinfer1::ILogger& logger); + static std::shared_ptr createPipelineComm(WorldConfig const& worldConfig); private: ncclComm_t mComm; diff --git a/cpp/tensorrt_llm/runtime/runtimeBuffers.cpp b/cpp/tensorrt_llm/runtime/runtimeBuffers.cpp index 4aed18c02526..b167c5416275 100644 --- a/cpp/tensorrt_llm/runtime/runtimeBuffers.cpp +++ b/cpp/tensorrt_llm/runtime/runtimeBuffers.cpp @@ -83,6 +83,9 @@ void RuntimeBuffers::clear() cacheIndirectionDecoderInput = nullptr; cacheIndirectionDecoderOutput = nullptr; + cumLogProbs = nullptr; + logProbs = nullptr; + hiddenStates = nullptr; allocated = false; @@ -155,10 +158,8 @@ void RuntimeBuffers::create(TllmRuntime& runtime, GptModelConfig const& modelCon if (modelConfig.useGptAttentionPlugin()) { pastKeyValueLengths = manager.emptyTensor(MemoryType::kCPU, nvinfer1::DataType::kINT32); - for (SizeType i = 0; i < modelConfig.getNbLayers(); ++i) - { - maxKvCacheLengths.emplace_back(manager.emptyTensor(MemoryType::kCPU, nvinfer1::DataType::kINT32)); - } + maxKvCacheLengths + = utils::createBufferVector(runtime, localNbLayers, MemoryType::kCPU, nvinfer1::DataType::kINT32); } else { @@ -238,11 +239,8 @@ void RuntimeBuffers::reshape(GptModelConfig const& modelConfig, WorldConfig cons if (modelConfig.useGptAttentionPlugin()) { pastKeyValueLengths->reshape(ITensor::makeShape({batchSize})); - for (SizeType i = 0; i < modelConfig.getNbLayers(); ++i) - { - maxKvCacheLengths[i]->reshape(ITensor::makeShape({1})); - } requestTypes->reshape(ITensor::makeShape({batchSize})); + utils::reshapeBufferVector(maxKvCacheLengths, ITensor::makeShape({1})); } else { @@ -511,6 +509,21 @@ void RuntimeBuffers::postContextStep(std::vector const& contextB TLLM_LOG_DEBUG("%s stop", __PRETTY_FUNCTION__); } +void RuntimeBuffers::postEachGenerationStep(BufferManager& manager, TensorPtr outputGenerationLogits, SizeType step, + SizeType firstBatchSlotIdx, SizeType microBatchSize, SizeType beamWidth, WorldConfig const& worldConfig) +{ + TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); + + if (worldConfig.isLastPipelineParallelRank()) + { + kernels::copyLatestTokenLogitsInGeneration( + *outputGenerationLogits, *logits, step, firstBatchSlotIdx, microBatchSize, beamWidth, manager.getStream()); + manager.getStream().synchronize(); + } + + TLLM_LOG_DEBUG("%s stop", __PRETTY_FUNCTION__); +} + void RuntimeBuffers::prepareContextStep(TensorPtr const& inputIds, TokenIdType const padId, BufferManager& manager, KvCacheManager const* kvCacheManager, SizeType firstBatchSlotIdx, GptModelConfig const& modelConfig, WorldConfig const& worldConfig) @@ -523,6 +536,9 @@ void RuntimeBuffers::prepareContextStep(TensorPtr const& inputIds, TokenIdType c // use context lengths only in context step sequenceLengths = contextLengthsDevice; + // get local number of layers. + auto const localNbLayers = modelConfig.getNbLayers(worldConfig.getPipelineParallelism()); + if (modelConfig.useGptAttentionPlugin()) { auto pastKeyValueLengthsPtr = bufferCast(*pastKeyValueLengths); @@ -534,7 +550,7 @@ void RuntimeBuffers::prepareContextStep(TensorPtr const& inputIds, TokenIdType c std::fill_n(RequestTypesPtr, batchSize, 0); // Set maxKvCacheLengths buffer to the same value currently. - for (auto layer = 0; layer < modelConfig.getNbLayers(); ++layer) + for (auto layer = 0; layer < localNbLayers; ++layer) { bufferCast(*maxKvCacheLengths[layer])[0] = generationConfig.maxKvCacheLength; } @@ -803,12 +819,7 @@ void RuntimeBuffers::getRuntimeBuffers(TensorMap& inputBuffers, TensorMap& outpu inputBuffers.insert_or_assign("host_past_key_value_lengths", pastKeyValueLengths); inputBuffers.insert_or_assign("host_request_types", requestTypes); inputBuffers.insert_or_assign("sequence_length", sequenceLengths); - - for (SizeType i = 0; i < modelConfig.getNbLayers(); ++i) - { - std::string name = "host_max_kv_cache_length_" + std::to_string(i); - inputBuffers.insert_or_assign(name, maxKvCacheLengths[i]); - } + utils::insertTensorVector(inputBuffers, "host_max_kv_cache_length_", maxKvCacheLengths, firstLayerId); if (modelConfig.usePackedInput()) { diff --git a/cpp/tensorrt_llm/runtime/runtimeBuffers.h b/cpp/tensorrt_llm/runtime/runtimeBuffers.h index efa669328dbf..5a98c6ee8705 100644 --- a/cpp/tensorrt_llm/runtime/runtimeBuffers.h +++ b/cpp/tensorrt_llm/runtime/runtimeBuffers.h @@ -106,6 +106,10 @@ class RuntimeBuffers // decoder TensorPtr nbFinished; + // Log probs + TensorPtr cumLogProbs; + TensorPtr logProbs; + // pipeline parallelism TensorPtr hiddenStates; @@ -135,6 +139,9 @@ class RuntimeBuffers void postContextStep(std::vector const& contextBuffers, BufferManager& manager, GptModelConfig const& modelConfig, WorldConfig const& worldConfig); + void postEachGenerationStep(BufferManager& manager, TensorPtr outputGenerationLogits, SizeType step, + SizeType firstBatchSlotIdx, SizeType microBatchSize, SizeType beamWidth, WorldConfig const& worldConfig); + void prepareContextStep(TensorPtr const& inputIds, TokenIdType padId, BufferManager& manager, KvCacheManager const* kvCacheManager, SizeType firstBatchSlotIdx, GptModelConfig const& modelConfig, WorldConfig const& worldConfig); diff --git a/cpp/tensorrt_llm/runtime/runtimeKernels.cu b/cpp/tensorrt_llm/runtime/runtimeKernels.cu index 45da61f20fbf..1dc955df6b5a 100644 --- a/cpp/tensorrt_llm/runtime/runtimeKernels.cu +++ b/cpp/tensorrt_llm/runtime/runtimeKernels.cu @@ -1013,4 +1013,87 @@ void gatherLastTokenLogits(ITensor& output, ITensor const& input, ITensor const& } } +// In the following kernel, we launch a grid with microBatchSize * beamWidth blocks of threads. Each thread block +// copies a `vocabSizePadded` length logits tensor from the "inputLogits (microBatchSize, beamWidth, vocabSizePadded)" +// to the "outputGenerationLogits (batchSize, beamWidth, outPutLen, vocabSizePadded)" +template +__global__ void copyLatestTokenLogitsInGenerationKernel(T* outputGenerationLogits, T const* inputLogits, int step, + int firstBatchSlotIdx, int beamWidth, int outPutLen, int vocabSizePadded) +{ + // The relatively batch slot index that this thread block in microBatchSize. + int relativeBatchSlotIdx = blockIdx.x / beamWidth; + + // The Absolute batch slot index in batchSize. + int absoluteBatchSlotIdx = firstBatchSlotIdx + relativeBatchSlotIdx; + + // The beam index that this thread block process + int mbeamIdx = blockIdx.x % beamWidth; + + // The output pointer. + const unsigned int outputOffset + = (absoluteBatchSlotIdx * beamWidth * outPutLen + mbeamIdx * outPutLen + step) * vocabSizePadded; + T* outputPtr = &outputGenerationLogits[outputOffset]; + + // The input pointer. + const unsigned int inputOffset = (relativeBatchSlotIdx * beamWidth + mbeamIdx) * vocabSizePadded; + T const* inputPtr = &inputLogits[inputOffset]; + + // The threads in the block collaborate to copy the logits. + for (int idx = threadIdx.x; idx < vocabSizePadded; idx += blockDim.x) + { + outputPtr[idx] = inputPtr[idx]; + } +} + +template +void invokeCopyLatestTokenLogitsInGeneration(ITensor& output, ITensor const& input, SizeType step, + SizeType firstBatchSlotIdx, SizeType microBatchSize, SizeType beamWidth, CudaStream const& stream) +{ + auto const& outputShape = output.getShape(); + auto const maxBatchSize = static_cast(outputShape.d[0]); + auto const _beamWidth = static_cast(outputShape.d[1]); + auto const outPutLen = static_cast(outputShape.d[2]); + auto const vocabSizePadded = static_cast(outputShape.d[3]); + + TLLM_CHECK_WITH_INFO(maxBatchSize >= microBatchSize, "Invalid output shape: dim[0]"); + TLLM_CHECK_WITH_INFO(_beamWidth == beamWidth, "Invalid output shape: dim[1]"); + TLLM_CHECK_WITH_INFO(outPutLen >= step, "Invalid output shape: dim[2]"); + + auto const& inputShape = input.getShape(); + TLLM_CHECK_WITH_INFO(inputShape.d[0] == microBatchSize, "Invalid input shape: dim[0]"); + TLLM_CHECK_WITH_INFO(inputShape.d[1] == beamWidth, "Invalid input shape: dim[1]"); + TLLM_CHECK_WITH_INFO(inputShape.d[2] == vocabSizePadded, "Invalid input shape: dim[2]"); + + dim3 const blockSize{256, 1}; + dim3 const gridSize{static_cast(microBatchSize * beamWidth), 1}; + + copyLatestTokenLogitsInGenerationKernel<<>>( + bufferCast(output), bufferCast(input), step, firstBatchSlotIdx, beamWidth, outPutLen, vocabSizePadded); +} + +void copyLatestTokenLogitsInGeneration(ITensor& output, ITensor const& input, SizeType step, SizeType firstBatchSlotIdx, + SizeType microBatchSize, SizeType beamWidth, CudaStream const& stream) +{ + switch (input.getDataType()) + { + case nvinfer1::DataType::kFLOAT: + invokeCopyLatestTokenLogitsInGeneration( + output, input, step, firstBatchSlotIdx, microBatchSize, beamWidth, stream); + break; + case nvinfer1::DataType::kHALF: + invokeCopyLatestTokenLogitsInGeneration( + output, input, step, firstBatchSlotIdx, microBatchSize, beamWidth, stream); + break; + case nvinfer1::DataType::kBF16: + invokeCopyLatestTokenLogitsInGeneration<__nv_bfloat16>( + output, input, step, firstBatchSlotIdx, microBatchSize, beamWidth, stream); + break; + case nvinfer1::DataType::kFP8: + invokeCopyLatestTokenLogitsInGeneration<__nv_fp8_e4m3>( + output, input, step, firstBatchSlotIdx, microBatchSize, beamWidth, stream); + break; + default: TLLM_CHECK_WITH_INFO(false, "data type not supported"); + } +} + } // namespace tensorrt_llm::runtime::kernels diff --git a/cpp/tensorrt_llm/runtime/runtimeKernels.h b/cpp/tensorrt_llm/runtime/runtimeKernels.h index 8b08d68ca018..6d2e3999f395 100644 --- a/cpp/tensorrt_llm/runtime/runtimeKernels.h +++ b/cpp/tensorrt_llm/runtime/runtimeKernels.h @@ -81,4 +81,7 @@ void tileTensorInplace(ITensor& tensor, SizeType beamWidth, CudaStream const& st void gatherLastTokenLogits( ITensor& output, ITensor const& input, ITensor const& lastTokenIds, CudaStream const& stream); +void copyLatestTokenLogitsInGeneration(ITensor& output, ITensor const& input, SizeType step, SizeType firstBatchSlotIdx, + SizeType microBatchSize, SizeType beamWidth, CudaStream const& stream); + } // namespace tensorrt_llm::runtime::kernels diff --git a/cpp/tensorrt_llm/runtime/statefulGptDecoder.cpp b/cpp/tensorrt_llm/runtime/statefulGptDecoder.cpp index 60b00963d17c..4cd496d7616c 100644 --- a/cpp/tensorrt_llm/runtime/statefulGptDecoder.cpp +++ b/cpp/tensorrt_llm/runtime/statefulGptDecoder.cpp @@ -120,7 +120,8 @@ void StatefulGptDecoder::reshapeBuffers( TLLM_LOG_DEBUG("%s stop", __PRETTY_FUNCTION__); } -void StatefulGptDecoder::newBatch(GenerationInput const& inputs, SamplingConfig const& samplingConfig) +void StatefulGptDecoder::newBatch( + GenerationInput const& inputs, GenerationOutput const& outputs, SamplingConfig const& samplingConfig) { TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); auto& manager = mBufferManager; @@ -132,7 +133,7 @@ void StatefulGptDecoder::newBatch(GenerationInput const& inputs, SamplingConfig auto const beamWidth = samplingConfig.beamWidth; reshapeBuffers(batchSize, beamWidth, mMaxKvCacheLength, mMaxSequenceLength); - mDecoder->setup(samplingConfig, batchSize); + mDecoder->setup(samplingConfig, batchSize, mMaxSequenceLength); // sanity checks, should always be true after reshape auto const& outputIdsShape = mDecodingOutput->ids->getShape(); @@ -189,6 +190,19 @@ void StatefulGptDecoder::newBatch(GenerationInput const& inputs, SamplingConfig manager.setZero(*dOutput.finished); manager.setZero(*dOutput.finishedSum); + // If outputs contains cumLogProbs, use that + if (outputs.cumLogProbs) + { + dOutput.cumLogProbs = outputs.cumLogProbs; + } + dOutput.logProbs = outputs.logProbs; + + if (dOutput.cumLogProbs) + manager.setZero(*dOutput.cumLogProbs); + + if (dOutput.logProbs) + manager.setZero(*dOutput.logProbs); + if (beamWidth > 1) { std::vector cumLogProbsHost(batchSize * beamWidth, DecodingOutput::kNegativeInfinity); @@ -199,13 +213,6 @@ void StatefulGptDecoder::newBatch(GenerationInput const& inputs, SamplingConfig } manager.copy(cumLogProbsHost.data(), *dOutput.cumLogProbs); - // kernels::invokeFill(*dOutput.cumLogProbs, DecodingOutput::kNegativeInfinity, *stream); - // for (SizeType batchIdx = 0; batchIdx < batchSize; ++batchIdx) - // { - // auto cumLogProbsSlice = ITensor::slice(dOutput.cumLogProbs, batchIdx, 1); - // manager.setZero(*IBuffer::slice(cumLogProbsSlice, 0, 1)); - // } - manager.setZero(*dOutput.parentIds); dOutput.beamHypotheses.init(manager, endId); } @@ -268,14 +275,14 @@ void StatefulGptDecoder::forwardSync() TLLM_LOG_DEBUG("%s stop", __PRETTY_FUNCTION__); } -IStatefulGptDecoder::TensorPtr StatefulGptDecoder::getFinalOutputIds() const +void StatefulGptDecoder::finalize() const { // TODO (rkobus) can we do this inplace? TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); auto& outputIds = mDecodingOutput->ids; auto finalOutputIds = mBufferManager.gpu(outputIds->getShape(), outputIds->getDataType()); - IGptDecoder::gatherTree(*finalOutputIds, *mDecodingOutput, *mDecodingInput, mBufferManager); + mDecoder->gatherTree(*finalOutputIds, *mDecodingOutput, *mDecodingInput, mBufferManager); mBufferManager.copy(*finalOutputIds, *outputIds); TLLM_LOG_DEBUG("%s stop", __PRETTY_FUNCTION__); - return outputIds; + return; } diff --git a/cpp/tensorrt_llm/runtime/statefulGptDecoder.h b/cpp/tensorrt_llm/runtime/statefulGptDecoder.h index 5244dceb3335..e60907e2e5bf 100644 --- a/cpp/tensorrt_llm/runtime/statefulGptDecoder.h +++ b/cpp/tensorrt_llm/runtime/statefulGptDecoder.h @@ -43,14 +43,15 @@ class StatefulGptDecoder : public IStatefulGptDecoder nvinfer1::DataType dtype) override; //! @brief Initialize the decoder with new batch of inputs. - void newBatch(GenerationInput const& input, SamplingConfig const& samplingConfig) override; + void newBatch( + GenerationInput const& input, GenerationOutput const& output, SamplingConfig const& samplingConfig) override; void forwardAsync(decoder::Output& output, decoder::Input const& input) override; void forwardSync() override; //! @brief Gather final results for all requests. - [[nodiscard]] TensorPtr getFinalOutputIds() const override; + void finalize() const override; //! @returns [batchSize, maxBeamWidth, maxInputLength + maxNewTokens], contains input token ids and generated token //! ids without padding, on gpu @@ -59,6 +60,18 @@ class StatefulGptDecoder : public IStatefulGptDecoder return mDecodingOutput->ids; } + //! @returns [batchSize, maxBeamWidth], cumulative log probabilities (per beam), on gpu + [[nodiscard]] TensorPtr getCumLogProbs() const override + { + return mDecodingOutput->cumLogProbs; + } + + //! @returns [batchSize, maxBeamWidth], cumulative log probabilities (per beam), on gpu + [[nodiscard]] TensorPtr getLogProbs() const override + { + return mDecodingOutput->logProbs; + } + //! @returns [batchSize, maxBeamWidth], tokens generated in last forward pass, on gpu [[nodiscard]] TensorPtr getNewTokens() const override { diff --git a/cpp/tensorrt_llm/runtime/utils/multiDeviceUtils.h b/cpp/tensorrt_llm/runtime/utils/multiDeviceUtils.h index 53e37ffef12d..39aaa779fc00 100644 --- a/cpp/tensorrt_llm/runtime/utils/multiDeviceUtils.h +++ b/cpp/tensorrt_llm/runtime/utils/multiDeviceUtils.h @@ -16,6 +16,7 @@ #pragma once +#include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/stringUtils.h" #include @@ -24,30 +25,21 @@ #include #endif // ENABLE_MULTI_DEVICE -#define TLLM_MPI_CHECK(cmd, logger) \ +#define TLLM_MPI_CHECK(cmd) \ do \ { \ auto e = cmd; \ - if (e != MPI_SUCCESS) \ - { \ - (logger).log(nvinfer1::ILogger::Severity::kERROR, \ - tensorrt_llm::common::fmtstr("Failed: MPI error %s:%d '%d'", __FILE__, __LINE__, e).c_str()); \ - exit(EXIT_FAILURE); \ - } \ + TLLM_CHECK_WITH_INFO(e == MPI_SUCCESS, \ + tensorrt_llm::common::fmtstr("Failed: MPI error %s:%d '%d'", __FILE__, __LINE__, e).c_str()); \ } while (0) #if ENABLE_MULTI_DEVICE -#define TLLM_NCCL_CHECK(cmd, logger) \ +#define TLLM_NCCL_CHECK(cmd) \ do \ { \ ncclResult_t r = cmd; \ - if (r != ncclSuccess) \ - { \ - (logger).log(nvinfer1::ILogger::Severity::kERROR, \ - tensorrt_llm::common::fmtstr( \ - "Failed, NCCL error %s:%d '%s'\n", __FILE__, __LINE__, ncclGetErrorString(r)) \ - .c_str()); \ - exit(EXIT_FAILURE); \ - } \ + TLLM_CHECK_WITH_INFO(r == ncclSuccess, \ + tensorrt_llm::common::fmtstr("Failed, NCCL error %s:%d '%s'\n", __FILE__, __LINE__, ncclGetErrorString(r)) \ + .c_str()); \ } while (0) #endif // ENABLE_MULTI_DEVICE diff --git a/cpp/tensorrt_llm/runtime/worldConfig.cpp b/cpp/tensorrt_llm/runtime/worldConfig.cpp index 15af55aece21..58ffe49800de 100644 --- a/cpp/tensorrt_llm/runtime/worldConfig.cpp +++ b/cpp/tensorrt_llm/runtime/worldConfig.cpp @@ -21,6 +21,7 @@ #include "tensorrt_llm/runtime/tllmLogger.h" #include "tensorrt_llm/runtime/utils/multiDeviceUtils.h" +#include #include #include @@ -40,15 +41,18 @@ void initMpi(nvinfer1::ILogger& logger, int threadMode = MPI_THREAD_FUNNELED) } int initialized = 0; - TLLM_MPI_CHECK(MPI_Initialized(&initialized), logger); + TLLM_MPI_CHECK(MPI_Initialized(&initialized)); if (!initialized) { logger.log( nvinfer1::ILogger::Severity::kINFO, tc::fmtstr("Initializing MPI with thread mode %d", threadMode).c_str()); int providedMode; - TLLM_MPI_CHECK(MPI_Init_thread(nullptr, nullptr, threadMode, &providedMode), logger); + TLLM_MPI_CHECK(MPI_Init_thread(nullptr, nullptr, threadMode, &providedMode)); TLLM_CHECK_WITH_INFO(providedMode >= threadMode, "MPI_Init_thread failed"); std::atexit([]() { MPI_Finalize(); }); + + auto previousHandler = std::signal(SIGABRT, [](int signal) { MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE); }); + TLLM_CHECK_WITH_INFO(previousHandler != SIG_ERR, "Signal handler setup failed"); } mpiInitialized = true; @@ -61,7 +65,7 @@ bool WorldConfig::validConfig(nvinfer1::ILogger& logger, SizeType tensorParallel initMpi(logger); int mpiSize; - TLLM_MPI_CHECK(MPI_Comm_size(MPI_COMM_WORLD, &mpiSize), logger); + TLLM_MPI_CHECK(MPI_Comm_size(MPI_COMM_WORLD, &mpiSize)); return mpiSize == tensorParallelism * pipelineParallelism; } @@ -71,8 +75,8 @@ WorldConfig WorldConfig::mpi(nvinfer1::ILogger& logger, SizeType gpusPerNode, st initMpi(logger); int mpiSize, mpiRank; - TLLM_MPI_CHECK(MPI_Comm_size(MPI_COMM_WORLD, &mpiSize), logger); - TLLM_MPI_CHECK(MPI_Comm_rank(MPI_COMM_WORLD, &mpiRank), logger); + TLLM_MPI_CHECK(MPI_Comm_size(MPI_COMM_WORLD, &mpiSize)); + TLLM_MPI_CHECK(MPI_Comm_rank(MPI_COMM_WORLD, &mpiRank)); logger.log(nvinfer1::ILogger::Severity::kINFO, tc::fmtstr("MPI size: %d, rank: %d", mpiSize, mpiRank).c_str()); auto pp = pipelineParallelism.value_or(1); diff --git a/cpp/tensorrt_llm/thop/ncclCommunicatorOp.cpp b/cpp/tensorrt_llm/thop/ncclCommunicatorOp.cpp index 3b549faabc8f..9188eb2230d2 100644 --- a/cpp/tensorrt_llm/thop/ncclCommunicatorOp.cpp +++ b/cpp/tensorrt_llm/thop/ncclCommunicatorOp.cpp @@ -15,18 +15,16 @@ */ #include "tensorrt_llm/thop/ncclCommunicatorOp.h" -#include "tensorrt_llm/runtime/tllmLogger.h" namespace torch_ext { NcclCommunicatorOp::NcclCommunicatorOp(int64_t tpSize, int64_t ppSize, int64_t rank) - : mLogger(std::make_shared()) - , mRank(static_cast(rank)) + : mRank(static_cast(rank)) { tensorrt_llm::runtime::WorldConfig worldConfig{ static_cast(tpSize), static_cast(ppSize), static_cast(rank)}; - mPipelineComm = tensorrt_llm::runtime::NcclCommunicator::createPipelineComm(worldConfig, *mLogger); + mPipelineComm = tensorrt_llm::runtime::NcclCommunicator::createPipelineComm(worldConfig); } void NcclCommunicatorOp::send(th::Tensor tensor, int64_t toRank) const @@ -34,7 +32,7 @@ void NcclCommunicatorOp::send(th::Tensor tensor, int64_t toRank) const auto ptr = reinterpret_cast(get_ptr(tensor)); size_t const size = tensor.numel() * th::elementSize(th::typeMetaToScalarType(tensor.dtype())); tensorrt_llm::runtime::CudaStream cudaStream{at::cuda::getCurrentCUDAStream().stream(), mRank, false}; - mPipelineComm->send(ptr, size, static_cast(toRank), cudaStream, *mLogger); + mPipelineComm->send(ptr, size, static_cast(toRank), cudaStream); } void NcclCommunicatorOp::recv(th::Tensor& tensor, int64_t fromRank) const @@ -42,7 +40,7 @@ void NcclCommunicatorOp::recv(th::Tensor& tensor, int64_t fromRank) const auto ptr = reinterpret_cast(get_ptr(tensor)); size_t const size = tensor.numel() * th::elementSize(th::typeMetaToScalarType(tensor.dtype())); tensorrt_llm::runtime::CudaStream cudaStream{at::cuda::getCurrentCUDAStream().stream(), mRank, false}; - mPipelineComm->receive(ptr, size, static_cast(fromRank), cudaStream, *mLogger); + mPipelineComm->receive(ptr, size, static_cast(fromRank), cudaStream); } } // namespace torch_ext diff --git a/cpp/tensorrt_llm/thop/ncclCommunicatorOp.h b/cpp/tensorrt_llm/thop/ncclCommunicatorOp.h index 9f08f1a845ab..cff669313c8f 100755 --- a/cpp/tensorrt_llm/thop/ncclCommunicatorOp.h +++ b/cpp/tensorrt_llm/thop/ncclCommunicatorOp.h @@ -33,7 +33,6 @@ class NcclCommunicatorOp : public th::jit::CustomClassHolder void recv(th::Tensor& tensor, int64_t fromRank) const; private: - std::shared_ptr mLogger; int32_t mRank; std::shared_ptr mPipelineComm; }; diff --git a/cpp/tests/kernels/weightOnly/weightOnlyKernelTest.cpp b/cpp/tests/kernels/weightOnly/weightOnlyKernelTest.cpp index 28bcb8c04404..6b58e593ab18 100644 --- a/cpp/tests/kernels/weightOnly/weightOnlyKernelTest.cpp +++ b/cpp/tests/kernels/weightOnly/weightOnlyKernelTest.cpp @@ -9,7 +9,6 @@ #include "tensorrt_llm/kernels/weightOnlyBatchedGemv/kernelLauncher.h" #include -#include #include #include #include @@ -64,11 +63,19 @@ struct BType struct CutlassKernel; struct CudaKernel; +void simple_assert(bool flag) +{ + if (!flag) + { + throw std::runtime_error("assert failed"); + } +} + template float benchmark_perchannel(void* act, void* weight, void* scales, void* zeros, void* bias, void* out, int m, int n, int k, int group_size, int warmup, int iter) { - assert(zeros == nullptr && bias == nullptr && group_size == 0); + simple_assert(zeros == nullptr && bias == nullptr && group_size == 0); cudaStream_t s; cudaStreamCreate(&s); cudaEvent_t begin, end; @@ -115,7 +122,6 @@ float benchmark_perchannel(void* act, void* weight, void* scales, void* zeros, v cudaEventSynchronize(end); float time; cudaEventElapsedTime(&time, begin, end); - fast_time = std::min(fast_time, time); if (time < fast_time) { fast_time = time; @@ -127,7 +133,6 @@ float benchmark_perchannel(void* act, void* weight, void* scales, void* zeros, v { gemm.gemm(act, weight, scales, out, m, n, k, best_config, ws_ptr, ws_bytes, s); } - cudaProfilerStart(); cudaEventRecord(begin, s); for (int i = 0; i < iter; ++i) { @@ -151,7 +156,7 @@ template ( - d_act.data(), d_weight.data(), d_scales.data(), p_zeros, p_bias, d_out.data(), m, n, k, 0, warmup, iter); + std::function)> benchmark_func_cuda + = benchmark_perchannel; + std::function)> benchmark_func_cutlass + = benchmark_perchannel; + if (group_size != 0) + { + benchmark_func_cuda = benchmark_groupwise; + benchmark_func_cutlass = benchmark_groupwise; + } + time1 = benchmark_func_cuda(d_act.data(), d_weight.data(), d_scales.data(), p_zeros, p_bias, d_out.data(), m, n, k, + group_size, warmup, iter); d_out.copy_to(h_out1.data()); - time2 = benchmark_perchannel( - d_act.data(), d_weight.data(), d_scales.data(), p_zeros, p_bias, d_out.data(), m, n, k, 0, warmup, iter); + time2 = benchmark_func_cutlass(d_act.data(), d_weight.data(), d_scales.data(), p_zeros, p_bias, d_out.data(), m, n, + k, group_size, warmup, iter); d_out.copy_to(h_out2.data()); float quant_scale = 1.f / (1 << (8 / elem_per_byte - 1)); bool pass = compare(h_out1.data(), h_out2.data(), m * n, quant_scale); diff --git a/cpp/tests/resources/scripts/build_chatglm_engines.py b/cpp/tests/resources/scripts/build_chatglm_engines.py index f3d50dcdf741..319f89ffc9c7 100755 --- a/cpp/tests/resources/scripts/build_chatglm_engines.py +++ b/cpp/tests/resources/scripts/build_chatglm_engines.py @@ -20,7 +20,6 @@ import subprocess as _sp import sys import typing as _tp -from collections import OrderedDict as _OrderedDict from pathlib import Path as _Path import torch.multiprocessing as _mp @@ -35,11 +34,11 @@ import build as _ecb -def build_engine(model_version: str, weight_dir: _pl.Path, engine_dir: _pl.Path, +def build_engine(model_name: str, weight_dir: _pl.Path, engine_dir: _pl.Path, world_size, *args): args = [ '-m', - str(model_version), + str(model_name), '--log_level=error', '--model_dir', str(weight_dir), @@ -47,6 +46,8 @@ def build_engine(model_version: str, weight_dir: _pl.Path, engine_dir: _pl.Path, str(engine_dir), '--max_batch_size=2', '--max_beam_width=2', + "--max_input_len=512", + "--max_output_len=512", '--builder_opt=0', f'--world_size={world_size}', ] + list(args) @@ -64,14 +65,8 @@ def run_command(command: _tp.Sequence[str], *, cwd=None, **kwargs) -> None: def build_engines(model_cache: _tp.Optional[str] = None, world_size: int = 1): - model_name_dict = _OrderedDict([ - ["chatglm-6b", "1"], - ["chatglm2-6b", "2"], - ["chatglm3-6b", "3"], - ]) - hf_dir_list = [ - resources_dir / model_name for model_name in model_name_dict.keys() - ] + model_name_list = ["chatglm_6b", "chatglm2_6b", "chatglm3_6b"] + hf_dir_list = [resources_dir / model_name for model_name in model_name_list] trt_dir = resources_dir / "trtModel" run_command( @@ -80,21 +75,23 @@ def build_engines(model_cache: _tp.Optional[str] = None, world_size: int = 1): cwd=resources_dir) # Clone the model directory - for model_name, hf_dir in zip(model_name_dict.keys(), hf_dir_list): + for model_name, hf_dir in zip(model_name_list, hf_dir_list): if not _Path(hf_dir).exists(): run_command( [ "git", "clone", - "https://huggingface.co/THUDM/" + model_name, + "https://huggingface.co/THUDM/" + + model_name.replace("_", "-"), + model_name, ], cwd=resources_dir, ) print("\nBuilding engines") - for model, hf_dir in zip(model_name_dict.items(), hf_dir_list): - print("Building %s" % model[0]) - build_engine(model[1], hf_dir, trt_dir, world_size) + for model_name, hf_dir in zip(model_name_list, hf_dir_list): + print("Building %s" % model_name) + build_engine(model_name, hf_dir, trt_dir, world_size) if not _Path(engine_target_path).exists(): _Path(engine_target_path).mkdir(parents=True, exist_ok=True) diff --git a/cpp/tests/resources/scripts/generate_expected_chatglm_output.py b/cpp/tests/resources/scripts/generate_expected_chatglm_output.py index 44c5920effb8..90e89f5b5f88 100755 --- a/cpp/tests/resources/scripts/generate_expected_chatglm_output.py +++ b/cpp/tests/resources/scripts/generate_expected_chatglm_output.py @@ -16,7 +16,6 @@ import json import sys -from collections import OrderedDict from pathlib import Path import numpy as np @@ -39,21 +38,14 @@ def generate(model_name, batch_size, beam_width): - model_name_dict = OrderedDict([ - ["chatglm-6b", "1"], - ["chatglm2-6b", "2"], - ["chatglm3-6b", "3"], - ]) - print("generate expected %s output BatchSize=%d, BeamWidth=%d" % (model_name, batch_size, beam_width)) - args = parse_arguments() + args = parse_arguments(['-m', model_name]) if batch_size == 1: args.input_text = args.input_text[:1] elif batch_size > 2: args.input_text += args.input_text[0] * (batch_size - 2) - args.model_version = model_name_dict[model_name] args.beam_width = beam_width args.tokenizer_dir = resources_dir / model_name args.engine_dir = Path(__file__).parent.parent / "models/rt_engine/chatglm" @@ -65,17 +57,22 @@ def generate(model_name, batch_size, beam_width): config = json.load(f) assert (config['builder_config']['name'] == model_name) dtype = config['builder_config']['precision'] - end_id = config['builder_config']['eos_token_id'] - pad_id = config['builder_config']['pad_token_id'] + config['builder_config']['max_batch_size'] + max_input_len = config['builder_config']['max_input_len'] + max_output_len = config['builder_config']['max_output_len'] + config['builder_config']['max_beam_width'] + remove_input_padding = config['builder_config']['remove_input_padding'] use_gpt_attention_plugin = config['plugin_config']['gpt_attention_plugin'] world_size = config['builder_config']['tensor_parallel'] assert world_size == tensorrt_llm.mpi_world_size( ), f'Engine world size ({world_size}) != Runtime world size ({tensorrt_llm.mpi_world_size()})' runtime_rank = tensorrt_llm.mpi_rank() - runtime_mapping = tensorrt_llm.Mapping(world_size, - runtime_rank, - tp_size=world_size) + runtime_mapping = tensorrt_llm.Mapping( + world_size, + runtime_rank, + tp_size=world_size, + ) torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node) serialize_path = find_engines( @@ -88,15 +85,51 @@ def generate(model_name, batch_size, beam_width): tokenizer = transformers.AutoTokenizer.from_pretrained( args.tokenizer_dir, trust_remote_code=True) + end_id = tokenizer.eos_token_id + pad_id = tokenizer.pad_token_id + if args.model_name in ["glm_10b"]: + sop_id = tokenizer.sop_token_id + eop_id = tokenizer.eop_token_id input_text = args.input_text tokenized = tokenizer(input_text, return_tensors="pt", padding=True, return_length=True) - input_ids = tokenized['input_ids'].int().contiguous().cuda() - input_lengths = tokenized['length'].int().contiguous().cuda() + input_ids = tokenized['input_ids'].int() + input_lengths = tokenized['length'].int() + max_input_len_real = torch.max(input_lengths) + if max_input_len_real > max_input_len: + print("Truncate input_length as %d" % max_input_len) + input_ids = input_ids[:, :max_input_len] + input_lengths = torch.where(input_lengths > max_input_len, + max_input_len, input_lengths) + else: + max_input_len = max_input_len_real + if args.model_name in ["glm_10b"]: + input_ids = torch.cat( + (input_ids, input_ids.new_full((batch_size, 1), sop_id)), + dim=-1, + ) + input_lengths += 1 + max_input_len_real += 1 + + if remove_input_padding: + input_ids_no_padding = torch.zeros(1, + torch.sum(input_lengths), + dtype=torch.int32) + lengths_acc = torch.cumsum( + torch.cat([torch.IntTensor([0]), input_lengths]), + dim=0, + ) + for i in range(len(input_ids)): + input_ids_no_padding[ + 0, lengths_acc[i]:lengths_acc[i + 1]] = torch.IntTensor( + input_ids[i, + max_input_len - input_lengths[i]:max_input_len]) - if use_gpt_attention_plugin: + input_ids = input_ids_no_padding + + elif use_gpt_attention_plugin: # when using gpt attention plugin, inputs needs to align at the head input_ids_padding_right = torch.zeros_like(input_ids) + end_id for i, sample in enumerate(input_ids): @@ -125,7 +158,7 @@ def generate(model_name, batch_size, beam_width): ) sampling_config = SamplingConfig( - end_id=end_id, + end_id=eop_id if args.model_name in ["glm_10b"] else end_id, pad_id=pad_id, num_beams=args.beam_width, temperature=args.temperature, @@ -136,23 +169,35 @@ def generate(model_name, batch_size, beam_width): with open(serialize_path, 'rb') as f: engine_buffer = f.read() - if model_name == 'chatglm-6b': - decoder = ChatGLMGenerationSession( - model_config, - engine_buffer, - runtime_mapping, - ) + + if args.model_name in ["chatglm_6b", "glm_10b"]: + session = ChatGLMGenerationSession else: - decoder = GenerationSession( - model_config, - engine_buffer, - runtime_mapping, - ) - decoder.setup(input_ids.size(0), input_ids.size(1), args.max_output_len, - args.beam_width) - output_ids = decoder.decode(input_ids, input_lengths, sampling_config) + session = GenerationSession + decoder = session( + model_config, + engine_buffer, + runtime_mapping, + ) + + decoder.setup( + len(input_text), + max_input_len, + max_output_len, + beam_width, + ) + output = decoder.decode( + input_ids.contiguous().cuda(), + input_lengths.contiguous().cuda(), + sampling_config, + output_sequence_lengths=True, + return_dict=True, + ) torch.cuda.synchronize() + output_ids = output["output_ids"] + output["sequence_lengths"] + data_path = Path(__file__).parent.parent / "data" / model_name data_path.mkdir(parents=True, exist_ok=True) nBS, nBM = input_ids.size(0), args.beam_width @@ -174,12 +219,13 @@ def generate(model_name, batch_size, beam_width): if __name__ == '__main__': - generate("chatglm-6b", batch_size=1, beam_width=1) - generate("chatglm-6b", batch_size=2, beam_width=1) - generate("chatglm2-6b", batch_size=1, beam_width=1) - generate("chatglm2-6b", batch_size=2, beam_width=1) - generate("chatglm2-6b", batch_size=1, beam_width=2) - generate("chatglm3-6b", batch_size=1, beam_width=1) - generate("chatglm3-6b", batch_size=2, beam_width=1) - generate("chatglm3-6b", batch_size=1, beam_width=2) + generate("chatglm_6b", batch_size=1, beam_width=1) + generate("chatglm2_6b", batch_size=1, beam_width=1) + generate("chatglm2_6b", batch_size=2, beam_width=1) + generate("chatglm2_6b", batch_size=1, beam_width=2) + generate("chatglm3_6b", batch_size=1, beam_width=1) + generate("chatglm3_6b", batch_size=2, beam_width=1) + generate("chatglm3_6b", batch_size=1, beam_width=2) + #generate("glm_10b", batch_size=1, beam_width=1) + #generate("glm_10b", batch_size=2, beam_width=1) print("Done.") diff --git a/cpp/tests/resources/scripts/test_cpp.py b/cpp/tests/resources/scripts/test_cpp.py index 7e553ad21d77..f35386b7b7d0 100755 --- a/cpp/tests/resources/scripts/test_cpp.py +++ b/cpp/tests/resources/scripts/test_cpp.py @@ -86,12 +86,14 @@ def run_tests(cuda_architectures: _tp.Optional[str] = None, build_dir: _tp.Optional[str] = None, dist_dir: _tp.Optional[str] = None, model_cache: _tp.Optional[str] = None, + skip_gpt=False, skip_gptj=False, skip_llama=False, skip_chatglm=False, only_fp8=False, only_multi_gpu=False, - trt_root: _tp.Optional[str] = None) -> None: + trt_root: _tp.Optional[str] = None, + build_only=False) -> None: root_dir = find_root_dir() _log.info("Using root directory: %s", str(root_dir)) @@ -114,27 +116,39 @@ def run_tests(cuda_architectures: _tp.Optional[str] = None, root_dir=root_dir, resources_dir=resources_dir, model_cache=model_cache, + skip_gpt=skip_gpt, skip_gptj=skip_gptj, skip_llama=skip_llama, skip_chatglm=skip_chatglm, only_fp8=only_fp8) + if build_only: + return + run_google_tests(build_dir=build_dir, + skip_gpt=skip_gpt, skip_gptj=skip_gptj, skip_llama=skip_llama, skip_chatglm=skip_chatglm, only_fp8=only_fp8) - run_benchmarks(python_exe=python_exe, - root_dir=root_dir, - build_dir=build_dir, - resources_dir=resources_dir) + if not skip_gpt: + run_benchmarks(python_exe=python_exe, + root_dir=root_dir, + build_dir=build_dir, + resources_dir=resources_dir) + else: + _log.info("Skipping benchmarks") + else: prepare_multi_gpu_model_tests(python_exe=python_exe, root_dir=root_dir, resources_dir=resources_dir, model_cache=model_cache) + if build_only: + return + run_multi_gpu_tests(build_dir=build_dir) @@ -142,6 +156,7 @@ def prepare_all_model_tests(python_exe: str, root_dir: _pl.Path, resources_dir: _pl.Path, model_cache: _tp.Optional[str] = None, + skip_gpt=False, skip_gptj=False, skip_llama=False, skip_chatglm=False, @@ -149,11 +164,14 @@ def prepare_all_model_tests(python_exe: str, model_cache_arg = ["--model_cache", model_cache] if model_cache else [] only_fp8_arg = ["--only_fp8"] if only_fp8 else [] - prepare_model_tests(model_name="gpt", - python_exe=python_exe, - root_dir=root_dir, - resources_dir=resources_dir, - model_cache_arg=model_cache_arg) + if not skip_gpt: + prepare_model_tests(model_name="gpt", + python_exe=python_exe, + root_dir=root_dir, + resources_dir=resources_dir, + model_cache_arg=model_cache_arg) + else: + _log.info("Skipping GPT tests") if not skip_gptj: prepare_model_tests(model_name="gptj", @@ -228,8 +246,8 @@ def prepare_model_tests(model_name: str, run_command(generate_expected_output, cwd=root_dir, env=model_env) -def run_google_tests(build_dir: _pl.Path, skip_gptj, skip_llama, skip_chatglm, - only_fp8): +def run_google_tests(build_dir: _pl.Path, skip_gpt, skip_gptj, skip_llama, + skip_chatglm, only_fp8): make_google_tests = [ "cmake", "--build", ".", "--config", "Release", "-j", "--target", "google-tests" @@ -239,6 +257,10 @@ def run_google_tests(build_dir: _pl.Path, skip_gptj, skip_llama, skip_chatglm, cpp_env = {**_os.environ} ctest = ["ctest", "--output-on-failure", "--output-junit", "results.xml"] excluded_tests = [] + if skip_gpt: + excluded_tests.append( + ".*GptTest.*|.*GptSessionTest.*|.*GptManagerTest.*|.*TrtGptModelTest.*" + ) if skip_gptj: excluded_tests.append(".*Gptj.*") if skip_llama: @@ -343,6 +365,21 @@ def run_benchmarks(python_exe: str, root_dir: _pl.Path, build_dir: _pl.Path, parser.add_argument("--model_cache", type=str, help="Directory where models are stored") + parser.add_argument("--only_gpt", + action="store_true", + help="Run only the tests for GPT") + parser.add_argument("--only_gptj", + action="store_true", + help="Run only the tests for GPT-J") + parser.add_argument("--only_llama", + action="store_true", + help="Run only the tests for Llama") + parser.add_argument("--only_chatglm", + action="store_true", + help="Run only the tests for ChatGLM") + parser.add_argument("--skip_gpt", + action="store_true", + help="Skip the tests for GPT") parser.add_argument("--skip_gptj", action="store_true", help="Skip the tests for GPT-J") @@ -360,5 +397,39 @@ def run_benchmarks(python_exe: str, root_dir: _pl.Path, build_dir: _pl.Path, "--only_multi_gpu", action="store_true", help="Run only mulit-GPU tests. Implemented for 4 GPUs.") + parser.add_argument("--build_only", + action="store_true", + help="Build only, do not run tests.") + + args = parser.parse_args() + + if (args.only_gpt + args.only_gptj + args.only_llama + args.only_chatglm > + 1): + parser.error('Cannot combine multiple only_* arguments.') + + if args.only_gpt: + args.skip_gptj = True + args.skip_llama = True + args.skip_chatglm = True + + if args.only_gptj: + args.skip_gpt = True + args.skip_llama = True + args.skip_chatglm = True + + if args.only_llama: + args.skip_gpt = True + args.skip_gptj = True + args.skip_chatglm = True + + if args.only_chatglm: + args.skip_gpt = True + args.skip_gptj = True + args.skip_llama = True + + del args.only_gpt + del args.only_gptj + del args.only_llama + del args.only_chatglm - run_tests(**vars(parser.parse_args())) + run_tests(**vars(args)) diff --git a/cpp/tests/runtime/bufferManagerTest.cpp b/cpp/tests/runtime/bufferManagerTest.cpp index a6a7c7a11f63..628fda884991 100644 --- a/cpp/tests/runtime/bufferManagerTest.cpp +++ b/cpp/tests/runtime/bufferManagerTest.cpp @@ -115,7 +115,7 @@ TEST_F(BufferManagerTest, Pointers) static_assert(static_cast(trtPointerType) == BufferDataType::kTrtPointerType); static_assert(trtPointerType == BufferDataType::kTrtPointerType); // uses implicit type conversion // The C++ type corresponding to the TensorRT type for storing pointers (int64_t) - using cppStorageType = CppDataType::type; + using cppStorageType = DataTypeTraits::type; static_assert(sizeof(cppStorageType) == sizeof(cppPointerType)); BufferManager manager(mStream); @@ -152,4 +152,21 @@ TEST_F(BufferManagerTest, MemPoolAttributes) std::uint64_t threshold{0}; TLLM_CUDA_CHECK(cudaMemPoolGetAttribute(memPool, cudaMemPoolAttrReleaseThreshold, &threshold)); EXPECT_EQ(threshold, std::numeric_limits::max()); + + manager.memoryPoolTrimTo(0); + auto const reserved = manager.memoryPoolReserved(); + auto const used = manager.memoryPoolUsed(); + auto const free = manager.memoryPoolFree(); + EXPECT_EQ(free, reserved - used); + auto constexpr kBytesToReserve = 1 << 20; + { + auto const mem = manager.allocate(MemoryType::kGPU, kBytesToReserve); + EXPECT_EQ(mem->getSize(), kBytesToReserve); + EXPECT_GE(manager.memoryPoolReserved(), reserved + kBytesToReserve); + EXPECT_GE(manager.memoryPoolUsed(), used + kBytesToReserve); + } + EXPECT_GE(manager.memoryPoolFree(), free + kBytesToReserve); + manager.memoryPoolTrimTo(0); + EXPECT_LE(manager.memoryPoolReserved(), reserved); + EXPECT_LE(manager.memoryPoolFree(), free); } diff --git a/cpp/tests/runtime/gptDecoderBatchTest.cpp b/cpp/tests/runtime/gptDecoderBatchTest.cpp index 34b0c14aea8a..3d9456b897c0 100644 --- a/cpp/tests/runtime/gptDecoderBatchTest.cpp +++ b/cpp/tests/runtime/gptDecoderBatchTest.cpp @@ -90,7 +90,8 @@ void verifyResults(BufferManager& manager, GptDecoderBatch const& decoder, } } -void testDecoder(nvinfer1::DataType const dtype, std::vector const& samplingConfigs, int maxBeamWidth) +void testDecoder(nvinfer1::DataType const dtype, std::vector const& samplingConfigs, int maxBeamWidth, + bool computeLogProbs) { TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); SizeType constexpr tensorParallelism{1}; @@ -172,7 +173,11 @@ void testDecoder(nvinfer1::DataType const dtype, std::vector con auto input = std::shared_ptr(manager.gpu(shape, TRTDataType::value)); kernels::invokeFill(*input, tokenId, *streamPtr); inputIds.emplace_back(input); - decoder.newRequest(b, decoder_batch::Request{inputIds[b], maxNewTokens, endId, padId}, samplingConfigs[b]); + + auto decoderRequest = decoder_batch::Request{inputIds[b], maxNewTokens, endId}; + decoderRequest.computeCumLogProbs = computeLogProbs; + decoderRequest.computeLogProbs = computeLogProbs; + decoder.newRequest(b, decoderRequest, samplingConfigs[b]); } cudaDeviceSynchronize(); @@ -206,13 +211,16 @@ void testDecoder(nvinfer1::DataType const dtype, std::vector con EXPECT_NO_THROW(decoder.forward(outputs, inputs)); EXPECT_THAT(decoder.getNbSteps(), ::testing::Each(maxNewTokens)); - decoder.newRequest(0, decoder_batch::Request{inputIds[0], maxNewTokens}, samplingConfigs[0]); + auto decoderRequest = decoder_batch::Request{inputIds[0], maxNewTokens}; + decoderRequest.computeCumLogProbs = computeLogProbs; + decoderRequest.computeLogProbs = computeLogProbs; + decoder.newRequest(0, decoderRequest, samplingConfigs[0]); EXPECT_FALSE(decoder.getFinished()[0]); EXPECT_EQ(decoder.getNbSteps()[0], 0); } -void testDecoderWavefront( - nvinfer1::DataType const dtype, std::vector const& samplingConfigs, int maxBeamWidth) +void testDecoderWavefront(nvinfer1::DataType const dtype, std::vector const& samplingConfigs, + int maxBeamWidth, bool computeLogProbs) { TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); SizeType constexpr tensorParallelism{1}; @@ -302,7 +310,11 @@ void testDecoderWavefront( auto input = std::shared_ptr(manager.gpu(shape, TRTDataType::value)); kernels::invokeFill(*input, tokenId, *streamPtr); inputIds.emplace_back(input); - decoder.newRequest(b, decoder_batch::Request{inputIds[b], maxNewTokens, endId, padId}, samplingConfigs[b]); + + auto decoderRequest = decoder_batch::Request{inputIds[b], maxNewTokens, endId}; + decoderRequest.computeCumLogProbs = computeLogProbs; + decoderRequest.computeLogProbs = computeLogProbs; + decoder.newRequest(b, decoderRequest, samplingConfigs[b]); decoder.forward(outputs, inputs); @@ -335,7 +347,7 @@ struct BeamConfig std::vector beamWidths; }; -class ParamTest : public ::testing::TestWithParam> +class ParamTest : public ::testing::TestWithParam> { }; @@ -343,32 +355,39 @@ TEST_P(ParamTest, Test) { nvinfer1::DataType const dtype{std::get<0>(GetParam())}; BeamConfig const beamConfig{std::get<1>(GetParam())}; + bool const computeLogProbs{std::get<2>(GetParam())}; std::vector samplingConfigs; for (auto const beamWidth : beamConfig.beamWidths) { samplingConfigs.emplace_back(beamWidth); } - testDecoder(dtype, samplingConfigs, beamConfig.maxBeamWidth); + testDecoder(dtype, samplingConfigs, beamConfig.maxBeamWidth, computeLogProbs); } -INSTANTIATE_TEST_SUITE_P(GptDecoderTest, ParamTest, +INSTANTIATE_TEST_SUITE_P(GptDecoderBatchTest, ParamTest, testing::Combine(testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kHALF), testing::Values(BeamConfig{1, {1, 1, 1}}, BeamConfig{3, {3, 3, 3, 3}}, BeamConfig{4, {1, 1}}, - BeamConfig{4, {3, 3, 3}}, BeamConfig{4, {1, 2, 3, 4}})), + BeamConfig{4, {3, 3, 3}}, BeamConfig{4, {1, 2, 3, 4}}), + testing::Values(false, true)), [](const testing::TestParamInfo& info) { std::string name{std::get<0>(info.param) == nvinfer1::DataType::kFLOAT ? "Float" : "Half"}; BeamConfig const beamConfig = std::get<1>(info.param); + bool const computeLogProbs = std::get<2>(info.param); name.append("MaxBeamWidth" + std::to_string(beamConfig.maxBeamWidth)); - for (auto const beamWdith : beamConfig.beamWidths) + for (auto const beamWidth : beamConfig.beamWidths) { - name.append("Bw" + std::to_string(beamWdith)); + name.append("Bw" + std::to_string(beamWidth)); + } + if (computeLogProbs) + { + name.append("LogProbs"); } return name; }); -class ParamWavefrontTest : public ::testing::TestWithParam> +class ParamWavefrontTest : public ::testing::TestWithParam> { }; @@ -376,27 +395,34 @@ TEST_P(ParamWavefrontTest, Test) { nvinfer1::DataType const dtype{std::get<0>(GetParam())}; BeamConfig const beamConfig{std::get<1>(GetParam())}; + bool const computeLogProbs{std::get<2>(GetParam())}; std::vector samplingConfigs; for (auto const beamWidth : beamConfig.beamWidths) { samplingConfigs.emplace_back(beamWidth); } - testDecoderWavefront(dtype, samplingConfigs, beamConfig.maxBeamWidth); + testDecoderWavefront(dtype, samplingConfigs, beamConfig.maxBeamWidth, computeLogProbs); } -INSTANTIATE_TEST_SUITE_P(GptDecoderTest, ParamWavefrontTest, +INSTANTIATE_TEST_SUITE_P(GptDecoderBatchTest, ParamWavefrontTest, testing::Combine(testing::Values(nvinfer1::DataType::kFLOAT, nvinfer1::DataType::kHALF), testing::Values(BeamConfig{1, {1, 1, 1}}, BeamConfig{3, {3, 3, 3, 3}}, BeamConfig{4, {1, 1}}, - BeamConfig{4, {3, 3, 3}}, BeamConfig{4, {1, 2, 3, 4}})), + BeamConfig{4, {3, 3, 3}}, BeamConfig{4, {1, 2, 3, 4}}), + testing::Values(false, true)), [](const testing::TestParamInfo& info) { std::string name{std::get<0>(info.param) == nvinfer1::DataType::kFLOAT ? "Float" : "Half"}; BeamConfig const beamConfig = std::get<1>(info.param); + bool const computeLogProbs = std::get<2>(info.param); name.append("MaxBeamWidth" + std::to_string(beamConfig.maxBeamWidth)); for (auto const beamWdith : beamConfig.beamWidths) { name.append("Bw" + std::to_string(beamWdith)); } + if (computeLogProbs) + { + name.append("LogProbs"); + } return name; }); diff --git a/cpp/tests/runtime/gptDecoderTest.cpp b/cpp/tests/runtime/gptDecoderTest.cpp index 3d1b96f2f23d..32db2ff08a92 100644 --- a/cpp/tests/runtime/gptDecoderTest.cpp +++ b/cpp/tests/runtime/gptDecoderTest.cpp @@ -55,18 +55,17 @@ void testDecoder(nvinfer1::DataType const dtype, SamplingConfig const& samplingC auto const beamWidth = samplingConfig.beamWidth; SizeType constexpr batchSize{4}; - decoder->setup(samplingConfig, batchSize); - - int constexpr endId{50257}; SizeType constexpr maxInputLength{8}; SizeType constexpr maxNewTokens{2}; auto constexpr maxSeqLength = maxInputLength + maxNewTokens; + decoder->setup(samplingConfig, batchSize, maxSeqLength); // set up inputs auto logits = std::shared_ptr( manager.gpu(ITensor::makeShape({batchSize, beamWidth, vocabSizePadded}), modelConfig.getDataType())); manager.setZero(*logits); + int constexpr endId{50257}; std::vector const endIdsVec(batchSize * beamWidth, endId); auto endIds = std::shared_ptr(manager.copyFrom(endIdsVec, ITensor::makeShape({batchSize, beamWidth}), MemoryType::kGPU)); diff --git a/cpp/tests/runtime/gptSessionTest.cpp b/cpp/tests/runtime/gptSessionTest.cpp index 9c4c697552a4..79c5e1dac3a5 100644 --- a/cpp/tests/runtime/gptSessionTest.cpp +++ b/cpp/tests/runtime/gptSessionTest.cpp @@ -40,7 +40,7 @@ namespace fs = std::filesystem; namespace { auto const TEST_RESOURCE_PATH = fs::path{TOP_LEVEL_DIR} / "cpp/tests/resources"; -auto const ENGINGE_PATH = TEST_RESOURCE_PATH / "models/rt_engine"; +auto const ENGINE_PATH = TEST_RESOURCE_PATH / "models/rt_engine"; auto const DATA_PATH = TEST_RESOURCE_PATH / "data"; auto const GPT_MODEL_DIR = "gpt2"; @@ -500,7 +500,7 @@ TEST_P(ParamTest, Test) std::ostringstream gpuSizePath; gpuSizePath << "tp" << modelSpec.mTPSize << "-pp" << modelSpec.mPPSize << "-gpu"; - auto const modelPath{ENGINGE_PATH / modelDir / modelSpec.mModelPath / gpuSizePath.str()}; + auto const modelPath{ENGINE_PATH / modelDir / modelSpec.mModelPath / gpuSizePath.str()}; auto const resultsPath = DATA_PATH / modelDir / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); fs::path const resultsFile{resultsPath / modelSpec.mResultsFile}; @@ -642,7 +642,7 @@ TEST_F(LlamaSessionOnDemandTest, SamplingFP16WithAttentionPlugin) GTEST_SKIP() << "Run only on demand"; auto const modelDir = "llama_7bf"; auto const engineDir = "llama_7bf_outputs_tp1"; - auto const modelPath{ENGINGE_PATH / modelDir / engineDir}; + auto const modelPath{ENGINE_PATH / modelDir / engineDir}; SizeType constexpr beamWidth{1}; fs::path resultsFile{DATA_PATH / modelDir / FP16_RESULT_FILE}; auto const batchSizes = {8}; @@ -659,7 +659,7 @@ TEST_F(LlamaSessionOnDemandTest, SamplingFP16AttentionPluginDecoderBatch) { GTEST_SKIP() << "Run only on demand"; auto const modelDir = "llamav2"; - auto const modelPath{ENGINGE_PATH / modelDir}; + auto const modelPath{ENGINE_PATH / modelDir}; SizeType constexpr beamWidth{1}; fs::path resultsFile{DATA_PATH / modelDir / FP16_RESULT_FILE}; auto const batchSizes = {8}; @@ -676,11 +676,11 @@ class ChatGlmSessionTest : public SessionTest // for ChatGLM-6B { }; -class ChatGlm2SessionTest : public SessionTest // for ChatGLM2-6B and ChatGLM2-6B-32k +class ChatGlm2SessionTest : public SessionTest // for ChatGLM2-6B { }; -class ChatGlm3SessionTest : public SessionTest // for ChatGLM3-6B and ChatGLM3-6B-32k +class ChatGlm3SessionTest : public SessionTest // for ChatGLM3-6B { }; @@ -691,7 +691,7 @@ namespace { // TODO: consolidate this function with testGptSession -// Notice: all ChatGLM models (ChatGLM-6B, ChatGLM2-6B, ChatGLM3-6B, ChatGLM2-6B-32k and ChatGLM3-6B-32k) use this +// Notice: all ChatGLM / GLM models use this // function The differences are GptModelConfig::ModelVariant void testChatGlmSession(fs::path const& modelPath, std::string const& modelName, ModelSpec const& modelSpec, ModelIds const modelIds, SizeType beamWidth, std::initializer_list const& batchSizes, @@ -704,7 +704,7 @@ void testChatGlmSession(fs::path const& modelPath, std::string const& modelName, std::string fileNameSuffix = std::string("-BS") + std::to_string(batchSize) + "-BM" + std::to_string(beamWidth) + std::string(".npy"); fs::path givenInputPath = DATA_PATH / modelName / (std::string("inputId") + fileNameSuffix); - auto const& givenInput = utils::loadNpy(manager, givenInputPath, MemoryType::kCPU); + auto const& givenInput = utils::loadNpy(manager, givenInputPath.string(), MemoryType::kCPU); auto const& inputShape = givenInput->getShape(); ASSERT_EQ(inputShape.nbDims, 2); ASSERT_GT(inputShape.d[0], 0); @@ -730,7 +730,7 @@ void testChatGlmSession(fs::path const& modelPath, std::string const& modelName, ASSERT_TRUE(fs::exists(enginePath)); auto const maxInputLength = static_cast(inputShape.d[1]); - auto const maxNewTokens = 1024; + auto const maxNewTokens = 512; auto const maxSeqLengthGroundTruth = static_cast(outputShape.d[2]); auto const maxSeqLength = maxInputLength + maxNewTokens; SamplingConfig samplingConfig{beamWidth}; @@ -866,8 +866,8 @@ void testChatGlmSession(fs::path const& modelPath, std::string const& modelName, TEST_F(ChatGlmSessionTest, SamplingFP16WithGptAttentionPluginBS1BM1) { - auto const modelName{"chatglm-6b"}; - auto const modelPath{ENGINGE_PATH / "chatglm"}; + auto const modelName{"chatglm_6b"}; + auto const modelPath{ENGINE_PATH / "chatglm"}; auto const batchSizes = {1}; auto constexpr dtype = nvinfer1::DataType::kHALF; auto const modelSpec = ModelSpec{"", "", dtype}.useGptAttentionPlugin(); @@ -876,22 +876,10 @@ TEST_F(ChatGlmSessionTest, SamplingFP16WithGptAttentionPluginBS1BM1) testChatGlmSession(modelPath, modelName, modelSpec, modeIds, 1, batchSizes, mLogger, false, MicroBatchSizes()); } -TEST_F(ChatGlmSessionTest, SamplingFP16WithGptAttentionPluginBS2BM1) -{ - auto const modelName{"chatglm-6b"}; - auto const modelPath{ENGINGE_PATH / "chatglm"}; - auto const batchSizes = {2}; - auto constexpr dtype = nvinfer1::DataType::kHALF; - auto const modelSpec = ModelSpec{"", "", dtype}.useGptAttentionPlugin(); - auto const modeIds = ModelIds{130005, 130005}; - - testChatGlmSession(modelPath, modelName, modelSpec, modeIds, 1, batchSizes, mLogger, false, MicroBatchSizes()); -} - TEST_F(ChatGlm2SessionTest, SamplingFP16WithGptAttentionPluginBS1BM1) { - auto const modelName{"chatglm2-6b"}; - auto const modelPath{ENGINGE_PATH / "chatglm"}; + auto const modelName{"chatglm2_6b"}; + auto const modelPath{ENGINE_PATH / "chatglm"}; auto const batchSizes = {1}; auto constexpr dtype = nvinfer1::DataType::kHALF; auto const modelSpec = ModelSpec{"", "", dtype}.useGptAttentionPlugin(); @@ -902,8 +890,8 @@ TEST_F(ChatGlm2SessionTest, SamplingFP16WithGptAttentionPluginBS1BM1) TEST_F(ChatGlm2SessionTest, SamplingFP16WithGptAttentionPluginBS2BM1) { - auto const modelName{"chatglm2-6b"}; - auto const modelPath{ENGINGE_PATH / "chatglm"}; + auto const modelName{"chatglm2_6b"}; + auto const modelPath{ENGINE_PATH / "chatglm"}; auto const batchSizes = {2}; auto constexpr dtype = nvinfer1::DataType::kHALF; auto const modelSpec = ModelSpec{"", "", dtype}.useGptAttentionPlugin(); @@ -914,8 +902,8 @@ TEST_F(ChatGlm2SessionTest, SamplingFP16WithGptAttentionPluginBS2BM1) TEST_F(ChatGlm2SessionTest, SamplingFP16WithGptAttentionPluginBS1BM2) { - auto const modelName{"chatglm2-6b"}; - auto const modelPath{ENGINGE_PATH / "chatglm"}; + auto const modelName{"chatglm2_6b"}; + auto const modelPath{ENGINE_PATH / "chatglm"}; auto const batchSizes = {1}; auto constexpr dtype = nvinfer1::DataType::kHALF; auto const modelSpec = ModelSpec{"", "", dtype}.useGptAttentionPlugin(); @@ -926,8 +914,8 @@ TEST_F(ChatGlm2SessionTest, SamplingFP16WithGptAttentionPluginBS1BM2) TEST_F(ChatGlm3SessionTest, SamplingFP16WithGptAttentionPluginBS1BM1) { - auto const modelName{"chatglm3-6b"}; - auto const modelPath{ENGINGE_PATH / "chatglm"}; + auto const modelName{"chatglm3_6b"}; + auto const modelPath{ENGINE_PATH / "chatglm"}; auto const batchSizes = {1}; auto constexpr dtype = nvinfer1::DataType::kHALF; auto const modelSpec = ModelSpec{"", "", dtype}.useGptAttentionPlugin(); diff --git a/cpp/tests/runtime/tllmBuffersTest.cpp b/cpp/tests/runtime/tllmBuffersTest.cpp index 40c207d641bd..ac6b3589a2da 100644 --- a/cpp/tests/runtime/tllmBuffersTest.cpp +++ b/cpp/tests/runtime/tllmBuffersTest.cpp @@ -294,7 +294,7 @@ void testBufferType() using limits = std::numeric_limits; static_assert(dataType.isPointer() || dataType.isUnsigned() != limits::is_signed); static_assert(std::is_same_v::type>); + typename DataTypeTraits::type>); IBuffer::SharedPtr buffer{std::make_shared(size, dataType, allocator)}; auto bufferPtr = bufferCast(*buffer); auto constexpr max = limits::max(); diff --git a/cpp/tests/runtime/torchTest.cpp b/cpp/tests/runtime/torchTest.cpp index e4f54a05b4e1..b4711bb985cd 100644 --- a/cpp/tests/runtime/torchTest.cpp +++ b/cpp/tests/runtime/torchTest.cpp @@ -57,7 +57,7 @@ void checkFilled(IBuffer& buffer, int fillValue) { if (DType == buffer.getDataType()) { - EXPECT_THAT(BufferRange::type>(buffer), ::testing::Each(fillValue)); + EXPECT_THAT(BufferRange::type>(buffer), ::testing::Each(fillValue)); } } } // namespace diff --git a/docker/Dockerfile.multi b/docker/Dockerfile.multi index 7fb0134507dc..6a13258207e5 100644 --- a/docker/Dockerfile.multi +++ b/docker/Dockerfile.multi @@ -1,6 +1,6 @@ # Multi-stage Dockerfile ARG BASE_IMAGE=nvcr.io/nvidia/pytorch -ARG BASE_TAG=23.08-py3 +ARG BASE_TAG=23.10-py3 FROM ${BASE_IMAGE}:${BASE_TAG} as base @@ -19,10 +19,16 @@ COPY docker/common/install_cmake.sh install_cmake.sh RUN bash ./install_cmake.sh && rm install_cmake.sh # Download & install internal TRT release -ARG RELEASE_URL_TRT -ARG TARGETARCH -ENV RELEASE_URL_TRT=$RELEASE_URL_TRT -ENV TRT_TARGETARCH=$TARGETARCH +ARG TRT_VER="9.1.0.4" +ENV TRT_VER=$TRT_VER +ARG CUDA_VER="12.2" +ENV CUDA_VER=$CUDA_VER +ARG CUDNN_VER="8.9.4.25-1+cuda12.2" +ENV CUDNN_VER=$CUDNN_VER +ARG NCCL_VER="2.18.3-1+cuda12.2" +ENV NCCL_VER=$NCCL_VER +ARG CUBLAS_VER="12.2.5.6-1" +ENV CUBLAS_VER=$CUBLAS_VER COPY docker/common/install_tensorrt.sh install_tensorrt.sh RUN bash ./install_tensorrt.sh && rm install_tensorrt.sh diff --git a/docker/Makefile b/docker/Makefile index e6646f5d8e58..7e3b365af56b 100644 --- a/docker/Makefile +++ b/docker/Makefile @@ -27,6 +27,11 @@ DOCKER_PROGRESS ?= auto CUDA_ARCHS ?= BUILD_WHEEL_ARGS ?= $(shell grep 'ARG BUILD_WHEEL_ARGS=' Dockerfile.multi | grep -o '=.*' | tr -d '="')$(if $(CUDA_ARCHS), --cuda_architectures $(CUDA_ARCHS)) TORCH_INSTALL_TYPE ?= skip +CUDA_VERSION ?= +CUDNN_VERSION ?= +NCCL_VERSION ?= +CUBLAS_VERSION ?= +TRT_VERSION ?= define add_local_user docker build \ @@ -50,6 +55,11 @@ endef $(if $(BASE_TAG), --build-arg BASE_TAG=$(BASE_TAG)) \ $(if $(BUILD_WHEEL_ARGS), --build-arg BUILD_WHEEL_ARGS="$(BUILD_WHEEL_ARGS)") \ $(if $(TORCH_INSTALL_TYPE), --build-arg TORCH_INSTALL_TYPE="$(TORCH_INSTALL_TYPE)") \ + $(if $(CUDA_VERSION), --build-arg CUDA_VER="$(CUDA_VERSION)") \ + $(if $(CUDNN_VERSION), --build-arg CUDNN_VER="$(CUDNN_VERSION)") \ + $(if $(NCCL_VERSION), --build-arg NCCL_VER="$(NCCL_VERSION)") \ + $(if $(CUBLAS_VERSION), --build-arg CUBLAS_VER="$(CUBLAS_VERSION)") \ + $(if $(TRT_VERSION), --build-arg TRT_VER="$(TRT_VERSION)") \ $(if $(STAGE), --target $(STAGE)) \ --file Dockerfile.multi \ --tag $(IMAGE_WITH_TAG) \ @@ -92,23 +102,29 @@ wheel_%: STAGE = wheel release_%: STAGE = release +# For x86_64 and aarch64 jenkins_%: IMAGE_WITH_TAG = $(shell grep 'LLM_DOCKER_IMAGE = ' ../jenkins/L0_MergeRequest.groovy | grep -o '".*"' | tr -d '"') jenkins_%: STAGE = devel +# For x86_64 centos7_%: IMAGE_WITH_TAG = $(shell grep 'LLM_CENTOS7_DOCKER_IMAGE = ' ../jenkins/L0_MergeRequest.groovy | grep -o '".*"' | tr -d '"') centos7_%: STAGE = devel -centos7_%: TORCH_INSTALL_TYPE = src_cxx11_abi centos7_%: BASE_IMAGE = nvidia/cuda -centos7_%: BASE_TAG = 12.2.0-devel-centos7 +centos7_%: BASE_TAG = 12.2.2-devel-centos7 +# For x86_64 and aarch64 ubuntu22_%: STAGE = devel -ubuntu22_%: TORCH_INSTALL_TYPE = src_cxx11_abi ubuntu22_%: BASE_IMAGE = nvidia/cuda -ubuntu22_%: BASE_TAG = 12.2.0-devel-ubuntu22.04 +ubuntu22_%: BASE_TAG = 12.2.2-devel-ubuntu22.04 +# For x86_64 and aarch64 old-cuda_%: IMAGE_WITH_TAG = $(shell grep 'LLM_OLD_CUDA_DOCKER_IMAGE = ' ../jenkins/L0_MergeRequest.groovy | grep -o '".*"' | tr -d '"') old-cuda_%: BASE_TAG = 23.07-py3 old-cuda_%: STAGE = devel +old-cuda_%: CUDA_VERSION = 12.1 +old-cuda_%: CUDNN_VERSION = 8.9.3.28-1+cuda12.1 +old-cuda_%: NCCL_VERSION = 2.18.3-1+cuda12.1 +old-cuda_%: CUBLAS_VERSION = 12.1.3.1-1 build: devel_build ; diff --git a/docker/common/install_base.sh b/docker/common/install_base.sh index b4e94897882b..314fc916aa1e 100644 --- a/docker/common/install_base.sh +++ b/docker/common/install_base.sh @@ -30,6 +30,7 @@ init_ubuntu() { fi apt-get clean rm -rf /var/lib/apt/lists/* + echo "export LD_LIBRARY_PATH=/usr/local/cuda/lib64:\$LD_LIBRARY_PATH" >> "${ENV}" # Remove previous TRT installation if [[ $(apt list --installed | grep libnvinfer) ]]; then apt-get remove --purge -y libnvinfer* @@ -63,10 +64,11 @@ init_centos() { yum -y update yum -y install centos-release-scl-rh epel-release # https://gitlab.com/nvidia/container-images/cuda + echo "export LD_LIBRARY_PATH=/usr/local/cuda/lib64:\$LD_LIBRARY_PATH" >> "${ENV}" CUDA_VERSION=$(nvcc --version | sed -n 's/^.*release \([0-9]\+\.[0-9]\+\).*$/\1/p') YUM_CUDA=${CUDA_VERSION/./-} # Consistent with manylinux2014 centos-7 based version - yum -y install wget rh-python${PY_VERSION} rh-python${PY_VERSION}-python-devel rh-git227 devtoolset-10 libffi-devel + yum -y install wget git-lfs rh-python${PY_VERSION} rh-python${PY_VERSION}-python-devel rh-git227 devtoolset-10 libffi-devel yum -y install openmpi3 openmpi3-devel echo "source scl_source enable rh-git227 rh-python38" >> "${ENV}" echo "source scl_source enable devtoolset-10" >> "${DEVTOOLSET_ENV_FILE}" diff --git a/docker/common/install_tensorrt.sh b/docker/common/install_tensorrt.sh index 882f645d9f77..1a735df3e5cd 100644 --- a/docker/common/install_tensorrt.sh +++ b/docker/common/install_tensorrt.sh @@ -2,28 +2,53 @@ set -ex +NVCC_VERSION_OUTPUT=$(nvcc --version) +if [[ $(echo $NVCC_VERSION_OUTPUT | grep -oP "\d+\.\d+" | head -n 1) != ${CUDA_VER} ]]; then + echo "The version of pre-installed CUDA is not equal to ${CUDA_VER}." + exit 1 +fi + install_ubuntu_requirements() { - CUDNN_VERSION="8" + apt-get update && apt-get install -y --no-install-recommends gnupg2 curl ca-certificates + ARCH=$(uname -m) + if [ "$ARCH" = "amd64" ];then ARCH="x86_64";fi + if [ "$ARCH" = "aarch64" ];then ARCH="sbsa";fi + curl -fsSLO https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/${ARCH}/cuda-keyring_1.0-1_all.deb + dpkg -i cuda-keyring_1.0-1_all.deb + apt-get update - apt-get install -y --no-install-recommends libcudnn${CUDNN_VERSION} libcudnn${CUDNN_VERSION}-dev libnccl-dev + if [[ $(apt list --installed | grep libcudnn8) ]]; then + apt-get remove --purge -y libcudnn8* + fi + if [[ $(apt list --installed | grep libnccl) ]]; then + apt-get remove --purge -y --allow-change-held-packages libnccl* + fi + if [[ $(apt list --installed | grep libcublas) ]]; then + apt-get remove --purge -y --allow-change-held-packages libcublas* + fi + CUBLAS_CUDA_VERSION=$(echo $CUDA_VER | sed 's/\./-/g') + apt-get install -y --no-install-recommends libcudnn8=${CUDNN_VER} libcudnn8-dev=${CUDNN_VER} + apt-get install -y --no-install-recommends libnccl2=${NCCL_VER} libnccl-dev=${NCCL_VER} + apt-get install -y --no-install-recommends libcublas-${CUBLAS_CUDA_VERSION}=${CUBLAS_VER} libcublas-dev-${CUBLAS_CUDA_VERSION}=${CUBLAS_VER} apt-get clean rm -rf /var/lib/apt/lists/* } install_centos_requirements() { - CUDNN_VERSION="8" + CUDNN_VER=$(echo $CUDNN_VER | sed 's/+/./g') + CUBLAS_CUDA_VERSION=$(echo $CUDA_VER | sed 's/\./-/g') yum -y update yum -y install epel-release - yum -y install libcudnn${CUDNN_VERSION} libcudnn${CUDNN_VERSION}-devel libnccl-devel + yum remove -y libcudnn* && yum -y install libcudnn8-${CUDNN_VER} libcudnn8-devel-${CUDNN_VER} + yum remove -y libnccl* && yum -y install libnccl-${NCCL_VER} libnccl-devel-${NCCL_VER} + yum remove -y libcublas* && yum -y install libcublas-${CUBLAS_CUDA_VERSION}-${CUBLAS_VER} libcublas-devel-${CUBLAS_CUDA_VERSION}-${CUBLAS_VER} yum clean all } install_tensorrt() { - TENSOR_RT_VERSION="9.1.0.4" - CUDA_VERSION="12.2" - PY_VERSION=$(python -c 'import sys; print(".".join(map(str, sys.version_info[0:2])))') PARSED_PY_VERSION=$(echo "${PY_VERSION//./}") + TRT_CUDA_VERSION="12.2" if [ -z "$RELEASE_URL_TRT" ];then ARCH=${TRT_TARGETARCH} @@ -32,11 +57,11 @@ install_tensorrt() { if [ "$ARCH" = "amd64" ];then ARCH="x86_64";fi if [ "$ARCH" = "x86_64" ];then DIR_NAME="x64-agnostic"; else DIR_NAME=${ARCH};fi if [ "$ARCH" = "aarch64" ];then OS="ubuntu-22.04"; else OS="linux";fi - RELEASE_URL_TRT=https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/secure/9.1.0/tars/tensorrt-${TENSOR_RT_VERSION}.${OS}.${ARCH}-gnu.cuda-${CUDA_VERSION}.tar.gz; + RELEASE_URL_TRT=https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/secure/9.1.0/tars/tensorrt-${TRT_VER}.${OS}.${ARCH}-gnu.cuda-${TRT_CUDA_VERSION}.tar.gz; fi wget --no-verbose ${RELEASE_URL_TRT} -O /tmp/TensorRT.tar tar -xf /tmp/TensorRT.tar -C /usr/local/ - mv /usr/local/TensorRT-${TENSOR_RT_VERSION} /usr/local/tensorrt + mv /usr/local/TensorRT-${TRT_VER} /usr/local/tensorrt pip install /usr/local/tensorrt/python/tensorrt-*-cp${PARSED_PY_VERSION}-*.whl rm -rf /tmp/TensorRT.tar echo 'export LD_LIBRARY_PATH=/usr/local/tensorrt/lib:$LD_LIBRARY_PATH' >> "${ENV}" diff --git a/docs/source/blogs/H200launch.md b/docs/source/blogs/H200launch.md index b20c91385fed..077f442f63da 100644 --- a/docs/source/blogs/H200launch.md +++ b/docs/source/blogs/H200launch.md @@ -33,7 +33,7 @@ For practical examples of H200's performance: **Max Throughput TP8:** an online chat agent scenario (ISL/OSL=80/200) with GPT3-175B on a full HGX (TP8) H200 is 1.6x more performant than H100. -max throughput llama TP1 +H200 TPS Preliminary measured performance, subject to change. TensorRT-LLM v0.5.0, TensorRT v9.1.0.4. | Llama-70B: H100 FP8 BS 8, H200 FP8 BS 32 | GPT3-175B: H100 FP8 BS 64, H200 FP8 BS 128 diff --git a/docs/source/gpt_runtime.md b/docs/source/gpt_runtime.md index 87782b02ec60..eeea339c7e9e 100644 --- a/docs/source/gpt_runtime.md +++ b/docs/source/gpt_runtime.md @@ -322,6 +322,11 @@ batchSize, beamWidth]`_. that enabling that computation may have an impact on performance (the final LM head has to perform a matrix multiplication on all the context tokens instead of a just the last one), + * `generationLogits`, is a tensor of values on the GPU (same datatype as the + computation type) to store the logits for the generation. Its shape is + `[batchSize, beamWidth, maxOutputLen-1, vocabSizePadded]`. This buffer will only be + filled in if the TensorRT engine was built with the + `gather_all_token_logits` parameter enabled. * `onTokenGenerated`, is a callback function invoked in the generation loop to pass newly generated tokens to the caller while the loop continues to execute. An implementation of that callback must accept the output `ids` diff --git a/docs/source/index.rst b/docs/source/index.rst index 57b4e66fc74b..6d15d7cea628 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -19,6 +19,7 @@ Welcome to TensorRT-LLM's documentation! 2023-05-19-how-to-debug.md 2023-05-17-how-to-add-a-new-model.md graph-rewriting.md + memory.md Python API ---------- diff --git a/docs/source/installation.md b/docs/source/installation.md index 9d785bcdcfa2..1cb6080844b9 100644 --- a/docs/source/installation.md +++ b/docs/source/installation.md @@ -145,12 +145,27 @@ example: ```bash # Build TensorRT-LLM for Ampere. -python3 ./scripts/build_wheel.py --cuda_architectures "80-real;86-real" +python3 ./scripts/build_wheel.py --cuda_architectures "80-real;86-real" --trt_root /usr/local/tensorrt ``` The list of supported architectures can be found in the [`CMakeLists.txt`](source:cpp/CMakeLists.txt) file. +### Build the Python Bindings for the C++ Runtime + +The C++ Runtime, in particular, [`GptSession`](../../cpp/include/tensorrt_llm/runtime/gptSession.h) can be exposed to +Python via [bindings](../../cpp/tensorrt_llm/pybind/bindings.cpp). This is currently an opt-in feature which needs to be +explicitly activated during compilation time. The corresponding option `--python_bindings` can be specified +to `build_wheel.py` in the standard way: + +```bash +python3 ./scripts/build_wheel.py --python_bindings --trt_root /usr/local/tensorrt +``` + +After installing the resulting wheel as described above, the C++ Runtime bindings will be available in +package `tensorrt_llm.bindings`. Running `help` on this package in a Python interpreter will provide on overview of the +relevant classes. The [associated unit tests](../../tests/bindings) should also be consulted for understanding the API. + ### Link with the TensorRT-LLM C++ Runtime The `build_wheel.py` script will also compile the library containing the C++ diff --git a/docs/source/memory.md b/docs/source/memory.md new file mode 100644 index 000000000000..069dfc40f7b0 --- /dev/null +++ b/docs/source/memory.md @@ -0,0 +1,114 @@ +# Memory Usage of TensorRT-LLM + + +This document summarizes the memory usage of TensorRT-LLM, and addresses common issues and questions reported by users. + + +## Understand inference time GPU memory usage + + +At inference time, there are 3 major contributors to GPU memory usage for a given TRT engine generated from a TensorRT-LLM model: weights, internal activation tensors, and IO tensors. For IO tensors, the major memory footprint comes from the KV cache tensor. + + +### Weights size + +Weights size is fixed depending on the model size, the chosen precision of the weights and the parallelization strategy. +Using lower precision like INT8 or FP8 can reduce the weights size. +When tensor parallelism or pipeline parallelism is used, each rank stores only some portion of the weights. +For example, each rank typically uses just 1/8 of the model weights when using 8-way tensor parallelism or 8-stages pipeline parallelism. + + +### Activation size + + +TensorRT can optimize the memory usage by reusing memory for different tensors based on live analysis and tensor size. To avoid out of memory errors at runtime and to reduce the runtime cost of switching optimization profiles and changing shapes, **TensorRT pre-computes the activation tensors memory requirement at build time**. The memory requirement is computed based on an optimized TensorRT graph, one profiles’ memory usage is computed by using the max tensor shape, and the memory requirement of one engine is computed by the maximum size between different profiles. There are external and internal factors that can affect the activation size returned by TensorRT, such as the network structure, kernel fusion, operation scheduling, etc. +Once the TensorRT engine is built, the activation memory size of that engine can be queried by the API `trt.ICudaEngine.device_memory_size`. + + +Practically, for a given model, specified precision and parallelization strategy, one can tune the activation memory usage by adjusting the max batch size, max input length, max beam width, max number of tokens, padding removal on/off flag, context FMHA on/off flag. +Here some explanations on how these values affect the memory: + + +1. Reduce build time max input tokens + + Most of the tensors inside a transformer network have a linear relationship with number of input tokens, so activation size will be close to `max number of input tokens * some constant factor`, the constant factor depends on the network structure and TRT internal optimization. The max number of input tokens is derived from build time arguments, one can change the parameters provided to the `prepare_inputs` function, like `GPTLMHeadModel.prepare_inputs` to affect the memory usage, or one can change the command line options of the `build.py` scripts used in the examples. + + When using the [padded tensors](./gpt_attention.md#padded-and-packed-tensors) format, the max number of input tokens equals to `max_batch_size*max_input_len`, so reducing `max_batch_size` and `max_input_len` can almost linearly reduce the activation memory size. + When using the [packed tensors](./gpt_attention.md#padded-and-packed-tensors) format and `max_num_tokens` is specified, reducing its value will also reduce activation memory size. If `max_num_tokens` is not specified, the max number of input tokens will be derived as `max_batch_size*max_input_len`. + + The packed tensors format is recommended, because it saves both memory and compute. + The beam width will be folded into the batch size dimension when passing the tensors range into TensorRT, so reducing `max_beam_width` can also reduce the memory usage. + + +2. Turn on context FMHA + + When the GPT attention plugin is used, turning on the `context_fmha_type` of the plugin will reduce the memory footprint significantly. See the [Context Phase](./gpt_attention.md#context-phase) for details. When the `context_fmha_type` is set to disabled, a workspace size of the plugin will quadratically depend on the sequence length. + + +3. Tensor parallelism and pipeline parallelism + + TensorRT will reuse memory between layers as much as possible, for a typical example, given *N* decoder blocks in one transformer network, TRT will not allocate *N* copies of the activation memory for each block, since the memory of tensors in the 1st block can be released after the execution, memory can be reused for later blocks, only 1 block’s memory is needed. + + + When using tensor parallelism, some tensors are split into smaller chunks and each rank only holds one chunk of the tensor, the activation memory size of each rank will be smaller than when executing the network on a single GPU. When using pipeline parallelism, each rank executes several decoder blocks, and all the tensors are full-size tensors, so the activation memory size is equal to 1 block’s memory size. Thus tensor parallelism normally has higher memory efficiency than pipeline parallelism when all other parameters are the same. + + +## KV cache tensor + +### Python runtime + +The Python runtime allocates KV cache tensors based on the parameters of the `GenerationSession.setup` function, the KV cache size is linearly dependent on the `batch_size` and `max_context_length+max_new_tokens`. **Note: This may change in the future, as the Python bindings of the C++ runtime may replace the current python runtime in the future. The Python bindings of C++ runtime behave like C++ runtime.** + +### C++ runtime + +* When paged KV cache is enabled + + TensorRT-LLM runtime pre-allocates KV cache tensors during initialization for a configured number of blocks and distributes them at runtime. + KV cache tensors are allocated based on the `KVCacheConfig` object when creating `GptSession`. If neither `maxTokens` nor `freeGpuMemoryFraction` is specified, KV cache will by default allocate 85% of the remaining free GPU memory. When either `maxTokens` or `freeGpuMemoryFraction` is specified, the specified value will be used to compute the KV cache memory size. And if both are specified, firstly the `freeGpuMemoryFraction` is used to compute the number of tokens in KV cache, and then the minimum between this computed number of tokens and `maxTokens` is used. + + In in-flight batching the scheduler can automatically schedule requests as long as enough KV cache space is available (exact behavior depends on the scheduler policy). + If paged KV cache is used in `GptSession` without in-flight batching, TensorRT-LLM may report OOM errors with message "Can't allocate new blocks. No free blocks left", if the paged KV cache is not large enough for the whole batch. + +* When paged KV cache is disabled + + C++ runtime allocates the KV cache tensors for each layer with shape `[batch size, 2, heads, max seq length, hidden dimension per head]`, where `max seq length` is specified by `GptSession::Config::maxSequenceLength` when creating `GptSession`. + +## Memory pool + +TensorRT-LLM C++ runtime is using stream-ordered memory allocator to allocate and free buffers, see [BufferManager::initMemoryPool](cpp/tensorrt_llm/runtime/bufferManager.cpp), which uses the default memory pool managed by the CUDA driver. When a `GptSession` object is destroyed, memory is returned to the memory pool and can be reused by the next instance of a `GptSession` object. Memory will be released from the pool if it is required for other memory allocations. +However, `nvidia-smi` may still show high memory occupation after memory is returned to the CUDA driver's memory pool. This should not be a concern and is intended behavior. The amount of reserved and free memory in the pool can be inspected by [BufferManager::memoryPoolReserved())](cpp/tensorrt_llm/runtime/bufferManager.cpp) and [BufferManager::memoryPoolFree())](cpp/tensorrt_llm/runtime/bufferManager.cpp), respectively. + +## Known Issues + + +1. When 2 optimization profiles are used, the weights memory size may be doubled in some cases if the underlying kernel choices require different weights layouts for different optimization profiles. This issue will be fixed in a future release. + +2. When FP8 GEMM is used, the activation memory might be larger than the theoretical optimized memory size, this will be enhanced in a future release. + +## FAQ + +1. Why is the memory size large even though a small batch size and sequence length are used in the runtime? + + As explained above, the activation memory size is computed based on the max tensor shapes at TensorRT engine building time, try to reduce the engine building time parameters like `max_num_token`, `max_batch_size`, `max_input_len`, see [Activation size](#activation-size) for details. + + +2. Why can the engine be generated, but the inference will run out of memory (OOM) at runtime? + + At engine building time, TensorRT will tune the kernel selection layer by layer, it does not necessarily allocate all the memory required to run the entire engine. If the activation tensors required to run a single layer are small, while the I/O tensor (like KV cache) sizes required to run the engine are large, building will succeed since it may not need to allocate the large I/O tensors, runtime may fail with OOM errors on allocating large IO tensors. + + TensorRT-LLM has provided a `check_gpt_mem_usage` utility function to check the upper bound of the memory size given an engine, and the related batch size, I/O sequence length, etc., when the upper boundary check exceeded the GPU physical memory size, warning messages will be printed. + +3. How to debug the memory usage of TensorRT-LLM? + + When the verbose logging level is used, TensorRT and TensorRT-LLM will print messages about memory usage details. + The line showing "Total Weights Memory" indicates the weights memory size, and the line "Total Activation Memory" indicates the activation memory size. + + Normally the weights memory size is close to the TensorRT engine size, since most of the content in the engine is from weights for LLM networks. + +4. For pipeline parallelism, is build time max batch size the limit of micro batch size? + + Yes, in pipeline parallel mode, TensorRT-LLM runtime will split the batch of requests into micro batches, and enqueue these micro batches into TRT engine sequentially. + The `max_batch_size` at build time means that batch size of one engine enqueue call shall be smaller than it. The total batch size before splitting into micro batches can be larger than the build time `max_batch_size`. + + For example, if you have 4-stages pipeline parallelism, and intend to run the engine using micro batch size 2 and run 16 micro batches (total batch size 32) in one `generate` call. + You could just set the `max_batch_size` at building time to 2, instead of 32. Setting build time `max_batch_size` 32 will occupy almost 16x more activation memory. diff --git a/docs/source/performance.md b/docs/source/performance.md index 6d7e38b19b7f..993ce0eb5041 100644 --- a/docs/source/performance.md +++ b/docs/source/performance.md @@ -142,3 +142,317 @@ The simplest implementation uses two Matmul operations and combines the results in a separate CUDA kernel. That's the current implementation in TensorRT-LLM. The next release will include a more efficient implementation that runs a single Matmul. + + +## Reproducing Benchmarked Results + +### Building the TensorRT-LLM Container +--- +In order to benchmark TensorRT-LLM, you will need to follow the [Quick Start](../../README.md#quick-start) +build process to create a baseline container for building a wheel. Additionally, the development +container needs a copy of the source code to build the wheel and the benchmarking script. Create the +right build environment, use the following : + +```shell +git clone https://github.com/NVIDIA/TensorRT-LLM.git +cd TensorRT-LLM +git submodule update --init --recursive +git lfs install +git lfs pull +make -C docker build +make -C docker run LOCAL_USER=1 +``` + +> [!WARNING] +> If you have elevated privileges on your system, then skip the `make -C docker run LOCAL_USER=1` +command above as it may make it so that you cannot access some required system libraries within the +container because the build forces your UID and GID to match those that are set for your non-elevated +user. There are cases where the container will be booted as root (i.e. on some SLURM systems with +the pyxis plugin) which will cause libraries to be missing. + +If you are benchmarking in a shared environment, you need to specify the GPU indices that you would +like the container to use, otherwise the Makefile defaults to loading the container with all GPUs on +the system. For example, if you only have the 4 higher indices of GPUs on your system you can +configure it using the following example: + +```shell +NV_GPU=0,1,2,3 +make -C docker run LOCAL_USER=1 GPU_OPTS='--gpus \"device=${NV_GPU}\"' +``` + +Additionally, if you'd like to mount external storage to access persistent storage, or previously +built engines, you can mount directories as follows (simply replace `source` and `destination` with +the appropriate paths): + +```shell +make -C docker run LOCAL_USER=1 DOCKER_RUN_ARGS="-v /source:/destination" +``` + +Once the container starts, you'll need to build the wheel and the benchmarking scripts. From the +code root (the default directory when the container is loaded), the following commands will build +the TensorRT-LLM wheel, install dependencies, and build the benchmark scripts: + +```shell +python3 ./scripts/build_wheel.py --benchmarks --trt_root /usr/local/tensorrt +pip install ./build/tensorrt_llm*.whl +``` + +## Methodology + +### Engine Building Setups + +Each engine needs to be built before they can be benchmarked, and requires the source code for each +of their respective build scripts. For smaller models, it is fine to build the engine on the fly in +container; however, for larger engines it is recommended to pre-build and mount a directory with the +engine because engine files are quite large and take time to repeatedly build. Additionally, built +engines can be used for input lengths, output lengths, and batch sizes *up to* their build options +meaning you can use an engine to benchmark multiple input configurations. + +In order to benchmark the various networks, our engine building scheme is as follows: +- For the GPT-J, Llama2-7b, and Llama2-70b benchmarks were ran using a single-setting engine build +for each network configured for our maximum expected throughput. +- For Falcon-180B, where memory limits and model size have a higher impact for running the model, +our benchmarks transition to a per-configuration engine build. + +Below we document how to benchmark each model on an H100-HBM3-80GB system and reproduce the throughput +numbers we document on our [Performance section](#performance of-tensorrt-llm). + +### Running on A100 + +To run the benchmarks below on A100, you will need to remove the `--enable_fp8 --fp8_kv_cache` options +from each engine build command because FP8 computation is a feature in H100 and newer GPUs. + +### Reproducing First Token Latency + +In order to test the latency to the first token, you can build the engines as specified below (or +with the tweaks specified above on A100) -- once built as described in the +[build steps](#engine-building-setups) above, you can then benchmark with a single output token in +order to find the time to first token latency. We provide the appropriate command lines below for +each of the benchmarked models, but you can use this same method to benchmark other models available +in [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM). + +## Benchmarking per Model + +#### GPT-J 6B +--- +```shell +python examples/gptj/build.py \ + --enable_context_fmha \ + --parallel_build \ + --output_dir /tmp/engines/gptj \ + --dtype float16 \ + --use_gpt_attention_plugin float16 \ + --world_size 1 \ + --max_batch_size 64 \ + --max_input_len 2048 \ + --max_output_len 2048 \ + --hidden_act gelu \ + --enable_fp8 \ + --fp8_kv_cache \ + --strongly_typed \ + --n_layer 28 \ + --n_head 16 \ + --n_embd 4096 \ + --n_positions 2048 \ + --enable_two_optimization_profiles +``` + +##### Throughput Benchmark + +```shell +in_out_sizes=("64:128,128" "64:128,2048" "64:2048,128" "64:2048,2048") +for in_out in ${in_out_sizes[@]} +do + batch_size=$(echo $in_out | awk -F':' '{ print $1 }') + in_out_dims=$(echo $in_out | awk -F':' '{ print $2 }') + echo "BS: $batch_size, ISL/OSL: $in_out_dims" + + ./cpp/build/benchmarks/gptSessionBenchmark --model gptj --engine_dir /tmp/engines/gptj/ --warm_up 1 --batch_size $batch_size --duration 0 --num_runs 5 --input_output_len $in_out_dims +done +``` + +##### First Token Latency Benchmark + +```shell +in_out_sizes=("64:128,1" "64:2048,1") +for in_out in ${in_out_sizes[@]} +do + batch_size=$(echo $in_out | awk -F':' '{ print $1 }') + in_out_dims=$(echo $in_out | awk -F':' '{ print $2 }') + echo "BS: $batch_size, ISL/OSL: $in_out_dims" + + ./cpp/build/benchmarks/gptSessionBenchmark --model gptj --engine_dir /tmp/engines/gptj/ --warm_up 1 --batch_size $batch_size --duration 0 --num_runs 5 --input_output_len $in_out_dims +done +``` + + +### Llama2-7b +--- +```shell +pip install -r examples/llama/requirements.txt +python examples/llama/build.py \ + --remove_input_padding \ + --enable_context_fmha \ + --parallel_build \ + --output_dir /tmp/engines/llama/7b \ + --dtype float16 \ + --use_gpt_attention_plugin float16 \ + --world_size 1 \ + --tp_size 1 \ + --pp_size 1 \ + --max_batch_size 64 \ + --max_input_len 2048 \ + --max_output_len 2048 \ + --enable_fp8 \ + --fp8_kv_cache \ + --strongly_typed \ + --n_layer 32 \ + --n_head 32 \ + --n_embd 4096 \ + --inter_size 11008 \ + --vocab_size 32000 \ + --n_positions 4096 \ + --hidden_act silu +``` + +##### Throughput Benchmark + +```shell +in_out_sizes=("64:128,128" "64:128,2048" "64:2048,128" "32:2048,2048") +for in_out in ${in_out_sizes[@]} +do + batch_size=$(echo $in_out | awk -F':' '{ print $1 }') + in_out_dims=$(echo $in_out | awk -F':' '{ print $2 }') + echo "BS: $batch_size, ISL/OSL: $in_out_dims" + + ./cpp/build/benchmarks/gptSessionBenchmark --model llama --engine_dir /tmp/engines/llama/7b --warm_up 1 --batch_size $batch_size --duration 0 --num_runs 5 --input_output_len $in_out_dims +done +``` +##### First Token Latency Benchmark + +```shell +in_out_sizes=("64:128,1" "32:2048,1") +for in_out in ${in_out_sizes[@]} +do + batch_size=$(echo $in_out | awk -F':' '{ print $1 }') + in_out_dims=$(echo $in_out | awk -F':' '{ print $2 }') + echo "BS: $batch_size, ISL/OSL: $in_out_dims" + + ./cpp/build/benchmarks/gptSessionBenchmark --model llama --engine_dir /tmp/engines/llama/7b --warm_up 1 --batch_size $batch_size --duration 0 --num_runs 5 --input_output_len $in_out_dims +done +``` + +### Llama2-70b + +```shell +pip install -r examples/llama/requirements.txt +python examples/llama/build.py \ + --remove_input_padding \ + --enable_context_fmha \ + --parallel_build \ + --output_dir /tmp/engines/llama/70b \ + --dtype float16 \ + --use_gpt_attention_plugin float16 \ + --world_size 4 \ + --tp_size 4 \ + --pp_size 1 \ + --max_batch_size 64 \ + --max_input_len 2048 \ + --max_output_len 2048 \ + --enable_fp8 \ + --fp8_kv_cache \ + --strongly_typed \ + --n_layer 80 \ + --n_head 64 \ + --n_kv_head 8 \ + --n_embd 8192 \ + --inter_size 28672 \ + --vocab_size 32000 \ + --n_positions 4096 \ + --hidden_act silu \ + --ffn_dim_multiplier 1.3 \ + --multiple_of 4096 +``` + +##### Throughput Benchmark + +```shell +in_out_sizes=("64:128,128" "64:128,2048" "64:2048,128" "64:2048,2048") +for in_out in ${in_out_sizes[@]} +do + batch_size=$(echo $in_out | awk -F':' '{ print $1 }') + in_out_dims=$(echo $in_out | awk -F':' '{ print $2 }') + echo "BS: $batch_size, ISL/OSL: $in_out_dims" + + mpirun -n 4 --allow-run-as-root --oversubscribe ./cpp/build/benchmarks/gptSessionBenchmark --model llama --engine_dir /tmp/engines/llama/70b --warm_up 1 --batch_size $batch_size --duration 0 --num_runs 5 --input_output_len $in_out_dims +done +``` + +##### First Token Latency Benchmark + +```shell +in_out_sizes=("64:128,1" "64:128,1") +for in_out in ${in_out_sizes[@]} +do + batch_size=$(echo $in_out | awk -F':' '{ print $1 }') + in_out_dims=$(echo $in_out | awk -F':' '{ print $2 }') + echo "BS: $batch_size, ISL/OSL: $in_out_dims" + + mpirun -n 4 --allow-run-as-root --oversubscribe ./cpp/build/benchmarks/gptSessionBenchmark --model llama --engine_dir /tmp/engines/llama/70b --warm_up 1 --batch_size $batch_size --duration 0 --num_runs 5 --input_output_len $in_out_dims +done +``` + + +### Falcon-180B +--- + +Benchmarking Falcon-180B requires a custom engine per batch size, input/output sequence length due +to the large footprint of the model and the large input size of 2048. You can build and benchmark +each engine one at a time with the following loop. + +```shell +# Benchmark specific batch size:isl:osl combinations. +in_out_sizes=("96:128,128" "96:128,2048" "64:2048,128") +for in_out in ${in_out_sizes[@]} +do + batch_size=$(echo $in_out | awk -F':' '{ print $1 }') + in_out_dims=$(echo $in_out | awk -F':' '{ print $2 }') + isl=$(echo $in_out_dims | awk -F',' '{ print $1 }') + osl=$(echo $in_out_dims | awk -F',' '{ print $2 }') + engine_path="/tmp/engines/falcon/180b/${batch_size}_${isl}_${osl}" + echo "BS: $batch_size, ISL/OSL: ${isl},${osl}" + + # Build the specific engine for the BS,ISL,OSL combination + python examples/falcon/build.py \ + --use_inflight_batching \ + --paged_kv_cache \ + --remove_input_padding \ + --enable_context_fmha \ + --parallel_build \ + --output_dir $engine_path \ + --dtype float16 \ + --use_gemm_plugin float16 \ + --use_gpt_attention_plugin float16 \ + --world_size 8 \ + --tp 8 \ + --max_batch_size $batch_size \ + --max_input_len $isl \ + --max_output_len $osl \ + --enable_fp8 \ + --fp8_kv_cache \ + --n_layer 80 \ + --n_head 232 \ + --n_kv_head 8 \ + --n_embd 14848 \ + --vocab_size 65024 \ + --new_decoder_architecture + # Throughput benchmark + mpirun -n 8 --allow-run-as-root --oversubscribe ./cpp/build/benchmarks/gptSessionBenchmark --model falcon --engine_dir $engine_path --warm_up 1 --batch_size $batch_size --duration 0 --num_runs 5 --input_output_len "${isl},${osl}" + # Time to first token benchmark + mpirun -n 8 --allow-run-as-root --oversubscribe ./cpp/build/benchmarks/gptSessionBenchmark --model falcon --engine_dir $engine_path --warm_up 1 --batch_size $batch_size --duration 0 --num_runs 5 --input_output_len "${isl},1" + + # The Falcon-180b engine is quite large, remove after the benchmark to free up space + # Remove this line if you'd like to save the engines. + rm -r $engine_path +done +``` diff --git a/docs/source/precision.md b/docs/source/precision.md index 86d193fe1053..76dc1f2a1605 100644 --- a/docs/source/precision.md +++ b/docs/source/precision.md @@ -118,21 +118,27 @@ This release of TensorRT-LLM contains the following examples: | :--------- | :---: | :---: | :---: | :---: | :-----: | :---: | :---: | :-------: | :--------: | | Baichuan | Y | Y | Y | . | Y | Y | Y | . | . | | BERT | Y | Y | Y | . | . | . | . | . | . | +| BLIP-2 | Y | Y | Y | . | . | . | . | . | . | | BLOOM | Y | Y | Y | . | Y | Y | Y | . | . | | ChatGLM | Y | Y | Y | . | . | . | . | . | . | | ChatGLM-v2 | Y | Y | Y | . | . | . | . | . | . | | ChatGLM-v3 | Y | Y | Y | . | . | . | . | . | . | -| Falcon | Y | Y | Y | . | . | . | . | . | . | +| Falcon | Y | Y | Y | Y | . | . | . | Y | . | +| Flan-T5 | Y | Y | Y | . | . | . | . | . | . | | GPT | Y | Y | Y | Y | Y | Y | Y | . | . | | GPT-J | Y | Y | Y | Y | Y | Y | Y | Y | . | | GPT-NeMo | Y | Y | Y | . | . | . | . | . | . | | GPT-NeoX | Y | Y | Y | . | . | . | . | . | Y | +| InternLM | Y | Y | Y | . | Y | Y | Y | . | . | | LLaMA | Y | Y | Y | . | Y | Y | Y | Y | Y | | LLaMA-v2 | Y | Y | Y | Y | Y | Y | Y | Y | Y | +| Mistral | Y | Y | Y | . | . | . | . | . | . | +| MPT | Y | Y | Y | Y | . | . | . | . | . | | OPT | Y | Y | Y | . | . | . | . | . | . | +| Replit Code| Y | Y | Y | . | . | . | . | . | . | | SantaCoder | Y | Y | Y | . | . | . | . | . | . | | StarCoder | Y | Y | Y | . | . | . | . | . | . | -| InternLM | Y | Y | Y | . | Y | Y | Y | . | . | +| T5 | Y | Y | Y | . | . | . | . | . | . | ## Technical Detail: The `QuantMode` Flags diff --git a/examples/baichuan/README.md b/examples/baichuan/README.md index 241bb540a30f..dadc0471597f 100644 --- a/examples/baichuan/README.md +++ b/examples/baichuan/README.md @@ -4,11 +4,11 @@ This document shows how to build and run a Baichuan models (including `v1_7b`/`v ## Overview -The TensorRT-LLM Baichuan implementation can be found in [tensorrt_llm/models/baichuan/model.py](../../tensorrt_llm/models/baichuan/model.py). The TensorRT-LLM Baichuan example code is located in [`examples/baichuan`](./). There are three main files in that folder:: +The TensorRT-LLM Baichuan implementation can be found in [tensorrt_llm/models/baichuan/model.py](../../tensorrt_llm/models/baichuan/model.py). The TensorRT-LLM Baichuan example code is located in [`examples/baichuan`](./). There are three main files: * [`build.py`](./build.py) to build the [TensorRT](https://developer.nvidia.com/tensorrt) engine(s) needed to run the Baichuan model, * [`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/cnn_dailymail) dataset using the model. + * and a shared [`../summarize.py`](../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/cnn_dailymail) dataset using the model. These scripts accept an argument named model_version, whose value should be `v1_7b`/`v1_13b`/`v2_7b`/`v2_13b` and the default value is `v1_13b`. @@ -193,26 +193,23 @@ mpirun -n 2 --allow-run-as-root \ ```bash # Run summarization using the Baichuan V1 13B model in FP16. -python summarize.py --model_version v1_13b \ - --test_trt_llm \ - --hf_model_location baichuan-inc/Baichuan-13B-Chat \ - --data_type fp16 \ - --engine_dir ./tmp/baichuan_v1_13b/trt_engines/fp16/1-gpu/ +python ../summarize.py --test_trt_llm \ + --hf_model_dir baichuan-inc/Baichuan-13B-Chat \ + --data_type fp16 \ + --engine_dir ./tmp/baichuan_v1_13b/trt_engines/fp16/1-gpu/ # Run summarization using the Baichuan V1 13B model quantized to INT8. -python summarize.py --model_version v1_13b \ - --test_trt_llm \ - --hf_model_location baichuan-inc/Baichuan-13B-Chat \ - --data_type fp16 \ - --engine_dir ./tmp/baichuan_v1_13b/trt_engines/int8_weight_only/1-gpu/ +python ../summarize.py --test_trt_llm \ + --hf_model_dir baichuan-inc/Baichuan-13B-Chat \ + --data_type fp16 \ + --engine_dir ./tmp/baichuan_v1_13b/trt_engines/int8_weight_only/1-gpu/ # Run summarization using the Baichuan V1 13B model in FP16 using two GPUs. mpirun -n 2 --allow-run-as-root \ - python summarize.py --model_version v1_13b \ - --test_trt_llm \ - --hf_model_location baichuan-inc/Baichuan-13B-Chat \ - --data_type fp16 \ - --engine_dir ./tmp/baichuan_v1_13b/trt_engines/fp16/2-gpu/ + python ../summarize.py --test_trt_llm \ + --hf_model_dir baichuan-inc/Baichuan-13B-Chat \ + --data_type fp16 \ + --engine_dir ./tmp/baichuan_v1_13b/trt_engines/fp16/2-gpu/ ``` ### Known Issues diff --git a/examples/baichuan/build.py b/examples/baichuan/build.py index 4e2074b75025..ad863139c674 100644 --- a/examples/baichuan/build.py +++ b/examples/baichuan/build.py @@ -536,6 +536,7 @@ def build(rank, args): max_position_embeddings=args.n_positions, max_batch_size=args.max_batch_size, max_input_len=args.max_input_len, + max_beam_width=args.max_beam_width, max_output_len=args.max_output_len, max_num_tokens=args.max_num_tokens, int8=int8_trt_flag, diff --git a/examples/baichuan/requirements.txt b/examples/baichuan/requirements.txt index 553812d05dc7..8311dc5e4dc5 100644 --- a/examples/baichuan/requirements.txt +++ b/examples/baichuan/requirements.txt @@ -1,4 +1,5 @@ datasets~=2.14.5 +evaluate~=0.4.1 rouge_score~=0.1.2 sentencepiece~=0.1.99 cpm-kernels~=1.0.11 diff --git a/examples/baichuan/summarize.py b/examples/baichuan/summarize.py deleted file mode 100644 index 723048aae9c5..000000000000 --- a/examples/baichuan/summarize.py +++ /dev/null @@ -1,401 +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 copy -import json -import os - -import numpy as np -import torch -from datasets import load_dataset, load_metric -from transformers import AutoModelForCausalLM, AutoTokenizer - -import tensorrt_llm -import tensorrt_llm.profiler as profiler -from tensorrt_llm.logger import logger -from tensorrt_llm.quantization import QuantMode - -from build import get_engine_name # isort:skip - - -def TRTBaichuan(args, config): - dtype = config['builder_config']['precision'] - world_size = config['builder_config']['tensor_parallel'] - 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['builder_config']['num_heads'] // world_size - hidden_size = config['builder_config']['hidden_size'] // world_size - vocab_size = config['builder_config']['vocab_size'] - num_layers = config['builder_config']['num_layers'] - use_gpt_attention_plugin = bool( - config['plugin_config']['gpt_attention_plugin']) - remove_input_padding = config['plugin_config']['remove_input_padding'] - paged_kv_cache = config['plugin_config']['paged_kv_cache'] - tokens_per_block = config['plugin_config']['tokens_per_block'] - quant_mode = QuantMode(config['builder_config']['quant_mode']) - - model_config = tensorrt_llm.runtime.ModelConfig( - vocab_size=vocab_size, - num_layers=num_layers, - num_heads=num_heads, - num_kv_heads=num_heads, - hidden_size=hidden_size, - gpt_attention_plugin=use_gpt_attention_plugin, - tokens_per_block=tokens_per_block, - remove_input_padding=remove_input_padding, - paged_kv_cache=paged_kv_cache, - dtype=dtype, - quant_mode=quant_mode) - - runtime_rank = tensorrt_llm.mpi_rank() - runtime_mapping = tensorrt_llm.Mapping(world_size, - runtime_rank, - tp_size=world_size) - torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node) - - engine_name = get_engine_name('baichuan', dtype, world_size, runtime_rank) - serialize_path = os.path.join(args.engine_dir, engine_name) - - tensorrt_llm.logger.set_level(args.log_level) - - profiler.start('load tensorrt_llm engine') - with open(serialize_path, 'rb') as f: - engine_buffer = f.read() - decoder = tensorrt_llm.runtime.GenerationSession(model_config, - engine_buffer, - runtime_mapping) - profiler.stop('load tensorrt_llm engine') - tensorrt_llm.logger.info( - f'Load engine takes: {profiler.elapsed_time_in_sec("load tensorrt_llm engine")} sec' - ) - return decoder - - -def main(args): - 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 - hf_model_location = args.hf_model_location - profiler.start('load tokenizer') - tokenizer = AutoTokenizer.from_pretrained(hf_model_location, - use_fast=False, - trust_remote_code=True) - profiler.stop('load tokenizer') - tensorrt_llm.logger.info( - f'Load tokenizer takes: {profiler.elapsed_time_in_sec("load tokenizer")} sec' - ) - tokenizer.pad_token = tokenizer.eos_token - - dataset_cnn = load_dataset("ccdv/cnn_dailymail", - '3.0.0', - cache_dir=args.dataset_path) - - max_batch_size = args.batch_size - - # runtime parameters - # repetition_penalty = 1 - top_k = args.top_k - output_len = 100 - test_token_num = 923 - # top_p = 0.0 - # random_seed = 5 - temperature = 1 - num_beams = args.num_beams - - pad_id = tokenizer.encode(tokenizer.pad_token, add_special_tokens=False)[0] - end_id = tokenizer.encode(tokenizer.eos_token, add_special_tokens=False)[0] - - if test_trt_llm: - config_path = os.path.join(args.engine_dir, 'config.json') - with open(config_path, 'r') as f: - config = json.load(f) - - tensorrt_llm_baichuan = TRTBaichuan(args, config) - - if test_hf: - profiler.start('load HF model') - model = AutoModelForCausalLM.from_pretrained(hf_model_location, - trust_remote_code=True) - profiler.stop('load HF model') - tensorrt_llm.logger.info( - f'Load HF model takes: {profiler.elapsed_time_in_sec("load HF model")} sec' - ) - if args.data_type == 'fp16': - model.half() - model.cuda() - - def summarize_tensorrt_llm(datapoint): - batch_size = len(datapoint['article']) - - line = copy.copy(datapoint['article']) - line_encoded = [] - input_lengths = [] - for i in range(batch_size): - line[i] = line[i] + ' TL;DR: ' - - line[i] = line[i].strip() - line[i] = line[i].replace(" n't", "n't") - - input_id = tokenizer.encode(line[i], - return_tensors='pt').type(torch.int32) - input_id = input_id[:, -test_token_num:] - - line_encoded.append(input_id) - input_lengths.append(input_id.shape[-1]) - - # do padding, should move outside the profiling to prevent the overhead - max_length = max(input_lengths) - if tensorrt_llm_baichuan.remove_input_padding: - line_encoded = [ - torch.tensor(t, dtype=torch.int32).cuda() for t in line_encoded - ] - else: - # do padding, should move outside the profiling to prevent the overhead - for i in range(batch_size): - pad_size = max_length - input_lengths[i] - - pad = torch.ones([1, pad_size]).type(torch.int32) * pad_id - line_encoded[i] = torch.cat( - [torch.tensor(line_encoded[i], dtype=torch.int32), pad], - axis=-1) - - line_encoded = torch.cat(line_encoded, axis=0).cuda() - input_lengths = torch.tensor(input_lengths, - dtype=torch.int32).cuda() - - sampling_config = tensorrt_llm.runtime.SamplingConfig( - end_id=end_id, pad_id=pad_id, top_k=top_k, num_beams=num_beams) - - with torch.no_grad(): - tensorrt_llm_baichuan.setup( - batch_size, - max_context_length=max_length, - max_new_tokens=output_len, - beam_width=num_beams, - max_kv_cache_length=args.max_kv_cache_len) - if tensorrt_llm_baichuan.remove_input_padding: - output_ids = tensorrt_llm_baichuan.decode_batch( - line_encoded, sampling_config) - else: - output_ids = tensorrt_llm_baichuan.decode( - line_encoded, - input_lengths, - sampling_config, - ) - - torch.cuda.synchronize() - - # Extract a list of tensors of shape beam_width x output_ids. - if tensorrt_llm_baichuan.mapping.is_first_pp_rank(): - output_beams_list = [ - tokenizer.batch_decode(output_ids[batch_idx, :, - input_lengths[batch_idx]:], - skip_special_tokens=True) - for batch_idx in range(batch_size) - ] - return output_beams_list, output_ids[:, :, max_length:].tolist() - return [], [] - - def summarize_hf(datapoint): - batch_size = len(datapoint['article']) - 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}" - ) - - line = copy.copy(datapoint['article']) - for i in range(batch_size): - line[i] = line[i] + ' TL;DR: ' - - line[i] = line[i].strip() - line[i] = line[i].replace(" n't", "n't") - - line_encoded = tokenizer(line, - return_tensors='pt', - padding=True, - truncation=True)["input_ids"].type(torch.int64) - - line_encoded = line_encoded[:, -test_token_num:] - line_encoded = line_encoded.cuda() - - with torch.no_grad(): - output = model.generate(line_encoded, - max_new_tokens=output_len, - top_k=top_k, - temperature=temperature, - eos_token_id=tokenizer.eos_token_id, - pad_token_id=tokenizer.pad_token_id, - num_beams=num_beams, - num_return_sequences=num_beams, - early_stopping=True) - - tokens_list = output[:, len(line_encoded[0]):].tolist() - output = output.reshape([batch_size, num_beams, -1]) - output_lines_list = [ - tokenizer.batch_decode(output[:, i, len(line_encoded[0]):], - skip_special_tokens=True) - for i in range(num_beams) - ] - - return output_lines_list, tokens_list - - if test_trt_llm: - datapoint = dataset_cnn['test'][0:1] - summary, _ = summarize_tensorrt_llm(datapoint) - if runtime_rank == 0: - logger.info( - "---------------------------------------------------------") - logger.info("TensorRT-LLM Generated : ") - logger.info(f" Article : {datapoint['article']}") - logger.info(f"\n Highlights : {datapoint['highlights']}") - logger.info(f"\n Summary : {summary}") - logger.info( - "---------------------------------------------------------") - - if test_hf: - datapoint = dataset_cnn['test'][0:1] - summary, _ = summarize_hf(datapoint) - logger.info("---------------------------------------------------------") - logger.info("HF Generated : ") - logger.info(f" Article : {datapoint['article']}") - logger.info(f"\n Highlights : {datapoint['highlights']}") - logger.info(f"\n Summary : {summary}") - logger.info("---------------------------------------------------------") - - metric_tensorrt_llm = [load_metric("rouge") for _ in range(num_beams)] - metric_hf = [load_metric("rouge") for _ in range(num_beams)] - for i in range(num_beams): - metric_tensorrt_llm[i].seed = 0 - metric_hf[i].seed = 0 - - ite_count = 0 - data_point_idx = 0 - while (data_point_idx < len(dataset_cnn['test'])) 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_cnn['test'][data_point_idx:(data_point_idx + - max_batch_size)] - - if test_trt_llm: - profiler.start('tensorrt_llm') - summary_tensorrt_llm, tokens_tensorrt_llm = summarize_tensorrt_llm( - datapoint) - profiler.stop('tensorrt_llm') - - if test_hf: - profiler.start('hf') - summary_hf, tokens_hf = summarize_hf(datapoint) - profiler.stop('hf') - - if runtime_rank == 0: - if test_trt_llm: - for batch_idx in range(len(summary_tensorrt_llm)): - for beam_idx in range(num_beams): - metric_tensorrt_llm[beam_idx].add_batch( - predictions=[ - summary_tensorrt_llm[batch_idx][beam_idx] - ], - references=[datapoint['highlights'][batch_idx]]) - if test_hf: - for beam_idx in range(num_beams): - for batch_idx in range(len(summary_hf[beam_idx])): - metric_hf[beam_idx].add_batch( - predictions=[summary_hf[beam_idx][batch_idx]], - references=[datapoint['highlights'][batch_idx]]) - - logger.debug('-' * 100) - logger.debug(f"Article : {datapoint['article']}") - if test_trt_llm: - logger.debug(f'TensorRT-LLM Summary: {summary_tensorrt_llm}') - if test_hf: - logger.debug(f'HF Summary: {summary_hf}') - logger.debug(f"highlights : {datapoint['highlights']}") - - data_point_idx += max_batch_size - ite_count += 1 - - if runtime_rank == 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)' - ) - for beam_idx in range(num_beams): - logger.info(f"TensorRT-LLM beam {beam_idx} result") - 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].mid[2]*100}' - ) - - if args.check_accuracy and beam_idx == 0: - assert computed_metrics_tensorrt_llm['rouge1'].mid[ - 2] * 100 > args.tensorrt_llm_rouge1_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)' - ) - for beam_idx in range(num_beams): - logger.info(f"HF beam {beam_idx} result") - computed_metrics_hf = metric_hf[beam_idx].compute() - for key in computed_metrics_hf.keys(): - logger.info( - f' {key} : {computed_metrics_hf[key].mid[2]*100}') - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--model_version', - type=str, - default='v1_13b', - choices=['v1_7b', 'v1_13b', 'v2_7b', 'v2_13b']) - parser.add_argument('--hf_model_location', - type=str, - default='baichuan-inc/Baichuan-13B-Chat') - parser.add_argument('--test_hf', action='store_true') - parser.add_argument('--test_trt_llm', action='store_true') - parser.add_argument('--data_type', - type=str, - choices=['fp32', 'fp16'], - default='fp16') - parser.add_argument('--dataset_path', type=str, default='') - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument('--engine_dir', type=str, default='baichuan_outputs') - parser.add_argument('--batch_size', type=int, default=1) - parser.add_argument('--max_ite', type=int, default=20) - parser.add_argument('--max_kv_cache_len', - type=int, - default=None, - help='The max kv cache length. \ - If the final sequence length exceeds the kv cache length, we will enable cyclic kv cache. \ - If it is set to None, we will use the max sequence length.') - parser.add_argument('--check_accuracy', action='store_true') - parser.add_argument('--tensorrt_llm_rouge1_threshold', - type=float, - default=15.0) - parser.add_argument('--num_beams', type=int, default=1) - parser.add_argument('--top_k', type=int, default=1) - - args = parser.parse_args() - - main(args) diff --git a/examples/blip2/run.py b/examples/blip2/run.py index 192d303b3eb1..d238d247b702 100644 --- a/examples/blip2/run.py +++ b/examples/blip2/run.py @@ -42,7 +42,8 @@ def TRTOPT(args, config): vocab_size = config['builder_config']['vocab_size'] num_layers = config['builder_config']['num_layers'] remove_input_padding = config['plugin_config']['remove_input_padding'] - use_prompt_tuning = config['builder_config']['use_prompt_tuning'] + max_prompt_embedding_table_size = config['builder_config'].get( + 'max_prompt_embedding_table_size', 0) model_config = tensorrt_llm.runtime.ModelConfig( vocab_size=vocab_size, @@ -52,7 +53,7 @@ def TRTOPT(args, config): hidden_size=hidden_size, gpt_attention_plugin=use_gpt_attention_plugin, remove_input_padding=remove_input_padding, - use_prompt_tuning=use_prompt_tuning, + max_prompt_embedding_table_size=max_prompt_embedding_table_size, dtype=dtype) runtime_rank = tensorrt_llm.mpi_rank() diff --git a/examples/bloom/README.md b/examples/bloom/README.md index 43f5c7798cd2..02dfbc61e196 100644 --- a/examples/bloom/README.md +++ b/examples/bloom/README.md @@ -4,11 +4,11 @@ This document shows how to build and run a BLOOM model in TensorRT-LLM on both s ## Overview -The TensorRT-LLM BLOOM implementation can be found in [tensorrt_llm/models/bloom/model.py](../../tensorrt_llm/models/bloom/model.py). The TensorRT-LLM BLOOM example code is located in [`examples/bloom`](./). There are three main files in that folder:: +The TensorRT-LLM BLOOM implementation can be found in [tensorrt_llm/models/bloom/model.py](../../tensorrt_llm/models/bloom/model.py). The TensorRT-LLM BLOOM example code is located in [`examples/bloom`](./). There are three main files: * [`build.py`](./build.py) to build the [TensorRT](https://developer.nvidia.com/tensorrt) engine(s) needed to run the BLOOM model, * [`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/cnn_dailymail) dataset using the model. + * and a shared [`../summarize.py`](../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/cnn_dailymail) dataset using the model. ## Support Matrix * FP16 @@ -174,25 +174,25 @@ Note we use `--bin_model_dir` instead of `--model_dir` since SmoothQuant model n ### 4. Run ```bash -python summarize.py --test_trt_llm \ - --hf_model_location ./bloom/560M/ \ - --data_type fp16 \ - --engine_dir ./bloom/560M/trt_engines/fp16/1-gpu/ +python ../summarize.py --test_trt_llm \ + --hf_model_dir ./bloom/560M/ \ + --data_type fp16 \ + --engine_dir ./bloom/560M/trt_engines/fp16/1-gpu/ -python summarize.py --test_trt_llm \ - --hf_model_location ./bloom/560M/ \ - --data_type fp16 \ - --engine_dir ./bloom/560M/trt_engines/int8_weight_only/1-gpu/ +python ../summarize.py --test_trt_llm \ + --hf_model_dir ./bloom/560M/ \ + --data_type fp16 \ + --engine_dir ./bloom/560M/trt_engines/int8_weight_only/1-gpu/ mpirun -n 2 --allow-run-as-root \ - python summarize.py --test_trt_llm \ - --hf_model_location ./bloom/560M/ \ - --data_type fp16 \ - --engine_dir ./bloom/560M/trt_engines/fp16/2-gpu/ + python ../summarize.py --test_trt_llm \ + --hf_model_dir ./bloom/560M/ \ + --data_type fp16 \ + --engine_dir ./bloom/560M/trt_engines/fp16/2-gpu/ mpirun -n 8 --allow-run-as-root \ - python summarize.py --test_trt_llm \ - --hf_model_location ./bloom/176B/ \ - --data_type fp16 \ - --engine_dir ./bloom/176B/trt_engines/fp16/8-gpu/ + python ../summarize.py --test_trt_llm \ + --hf_model_dir ./bloom/176B/ \ + --data_type fp16 \ + --engine_dir ./bloom/176B/trt_engines/fp16/8-gpu/ ``` diff --git a/examples/bloom/build.py b/examples/bloom/build.py index 75ceef728f60..ad8e27b7d963 100644 --- a/examples/bloom/build.py +++ b/examples/bloom/build.py @@ -17,12 +17,15 @@ import time from pathlib import Path +import onnx import tensorrt as trt import torch import torch.multiprocessing as mp +from onnx import TensorProto, helper from transformers import BloomConfig, BloomForCausalLM import tensorrt_llm +from tensorrt_llm import profiler from tensorrt_llm._utils import str_dtype_to_trt from tensorrt_llm.builder import Builder from tensorrt_llm.logger import logger @@ -32,13 +35,13 @@ from tensorrt_llm.plugin.plugin import ContextFMHAType from tensorrt_llm.quantization import QuantMode -from weight import load_from_hf_bloom, load_from_bin, parse_config, check_embedding_share # isort:skip +# isort: off +from weight import (check_embedding_share, load_from_bin, load_from_hf_bloom, + load_from_hf_checkpoint, parse_config) -MODEL_NAME = "bloom" +# isort: on -import onnx -import tensorrt as trt -from onnx import TensorProto, helper +MODEL_NAME = "bloom" def trt_dtype_to_onnx(dtype): @@ -178,6 +181,9 @@ def parse_arguments(): ) parser.add_argument('--parallel_build', default=False, action='store_true') parser.add_argument('--visualize', default=False, action='store_true') + parser.add_argument('--load_by_shard', + action='store_true', + help='Load a pretrained model shard-by-shard.') parser.add_argument('--enable_debug_output', default=False, action='store_true') @@ -325,6 +331,8 @@ def build_rank_engine(builder: Builder, ''' kv_dtype = str_dtype_to_trt(args.dtype) + profiler.print_memory_usage(f'Rank {rank} Engine build starts') + # Share_embedding_table can be set True only when: # 1) the weight for lm_head() does not exist while other weights exist # 2) For multiple-processes, use_parallel_embedding=True and embedding_sharding_dim == 0. @@ -366,20 +374,32 @@ def build_rank_engine(builder: Builder, if args.model_dir is not None: logger.info(f'Loading HF BLOOM ... from {args.model_dir}') tik = time.time() - hf_bloom = BloomForCausalLM.from_pretrained(args.model_dir, - torch_dtype="auto") + if not args.load_by_shard: + hf_bloom = BloomForCausalLM.from_pretrained(args.model_dir, + torch_dtype="auto") + print(hf_bloom) + load_from_hf_bloom( + tensorrt_llm_bloom, + hf_bloom, + rank, + args.world_size, + fp16=(args.dtype == 'float16'), + use_parallel_embedding=args.use_parallel_embedding, + sharding_dim=args.embedding_sharding_dim, + share_embedding_table=share_embedding_table) + del hf_bloom + else: + load_from_hf_checkpoint( + tensorrt_llm_bloom, + model_dir=args.model_dir, + dtype=args.dtype, + use_parallel_embedding=args.use_parallel_embedding, + sharding_dim=args.embedding_sharding_dim, + share_embedding_table=share_embedding_table) tok = time.time() t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) logger.info(f'HF BLOOM loaded. Total time: {t}') - print(hf_bloom) - load_from_hf_bloom(tensorrt_llm_bloom, - hf_bloom, - rank, - args.world_size, - fp16=(args.dtype == 'float16'), - use_parallel_embedding=args.use_parallel_embedding, - sharding_dim=args.embedding_sharding_dim, - share_embedding_table=share_embedding_table) + elif args.bin_model_dir is not None: load_from_bin(tensorrt_llm_bloom, args.bin_model_dir, @@ -389,6 +409,7 @@ def build_rank_engine(builder: Builder, use_parallel_embedding=args.use_parallel_embedding, sharding_dim=args.embedding_sharding_dim, share_embedding_table=share_embedding_table) + profiler.print_memory_usage(f'Rank {rank} model weight loaded.') # Module -> Network network = builder.create_network() @@ -490,6 +511,7 @@ def build(rank, args): vocab_size=args.vocab_size, max_position_embeddings=args.n_positions, max_batch_size=args.max_batch_size, + max_beam_width=args.max_beam_width, max_input_len=args.max_input_len, max_output_len=args.max_output_len, int8=int8_trt_flag, @@ -509,6 +531,7 @@ def build(rank, args): serialize_engine(engine, os.path.join(args.output_dir, engine_name)) del engine + profiler.print_memory_usage(f'Rank {cur_rank} Engine serialized') if rank == 0: ok = builder.save_timing_cache( diff --git a/examples/bloom/requirements.txt b/examples/bloom/requirements.txt index 4c61cfa8f1da..ba54c3ef5ac2 100644 --- a/examples/bloom/requirements.txt +++ b/examples/bloom/requirements.txt @@ -1,3 +1,4 @@ datasets~=2.14.5 +evaluate~=0.4.1 rouge_score~=0.1.2 sentencepiece~=0.1.99 diff --git a/examples/bloom/summarize.py b/examples/bloom/summarize.py deleted file mode 100644 index 0db3b859d74b..000000000000 --- a/examples/bloom/summarize.py +++ /dev/null @@ -1,377 +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. -# TODO Just a copy paste, needs work - -import argparse -import copy -import json -import os - -import numpy as np -import torch -from datasets import load_dataset, load_metric -from transformers import AutoModelForCausalLM, BloomTokenizerFast - -import tensorrt_llm -import tensorrt_llm.profiler as profiler -from tensorrt_llm.logger import logger - -from build import get_engine_name # isort:skip - - -def TRTBloom(args, config): - dtype = config['builder_config']['precision'] - world_size = config['builder_config']['tensor_parallel'] - assert world_size == tensorrt_llm.mpi_world_size(), \ - f'Engine world size ({world_size}) != Runtime world size ({tensorrt_llm.mpi_world_size()})' - - world_size = config['builder_config']['tensor_parallel'] - num_heads = config['builder_config']['num_heads'] // world_size - hidden_size = config['builder_config']['hidden_size'] // world_size - vocab_size = config['builder_config']['vocab_size'] - num_layers = config['builder_config']['num_layers'] - use_gpt_attention_plugin = bool( - config['plugin_config']['gpt_attention_plugin']) - - model_config = tensorrt_llm.runtime.ModelConfig( - vocab_size=vocab_size, - num_layers=num_layers, - num_heads=num_heads, - num_kv_heads=num_heads, - hidden_size=hidden_size, - gpt_attention_plugin=use_gpt_attention_plugin, - dtype=dtype) - - runtime_rank = tensorrt_llm.mpi_rank() - runtime_mapping = tensorrt_llm.Mapping(world_size, - runtime_rank, - tp_size=world_size) - torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node) - - engine_name = get_engine_name('bloom', dtype, world_size, runtime_rank) - serialize_path = os.path.join(args.engine_dir, engine_name) - - tensorrt_llm.logger.set_level(args.log_level) - - profiler.start('load tensorrt_llm engine') - with open(serialize_path, 'rb') as f: - engine_buffer = f.read() - decoder = tensorrt_llm.runtime.GenerationSession(model_config, - engine_buffer, - runtime_mapping) - profiler.stop('load tensorrt_llm engine') - tensorrt_llm.logger.info( - f'Load engine takes: {profiler.elapsed_time_in_sec("load tensorrt_llm engine")} sec' - ) - return decoder - - -def main(args): - 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 - hf_model_location = args.hf_model_location - profiler.start('load tokenizer') - tokenizer = BloomTokenizerFast.from_pretrained(hf_model_location, - padding_side='left') - profiler.stop('load tokenizer') - tensorrt_llm.logger.info( - f'Load tokenizer takes: {profiler.elapsed_time_in_sec("load tokenizer")} sec' - ) - tokenizer.pad_token = tokenizer.eos_token - - dataset_cnn = load_dataset("ccdv/cnn_dailymail", - '3.0.0', - cache_dir=args.dataset_path) - - max_batch_size = args.batch_size - - # runtime parameters - # repetition_penalty = 1 - top_k = args.top_k - output_len = 100 - test_token_num = 923 - # top_p = 0.0 - # random_seed = 5 - temperature = 1 - num_beams = args.num_beams - - pad_id = tokenizer.encode(tokenizer.pad_token, add_special_tokens=False)[0] - end_id = tokenizer.encode(tokenizer.eos_token, add_special_tokens=False)[0] - - if test_trt_llm: - config_path = os.path.join(args.engine_dir, 'config.json') - with open(config_path, 'r') as f: - config = json.load(f) - - tensorrt_llm_bloom = TRTBloom(args, config) - - if test_hf: - profiler.start('load HF model') - model = AutoModelForCausalLM.from_pretrained(hf_model_location) - profiler.stop('load HF model') - tensorrt_llm.logger.info( - f'Load HF model takes: {profiler.elapsed_time_in_sec("load HF model")} sec' - ) - if args.data_type == 'fp16': - model.half() - model.cuda() - - def summarize_tensorrt_llm(datapoint): - batch_size = len(datapoint['article']) - - line = copy.copy(datapoint['article']) - line_encoded = [] - input_lengths = [] - for i in range(batch_size): - line[i] = line[i] + ' TL;DR: ' - - line[i] = line[i].strip() - line[i] = line[i].replace(" n't", "n't") - - input_id = tokenizer.encode(line[i], - return_tensors='pt').type(torch.int32) - input_id = input_id[:, -test_token_num:] - - line_encoded.append(input_id) - input_lengths.append(input_id.shape[-1]) - - # do padding, should move outside the profiling to prevent the overhead - max_length = max(input_lengths) - for i in range(batch_size): - pad_size = max_length - input_lengths[i] - - pad = torch.ones([1, pad_size]).type(torch.int32) * pad_id - line_encoded[i] = torch.cat( - [torch.tensor(line_encoded[i], dtype=torch.int32), pad], - axis=-1) - - line_encoded = torch.cat(line_encoded, axis=0).cuda() - input_lengths = torch.tensor(input_lengths, dtype=torch.int32).cuda() - - sampling_config = tensorrt_llm.runtime.SamplingConfig( - end_id=end_id, pad_id=pad_id, top_k=top_k, num_beams=num_beams) - - with torch.no_grad(): - tensorrt_llm_bloom.setup(line_encoded.size(0), - max_context_length=line_encoded.size(1), - max_new_tokens=output_len, - beam_width=num_beams, - max_kv_cache_length=args.max_kv_cache_len) - - output_ids = tensorrt_llm_bloom.decode( - line_encoded, - input_lengths, - sampling_config, - ) - - torch.cuda.synchronize() - - # Extract a list of tensors of shape beam_width x output_ids. - if tensorrt_llm_bloom.mapping.is_first_pp_rank(): - output_beams_list = [ - tokenizer.batch_decode(output_ids[batch_idx, :, - input_lengths[batch_idx]:], - skip_special_tokens=True) - for batch_idx in range(batch_size) - ] - return output_beams_list, output_ids[:, :, max_length:].tolist() - return [], [] - - def summarize_hf(datapoint): - batch_size = len(datapoint['article']) - 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}" - ) - - line = copy.copy(datapoint['article']) - for i in range(batch_size): - line[i] = line[i] + ' TL;DR: ' - - line[i] = line[i].strip() - line[i] = line[i].replace(" n't", "n't") - - line_encoded = tokenizer(line, - return_tensors='pt', - padding=True, - truncation=True)["input_ids"].type(torch.int64) - - line_encoded = line_encoded[:, -test_token_num:] - line_encoded = line_encoded.cuda() - - with torch.no_grad(): - output = model.generate(line_encoded, - max_length=len(line_encoded[0]) + - output_len, - top_k=top_k, - temperature=temperature, - eos_token_id=tokenizer.eos_token_id, - pad_token_id=tokenizer.pad_token_id, - num_beams=num_beams, - num_return_sequences=num_beams, - early_stopping=True) - - tokens_list = output[:, len(line_encoded[0]):].tolist() - output = output.reshape([batch_size, num_beams, -1]) - output_lines_list = [ - tokenizer.batch_decode(output[:, i, len(line_encoded[0]):], - skip_special_tokens=True) - for i in range(num_beams) - ] - - return output_lines_list, tokens_list - - if test_trt_llm: - datapoint = dataset_cnn['test'][0:1] - summary, _ = summarize_tensorrt_llm(datapoint) - if runtime_rank == 0: - logger.info( - "---------------------------------------------------------") - logger.info("TensorRT-LLM Generated : ") - logger.info(f" Article : {datapoint['article']}") - logger.info(f"\n Highlights : {datapoint['highlights']}") - logger.info(f"\n Summary : {summary}") - logger.info( - "---------------------------------------------------------") - - if test_hf: - datapoint = dataset_cnn['test'][0:1] - summary, _ = summarize_hf(datapoint) - logger.info("---------------------------------------------------------") - logger.info("HF Generated : ") - logger.info(f" Article : {datapoint['article']}") - logger.info(f"\n Highlights : {datapoint['highlights']}") - logger.info(f"\n Summary : {summary}") - logger.info("---------------------------------------------------------") - - metric_tensorrt_llm = [load_metric("rouge") for _ in range(num_beams)] - metric_hf = [load_metric("rouge") for _ in range(num_beams)] - for i in range(num_beams): - metric_tensorrt_llm[i].seed = 0 - metric_hf[i].seed = 0 - - ite_count = 0 - data_point_idx = 0 - while (data_point_idx < len(dataset_cnn['test'])) 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_cnn['test'][data_point_idx:(data_point_idx + - max_batch_size)] - - if test_trt_llm: - profiler.start('tensorrt_llm') - summary_tensorrt_llm, tokens_tensorrt_llm = summarize_tensorrt_llm( - datapoint) - profiler.stop('tensorrt_llm') - - if test_hf: - profiler.start('hf') - summary_hf, tokens_hf = summarize_hf(datapoint) - profiler.stop('hf') - - if runtime_rank == 0: - if test_trt_llm: - for batch_idx in range(len(summary_tensorrt_llm)): - for beam_idx in range(num_beams): - metric_tensorrt_llm[beam_idx].add_batch( - predictions=[ - summary_tensorrt_llm[batch_idx][beam_idx] - ], - references=[datapoint['highlights'][batch_idx]]) - if test_hf: - for beam_idx in range(num_beams): - for batch_idx in range(len(summary_hf[beam_idx])): - metric_hf[beam_idx].add_batch( - predictions=[summary_hf[beam_idx][batch_idx]], - references=[datapoint['highlights'][batch_idx]]) - - logger.debug('-' * 100) - logger.debug(f"Article : {datapoint['article']}") - if test_trt_llm: - logger.debug(f'TensorRT-LLM Summary: {summary_tensorrt_llm}') - if test_hf: - logger.debug(f'HF Summary: {summary_hf}') - logger.debug(f"highlights : {datapoint['highlights']}") - - data_point_idx += max_batch_size - ite_count += 1 - - if runtime_rank == 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)' - ) - for beam_idx in range(num_beams): - logger.info(f"TensorRT-LLM beam {beam_idx} result") - 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].mid[2]*100}' - ) - - if args.check_accuracy and beam_idx == 0: - assert computed_metrics_tensorrt_llm['rouge1'].mid[ - 2] * 100 > args.tensorrt_llm_rouge1_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)' - ) - for beam_idx in range(num_beams): - logger.info(f"HF beam {beam_idx} result") - computed_metrics_hf = metric_hf[beam_idx].compute() - for key in computed_metrics_hf.keys(): - logger.info( - f' {key} : {computed_metrics_hf[key].mid[2]*100}') - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--hf_model_location', type=str, default='./bloom/560M') - parser.add_argument('--test_hf', action='store_true') - parser.add_argument('--test_trt_llm', action='store_true') - parser.add_argument('--data_type', - type=str, - choices=['fp32', 'fp16'], - default='fp16') - parser.add_argument('--dataset_path', type=str, default='') - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument('--engine_dir', type=str, default='bloom_outputs') - parser.add_argument('--batch_size', type=int, default=1) - parser.add_argument('--max_ite', type=int, default=20) - parser.add_argument('--max_kv_cache_len', - type=int, - default=None, - help='The max kv cache length. \ - If the final sequence length exceeds the kv cache length, we will enable cyclic kv cache. \ - If it is set to None, we will use the max sequence length.') - parser.add_argument('--check_accuracy', action='store_true') - parser.add_argument('--tensorrt_llm_rouge1_threshold', - type=float, - default=15.0) - parser.add_argument('--num_beams', type=int, default=1) - parser.add_argument('--top_k', type=int, default=1) - - args = parser.parse_args() - - main(args) diff --git a/examples/bloom/weight.py b/examples/bloom/weight.py index df5de3557587..2313fa7c2848 100644 --- a/examples/bloom/weight.py +++ b/examples/bloom/weight.py @@ -15,12 +15,14 @@ import configparser import time from pathlib import Path +from typing import Union import numpy as np import torch import tensorrt_llm from tensorrt_llm._utils import str_dtype_to_np +from tensorrt_llm.logger import logger from tensorrt_llm.models import BloomForCausalLM from tensorrt_llm.quantization import QuantMode @@ -195,8 +197,17 @@ def load_from_hf_bloom(tensorrt_llm_bloom, embed_w = get_weight(model_params, 'transformer.word_embeddings', dtype) if not share_embedding_table: + vocab_size = embed_w.shape[0] + lm_head_weight = embed_w.copy() + if vocab_size % tensor_parallel != 0: + # padding + vocab_size_padded = tensorrt_llm_bloom.lm_head.out_features * tensor_parallel + pad_width = vocab_size_padded - vocab_size + lm_head_weight = np.pad(lm_head_weight, ((pad_width, 0), (0, 0)), + 'constant', + constant_values=0) tensorrt_llm_bloom.lm_head.weight.value = split_matrix_tp( - embed_w.copy(), tensor_parallel, rank, dim=0) + lm_head_weight, tensor_parallel, rank, dim=0) if not use_parallel_embedding: tensorrt_llm_bloom.embedding.weight.value = embed_w @@ -220,6 +231,108 @@ def load_from_hf_bloom(tensorrt_llm_bloom, tensorrt_llm.logger.info(f'Weights loaded. Total time: {t}') +def load_from_hf_checkpoint( + tensorrt_llm_bloom: tensorrt_llm.models.BloomForCausalLM, + model_dir: Union[str, Path], + dtype: Union[str, torch.dtype] = torch.float32, + use_parallel_embedding: bool = False, + sharding_dim: int = 0, + share_embedding_table: bool = False, +): + tensorrt_llm.logger.info('Loading weights from HF BLOOM...') + tik = time.time() + + quant_mode = getattr(tensorrt_llm_bloom, 'quant_mode', QuantMode(0)) + mapping = tensorrt_llm_bloom.mapping + tp_size = mapping.tp_size + tp_rank = mapping.tp_rank + + if isinstance(dtype, str): + dtype = tensorrt_llm._utils.str_dtype_to_torch(dtype) + + def is_bias(_name): + return 'bias' in _name + + # Load examples/common/utils.py + import sys + sys.path.append(str(Path(__file__).parent.parent)) + from common import utils + + for model_file in utils.iterate_shard_files(model_dir, mapping.tp_rank): + logger.debug(f'Loading file {str(model_file)}...') + model_params = utils.load_state_dict(model_file, dtype=dtype) + for name, param in model_params.items(): + logger.debug(f'Converting weight {name}...') + i = utils.retrieved_layer_index_from_name(name) + layer = tensorrt_llm_bloom.layers[i] if i is not None else None + param = param.detach().cpu().numpy() + if 'self_attention.query_key_value' in name: + if not is_bias(name): + split_v = split_qkv_tp(tensorrt_llm_bloom, param, tp_size, + tp_rank) + set_layer_weight(layer.attention.qkv, split_v, quant_mode) + else: + layer.attention.qkv.bias.value = split_qkv_bias_tp( + tensorrt_llm_bloom, param, tp_size, tp_rank) + elif 'self_attention.dense' in name: + if not is_bias(name): + split_v = split_matrix_tp(param, tp_size, tp_rank, dim=1) + set_layer_weight(layer.attention.dense, split_v, quant_mode) + else: + layer.attention.dense.bias.value = param + + elif 'mlp.dense_h_to_4h' in name: + if not is_bias(name): + split_v = split_matrix_tp(param, tp_size, tp_rank, dim=0) + set_layer_weight(layer.mlp.fc, split_v, quant_mode) + else: + layer.mlp.fc.bias.value = split_matrix_tp(param, + tp_size, + tp_rank, + dim=0) + elif 'mlp.dense_4h_to_h' in name: + if not is_bias(name): + split_v = split_matrix_tp(param, tp_size, tp_rank, dim=1) + set_layer_weight(layer.mlp.proj, split_v, quant_mode) + else: + layer.mlp.proj.bias.value = param + elif 'input_layernorm' in name: + if not is_bias(name): + layer.input_layernorm.weight.value = param + else: + layer.input_layernorm.bias.value = param + elif 'post_attention_layernorm' in name: + if not is_bias(name): + layer.post_layernorm.weight.value = param + else: + layer.post_layernorm.bias.value = param + elif 'word_embeddings.' in name: + if not share_embedding_table: + tensorrt_llm_bloom.lm_head.weight.value = split_matrix_tp( + param, tp_size, tp_rank, dim=0) + if not use_parallel_embedding: + tensorrt_llm_bloom.embedding.weight.value = param + else: + assert tensorrt_llm_bloom._vocab_size % tp_size == 0 + tensorrt_llm_bloom.embedding.weight.value = split_matrix_tp( + param, tp_size, tp_rank, dim=sharding_dim) + elif 'word_embeddings_layernorm.' in name: + if not is_bias(name): + tensorrt_llm_bloom.ln_embed.weight.value = param + else: + tensorrt_llm_bloom.ln_embed.bias.value = param + elif 'ln_f.' in name: + if not is_bias(name): + tensorrt_llm_bloom.ln_f.weight.value = param + else: + tensorrt_llm_bloom.ln_f.bias.value = param + del model_params + + tok = time.time() + t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) + tensorrt_llm.logger.info(f'Weights loaded. Total time: {t}') + + def gen_suffix(rank, use_smooth_quant, quant_per_channel): suffix = f"{rank}.bin" if use_smooth_quant: diff --git a/examples/chatglm/.gitignore b/examples/chatglm/.gitignore index 979e236242a1..3d8bd4430896 100644 --- a/examples/chatglm/.gitignore +++ b/examples/chatglm/.gitignore @@ -1,6 +1,7 @@ __pycache__/ -chatglm*-6b/ -chatglm*-6b-32k/ -trtModel/ -dataset/ .vscode/ +awq/ +chatglm*_6b*/ +dataset/ +glm_10b/ +trtModel/ diff --git a/examples/chatglm/README.md b/examples/chatglm/README.md index 74042ce268bc..9360f5622bd5 100644 --- a/examples/chatglm/README.md +++ b/examples/chatglm/README.md @@ -1,24 +1,36 @@ # ChatGLM -This document explains how to build the [ChatGLM-6B](https://huggingface.co/THUDM/chatglm-6b), [ChatGLM2-6B](https://huggingface.co/THUDM/chatglm2-6b) and [ChatGLM3-6B](https://huggingface.co/THUDM/chatglm3-6b), [ChatGLM2-6B-32k](https://huggingface.co/THUDM/chatglm2-6b-32k), [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. +This document explains how to build the [ChatGLM-6B](https://huggingface.co/THUDM/chatglm-6b), [ChatGLM2-6B](https://huggingface.co/THUDM/chatglm2-6b), [ChatGLM2-6B-32k](https://huggingface.co/THUDM/chatglm2-6b-32k), [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. ## 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/chatglm`](./). There are 3 main files in that folder: +The TensorRT-LLM ChatGLM example code is located in [`examples/chatglm`](./). There are three main files: * [`build.py`](./build.py) to build the [TensorRT](https://developer.nvidia.com/tensorrt) engine(s) needed to run the ChatGLM model. * [`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/cnn_dailymail) dataset using the model. +* and a shared [`../summarize.py`](../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/cnn_dailymail) dataset using the model. ## Support Matrix -* FP16 -* Weight Only Quantization (int8 / int4) -* Paged KV cache -* Remove Input Padding -* Tensor Parallel -* Strongly Typed +| Model Name | FP16 | FMHA | WO | AWQ | SQ | TP | PP | Strongly Typed | C++ Runtime | benchmark | IFB | +| :--------------: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :------------: | :---------: | :-------: | :---: | +| chatglm_6b | Y | Y | Y | | | Y | | Y | Y | Y | | +| chatglm2_6b | Y | Y | Y | | | Y | | Y | Y | Y | | +| chatglm2-6b_32k | Y | Y | Y | | | Y | | Y | Y | Y | | +| chatglm3_6b | Y | Y | Y | | | Y | | Y | Y | Y | | +| chatglm3_6b_base | Y | Y | Y | | | Y | | Y | Y | Y | | +| chatglm3_6b_32k | Y | Y | Y | | | Y | | Y | Y | Y | | +| glm_10b | 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) +* AWQ: Activation Aware Weight Quantization +* SQ:Smooth Quantization +* TP: Tensor Parallel +* PP: Pipeline Parallel +* IFB: In-flight Batching (see introduction below) ## Usage @@ -33,11 +45,13 @@ 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 -git clone https://huggingface.co/THUDM/chatglm2-6b -git clone https://huggingface.co/THUDM/chatglm3-6b -git clone https://huggingface.co/THUDM/chatglm2-6b-32k -git clone https://huggingface.co/THUDM/chatglm3-6b-32k +git clone https://huggingface.co/THUDM/chatglm-6b chatglm_6b +git clone https://huggingface.co/THUDM/chatglm2-6b chatglm2_6b +git clone https://huggingface.co/THUDM/chatglm2-6b-32k chatglm2_6b_32k +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 +git clone https://huggingface.co/THUDM/glm-10b glm_10b ``` ### 2. Build TensorRT engine(s) @@ -48,45 +62,51 @@ git clone https://huggingface.co/THUDM/chatglm3-6b-32k * You can enable parallel builds to accelerate the engine building process if you have more than one GPU in your system (of the same model). * For parallel building, add the `--parallel_build` argument to the build command (this feature cannot take advantage of more than a single node). * The number of TensorRT engines depends on the number of GPUs that will be used to run inference. -* argument [--model_version/-m] is required, which can be one of "1", "2", "3", "2-32k" or "3-32k" for ChatGLM-6B, ChatGLM2-6B, ChatGLM3-6B, ChatGLM2-6B-32K or ChatGLM3-6B-32K respectively. +* argument [--model_name/-m] is required, which can be one of "chatglm_6b", "chatglm2_6b", "chatglm2_6b_32k", "chatglm3_6b", "chatglm3_6b_base", "chatglm3_6b_32k" or "glm-10b" (use "_" rather than "-") for ChatGLM-6B, ChatGLM2-6B, ChatGLM2-6B-32K ChatGLM3-6B, ChatGLM3-6B-Base, ChatGLM3-6B-32K or GLM-10B model respectively. #### Examples of build invocations ```bash # Build a default engine of ChatGLM3-6B on single GPU with FP16, GPT Attention plugin, Gemm plugin, RMS Normolization plugin -python3 build.py -m 3 +python3 build.py -m chatglm3_6b # Build a engine on single GPU with FMHA kernels (see introduction below), other configurations are the same as default example -python3 build.py -m 3 --enable_context_fmha # or --enable_context_fmha_fp32_acc +python3 build.py -m chatglm3_6b --enable_context_fmha # or --enable_context_fmha_fp32_acc # Build a engine on single GPU with int8/int4 Weight-Only quantization, other configurations are the same as default example -python3 build.py -m 3 --use_weight_only # or --use_weight_only --weight_only_precision int4 +python3 build.py -m chatglm3_6b --use_weight_only # or --use_weight_only --weight_only_precision int4 # Build a engine on single GPU with int8_kv_cache and remove_input_padding, other configurations are the same as default example -python3 build.py -m 3 --paged_kv_cache --remove_input_padding +python3 build.py -m chatglm3_6b --paged_kv_cache --remove_input_padding # Build a engine on two GPU, other configurations are the same as default example -python3 build.py -m 3 --world_size 2 +python3 build.py -m chatglm3_6b --world_size 2 -# Build a engine of ChatGLM-6B on single GPU, other configurations are the same as default example -python3 build.py -m 1 +# Build a engine of Chatglm-6B on single GPU, other configurations are the same as default example +python3 build.py -m chatglm_6b -# Build a engine of ChatGLM2-6B on single GPU, other configurations are the same as default example -python3 build.py -m 2 +# Build a engine of Chatglm2-6B on single GPU, other configurations are the same as default example +python3 build.py -m chatglm2_6b # Build a engine of ChatGLM2-6B-32k on single GPU, other configurations are the same as default example -python3 build.py -m 2-32k +python3 build.py -m chatglm2_6b-32k + +# Build a engine of ChatGLM3-6B-Base on single GPU, other configurations are the same as default example +python3 build.py -m chatglm3_6b_base # Build a engine of ChatGLM3-6B-32k on single GPU, other configurations are the same as default example -python3 build.py -m 3-32k +python3 build.py -m chatglm3_6b-32k + +# Build a engine of GLM-10B on single GPU, other configurations are the same as default example +python3 build.py -m glm_10b ``` #### Enabled plugins * Use `--use_gemm_plugin ` to configure GPT Attention plugin (default as float16) * Use `--use_gemm_plugin ` to configure GEMM normolization plugin (default as float16) -* Use `--use_layernorm_plugin ` (for ChatGLM-6B) to configure RMS normolization plugin (default as float16) -* Use `--use_rmsnorm_plugin ` (for ChatGLM2-6B and ChatGLM3-6B) to configure RMS normolization plugin (default as float16) +* Use `--use_layernorm_plugin ` (for ChatGLM-6B and GLM-10B models) to configure RMS normolization plugin (default as float16) +* Use `--use_rmsnorm_plugin ` (for ChatGLM2-6B\* and ChatGLM3-6B\* models) to configure RMS normolization plugin (default as float16) #### Fused MultiHead Attention (FMHA) @@ -102,7 +122,7 @@ python3 build.py -m 3-32k * Furthermore, use `--weight_only_precision int8` or `--weight_only_precision int4` to configure the data type of the weights. -#### In-flight batching and paged KV cache [TODO] +#### In-flight batching * The engine must be built accordingly if [in-flight batching in C++ runtime](../../docs/in_flight_batching.md) will be used. @@ -119,15 +139,15 @@ python3 build.py -m 3-32k #### Single node, single GPU ```bash -# Run the default engine of ChatGLM3-6B on single GPU, other model version is available if built. -python3 run.py -m 3 +# Run the default engine of ChatGLM3-6B on single GPU, other model name is available if built. +python3 run.py -m chatglm3_6b ``` #### Single node, multi GPU ```bash -# Run the Tensor Parallel 2 engine of ChatGLM3-6B on two GPU, other model version is available if built. -mpirun -n 2 python run.py -m 3 +# Run the Tensor Parallel 2 engine of ChatGLM3-6B on two GPU, other model name is available if built. +mpirun -n 2 python run.py -m chatglm3_6b ``` * `--allow-run-as-root` might be needed if using `mpirun` as root. @@ -135,8 +155,8 @@ mpirun -n 2 python run.py -m 3 #### Run comparison of performance and accuracy ```bash -# Run the summarization of ChatGLM3-6B task, other model version is available if built. -python3 summarize.py -m 3 +# Run the summarization of ChatGLM3-6B task, other model name is available if built. +python3 ../summarize.py -m chatglm3_6b ``` ## Benchmark diff --git a/examples/chatglm/build.py b/examples/chatglm/build.py index 414b2ea10ccb..865671b4674f 100644 --- a/examples/chatglm/build.py +++ b/examples/chatglm/build.py @@ -18,20 +18,27 @@ from pathlib import Path from typing import List +import onnx +import tensorrt as trt import torch import torch.multiprocessing as mp -import transformers +from onnx import TensorProto, helper from weight import load_from_hf import tensorrt_llm +from tensorrt_llm._utils import str_dtype_to_trt from tensorrt_llm.builder import Builder from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from tensorrt_llm.models import ChatGLMHeadModel, quantize_model from tensorrt_llm.network import net_guard from tensorrt_llm.plugin.plugin import ContextFMHAType +from tensorrt_llm.profiler import check_gpt_mem_usage from tensorrt_llm.quantization import QuantMode +from weight import get_scaling_factors # isort:skip +from weight import load_from_hf_checkpoint # isort:skip + def get_engine_name(model, dtype, tp_size, rank): return '{}_{}_tp{}_rank{}.engine'.format(model, dtype, tp_size, rank) @@ -46,6 +53,61 @@ def find_engines(dir: Path, return list(dir.glob(template)) +def trt_dtype_to_onnx(dtype): + if dtype == trt.float16: + return TensorProto.DataType.FLOAT16 + elif dtype == trt.float32: + return TensorProto.DataType.FLOAT + elif dtype == trt.int32: + return TensorProto.DataType.INT32 + else: + raise TypeError("%s is not supported" % dtype) + + +def to_onnx(network, path): + inputs = [] + for i in range(network.num_inputs): + network_input = network.get_input(i) + inputs.append( + helper.make_tensor_value_info( + network_input.name, trt_dtype_to_onnx(network_input.dtype), + list(network_input.shape))) + + outputs = [] + for i in range(network.num_outputs): + network_output = network.get_output(i) + outputs.append( + helper.make_tensor_value_info( + network_output.name, trt_dtype_to_onnx(network_output.dtype), + list(network_output.shape))) + + nodes = [] + for i in range(network.num_layers): + layer = network.get_layer(i) + layer_inputs = [] + for j in range(layer.num_inputs): + ipt = layer.get_input(j) + if ipt is not None: + layer_inputs.append(layer.get_input(j).name) + layer_outputs = [ + layer.get_output(j).name for j in range(layer.num_outputs) + ] + nodes.append( + helper.make_node(str(layer.type), + name=layer.name, + inputs=layer_inputs, + outputs=layer_outputs, + domain="com.nvidia")) + + onnx_model = helper.make_model(helper.make_graph(nodes, + 'attention', + inputs, + outputs, + initializer=None), + producer_name='NVIDIA') + onnx.save(onnx_model, path) + + def serialize_engine(engine, path): logger.info(f'Serializing engine to {path}...') tik = time.time() @@ -56,21 +118,46 @@ def serialize_engine(engine, path): logger.info(f'Engine serialized. Total time: {t}') +def truncate_input_output( + max_input_len, + max_output_len, + max_seq_length_from_config, + is_fixed_max_position_length=False, +): + max_seq_length = max_seq_length_from_config + if max_input_len >= max_seq_length_from_config: + print("Truncate max_input_len as %d" % (max_seq_length_from_config - 1)) + max_input_len = max_seq_length_from_config - 1 + max_output_len = 1 + elif max_input_len + max_output_len > max_seq_length_from_config: + print("Truncate max_output_len as %d" % + (max_seq_length_from_config - max_input_len)) + max_output_len = max_seq_length_from_config - max_input_len + elif not is_fixed_max_position_length: + max_seq_length = max_input_len + max_output_len + return max_input_len, max_output_len, max_seq_length + + def parse_arguments(args): parser = argparse.ArgumentParser() parser.add_argument( - '--model_version', + '--model_name', '-m', type=str, required=True, - choices=["1", "2", "3", "2-32k", "3-32k"], + choices=[ + "chatglm_6b", "chatglm2_6b", "chatglm2_6b_32k", "chatglm3_6b", + "chatglm3_6b_base", "chatglm3_6b_32k", "glm_10b" + ], help= - '1, 2, 3, 2-32k, 3-32k for ChatGLM-6B, ChatGLM2-6B, ChatGLM3-6B, ChatGLM2-32k and ChatGLM3-32k respectively' + 'the name of the model, use "_" rather than "-" to connect the name parts' ) parser.add_argument('--world_size', type=int, default=1, help='world size, only support tensor parallelism now') + parser.add_argument('--tp_size', type=int, default=1) + parser.add_argument('--pp_size', type=int, default=1) parser.add_argument('--model_dir', type=str, default=None) parser.add_argument('--dtype', type=str, @@ -123,7 +210,7 @@ def parse_arguments(args): default='float16', choices=['float32', 'float16', 'bfloat16', False], help= - "Activates layernorm plugin for ChatGLM-6B. You can specify the plugin dtype or leave blank to use the model dtype." + "Activates layernorm plugin for ChatGLM-6B / GLM-10B models. You can specify the plugin dtype or leave blank to use the model dtype." ) parser.add_argument( '--use_rmsnorm_plugin', @@ -133,7 +220,7 @@ def parse_arguments(args): default='float16', choices=['float32', 'float16', 'bfloat16', False], help= - "Activates rmsnorm plugin for ChatGLM2-6B / ChatGLM3-6B. You can specify the plugin dtype or leave blank to use the model dtype." + "Activates rmsnorm plugin for ChatGLM2-6B* / ChatGLM3-6B* models. You can specify the plugin dtype or leave blank to use the model dtype." ) parser.add_argument('--gather_all_token_logits', action='store_true', @@ -153,6 +240,13 @@ def parse_arguments(args): 'Split long kv sequence into multiple blocks (applied to generation MHA kernels). \ It is beneifical when batchxnum_heads cannot fully utilize GPU.' ) + parser.add_argument('--load_by_shard', + action='store_true', + help='Load a pretrained model shard-by-shard.') + parser.add_argument('--visualize', default=False, action='store_true') + parser.add_argument('--enable_debug_output', + default=False, + action='store_true') parser.add_argument('--gpus_per_node', type=int, default=8) parser.add_argument('--builder_opt', type=int, default=None) parser.add_argument( @@ -162,9 +256,23 @@ def parse_arguments(args): help= 'The path to save the serialized engine files, timing cache file and model configs' ) + parser.add_argument( + '--strongly_typed', + default=False, + action="store_true", + help= + 'This option is introduced with trt 9.1.0.1+ and will reduce the building time significantly for fp8.' + ) parser.add_argument('--remove_input_padding', default=False, action='store_true') + parser.add_argument( + '--paged_kv_cache', + action="store_true", + default=False, + help= + 'By default we use contiguous KV cache. By setting this flag you enable paged KV cache' + ) parser.add_argument( '--use_inflight_batching', action="store_true", @@ -226,13 +334,6 @@ def parse_arguments(args): default=None, help= 'Seed to use when initializing the random number generator for torch.') - parser.add_argument( - '--paged_kv_cache', - action="store_true", - default=False, - help= - 'By default we use contiguous KV cache. By setting this flag you enable paged KV cache' - ) parser.add_argument('--tokens_per_block', type=int, default=64, @@ -255,13 +356,6 @@ def parse_arguments(args): type=int, default=None, help='Define the max number of tokens supported by the engine') - parser.add_argument( - '--strongly_typed', - default=False, - action="store_true", - help= - 'This option is introduced with trt 9.1.0.1+ and will reduce the building time significantly for fp8.' - ) parser.add_argument( '--use_custom_all_reduce', action='store_true', @@ -283,58 +377,83 @@ def parse_arguments(args): ) setattr(args, plugin_arg, args.dtype) - if args.model_version == "1": - args.model_name = "chatglm-6b" - elif args.model_version in ["2", "3"]: - args.model_name = "chatglm%s-6b" % args.model_version - else: - args.model_name = "chatglm%s-6b-32k" % args.model_version.split("-")[0] if args.model_dir is None: args.model_dir = args.model_name with open(Path(args.model_dir) / "config.json", "r") as f: js = json.loads(f.read()) - if args.model_version == "1": + + if args.model_name in ["chatglm_6b", "glm_10b"]: assert args.max_input_len < js["max_sequence_length"] - args.apply_query_key_layer_scaling = False # always False in TRT-LLM - args.eos_token_id = js["eos_token_id"] - args.hidden_size = js["hidden_size"] - args.multi_block_mode = False - args.norm_epsilon = js["layernorm_epsilon"] - args.num_heads = js["num_attention_heads"] - args.num_layers = js["num_layers"] - args.pad_token_id = js["pad_token_id"] - args.use_cache = js["use_cache"] - if args.model_version == "1": + + if args.model_name in ["chatglm_6b"]: args.ffn_hidden_size = js["inner_hidden_size"] + args.hidden_size = js["hidden_size"] + args.norm_epsilon = js["layernorm_epsilon"] + args.num_heads = js["num_attention_heads"] + args.num_layers = js["num_layers"] + args.vocab_size = js["vocab_size"] + args.max_input_len, args.max_output_len, args.max_seq_length = truncate_input_output( + args.max_input_len, args.max_output_len, js["max_sequence_length"]) + args.apply_query_key_layer_scaling = False args.hidden_act = 'gelu' - args.linear_bias = True # always True in ChatGLM-6B - args.max_seq_length = min(args.max_input_len + args.max_output_len, - js["max_sequence_length"]) - args.multi_query_mode = False # always False in ChatGLM-6B + args.linear_bias = True + args.multi_block_mode = False + args.multi_query_mode = False args.num_kv_heads = js["num_attention_heads"] - args.qkv_bias = True # always True in ChatGLM-6B + args.qkv_bias = True + args.use_cache = js["use_cache"] + elif args.model_name in ["glm_10b"]: + args.hidden_size = js["hidden_size"] + args.num_attention_heads = js["num_attention_heads"] + args.num_heads = js["num_attention_heads"] + args.num_layers = js["num_layers"] args.vocab_size = js["vocab_size"] - else: - #args.kv_channels = js["kv_channels"] # useless + args.max_input_len, args.max_output_len, args.max_seq_length = truncate_input_output( + args.max_input_len, args.max_output_len, js["max_sequence_length"], + True) + args.apply_query_key_layer_scaling = False + args.apply_residual_connection_post_layernorm = False + args.ffn_hidden_size = 4 * args.hidden_size + args.hidden_act = 'gelu' + args.linear_bias = True + args.multi_block_mode = False + args.multi_query_mode = False + args.norm_epsilon = 1.0e-5 + args.num_kv_heads = js["num_attention_heads"] + args.qkv_bias = True + args.use_cache = True + elif args.model_name in [ + "chatglm2_6b", "chatglm2_6b_32k", "chatglm3_6b", "chatglm3_6b_base", + "chatglm3_6b_32k" + ]: + args.apply_query_key_layer_scaling = False args.apply_residual_connection_post_layernorm = js[ "apply_residual_connection_post_layernorm"] args.ffn_hidden_size = js["ffn_hidden_size"] - args.hidden_act = 'swiglu' + args.hidden_size = js["hidden_size"] args.linear_bias = js["add_bias_linear"] - args.max_seq_length = min(args.max_input_len + args.max_output_len, - js["seq_length"]) args.multi_query_mode = js["multi_query_attention"] + args.norm_epsilon = js["layernorm_epsilon"] + args.num_heads = js["num_attention_heads"] args.num_kv_heads = js["multi_query_group_num"] + args.num_layers = js["num_layers"] args.qkv_bias = js["add_qkv_bias"] args.rmsnorm = js["rmsnorm"] + args.use_cache = js["use_cache"] args.vocab_size = js["padded_vocab_size"] + args.max_seq_length = min(args.max_input_len + args.max_output_len, + js["seq_length"]) + if args.model_name in ["chatglm2_6b_32k", "chatglm3_6b_32k"]: + args.rotary_embedding_scaling = js["rope_ratio"] + args.hidden_act = 'swiglu' + args.multi_block_mode = False if args.use_inflight_batching: if not args.use_gpt_attention_plugin: args.use_gpt_attention_plugin = 'float16' logger.info( - f"Using GPT attention plugin for inflight batching mode. Setting to default '{args.use_gpt_attention_plugin}'" - ) + f"Using GPT attention plugin for inflight batching mode. " + f"Setting to default '{args.use_gpt_attention_plugin}'") if not args.remove_input_padding: args.remove_input_padding = True logger.info( @@ -374,37 +493,57 @@ def parse_arguments(args): return args -def build_rank_engine(builder: Builder, - builder_config: tensorrt_llm.builder.BuilderConfig, - engine_name, rank, args): +def build_rank_engine( + builder: Builder, + builder_config: tensorrt_llm.builder.BuilderConfig, + engine_name: str, + rank: int, + args: argparse.Namespace, +) -> trt.IHostMemory: ''' @brief: Build the engine on the given rank. @param rank: The rank to build the engine. @param args: The cmd line arguments. @return: The built engine. ''' - # Initialize Module args.mapping = Mapping( world_size=args.world_size, rank=rank, tp_size=args.world_size, ) + assert args.num_layers % args.pp_size == 0, \ + f"num_layers {args.n_layer} must be a multiple of pipeline "\ + f"parallelism size {args.pp_size}" trtllm_model = ChatGLMHeadModel(args=args) if args.use_smooth_quant or args.use_weight_only: trtllm_model = quantize_model(trtllm_model, args.quant_mode) - if args.model_dir is not None: - hf_model = transformers.AutoModel.from_pretrained( - args.model_dir, trust_remote_code=True).cpu() + if args.enable_fp8 or args.fp8_kv_cache: + logger.info(f'Loading scaling factors from ' + f'{args.quantized_fp8_model_path}') + quant_scales = get_scaling_factors(args.quantized_fp8_model_path, + num_layers=args.n_layer, + quant_mode=args.quant_mode) + tensorrt_llm_falcon = quantize_model(tensorrt_llm_falcon, + quant_mode=args.quant_mode, + quant_scales=quant_scales) + if not args.load_by_shard: trtllm_model = load_from_hf( trtllm_model, - hf_model, + args.model_dir, + mapping=args.mapping, + dtype=args.dtype, + model_name=args.model_name, + ) + else: + trtllm_model = load_from_hf_checkpoint( + trtllm_model, + args.model_dir, mapping=args.mapping, dtype=args.dtype, - model_version=args.model_version, + model_name=args.model_name, ) - del hf_model # Module -> Network network = builder.create_network() @@ -427,6 +566,10 @@ def build_rank_engine(builder: Builder, ContextFMHAType.enabled_with_fp32_acc) if args.multi_block_mode: network.plugin_config.enable_mmha_multi_block_mode() + + if args.world_size > 1: + network.plugin_config.set_nccl_plugin(args.dtype, + args.use_custom_all_reduce) if args.remove_input_padding: network.plugin_config.enable_remove_input_padding() if args.paged_kv_cache: @@ -437,7 +580,6 @@ def build_rank_engine(builder: Builder, network.plugin_config.set_smooth_quant_gemm_plugin(dtype=args.dtype) network.plugin_config.set_layernorm_quantization_plugin( dtype=args.dtype) - network.plugin_config.set_quantize_tensor_plugin() network.plugin_config.set_quantize_per_token_plugin() elif args.use_weight_only: @@ -461,6 +603,16 @@ def build_rank_engine(builder: Builder, max_beam_width=args.max_beam_width, ) trtllm_model(*inputs) + if args.enable_debug_output: + # mark intermediate nodes' outputs + for k, v in tensorrt_llm_falcon.named_network_outputs(): + v = v.trt_tensor + v.name = k + network.trt_network.mark_output(v) + v.dtype = str_dtype_to_trt(args.dtype) + if args.visualize: + model_path = args.output_dir / 'test.onnx' + to_onnx(network.trt_network, model_path) tensorrt_llm.graph_rewriting.optimize(network) @@ -499,8 +651,8 @@ def build(rank, args): fp8=args.enable_fp8, strongly_typed=args.strongly_typed, opt_level=args.builder_opt, + hardware_compatibility=None, apply_query_key_layer_scaling=args.apply_query_key_layer_scaling, - eos_token_id=args.eos_token_id, gather_all_token_logits=args.gather_all_token_logits, hidden_act=args.hidden_act, hidden_size=args.hidden_size, @@ -515,7 +667,6 @@ def build(rank, args): num_heads=args.num_heads, num_kv_heads=args.num_kv_heads, num_layers=args.num_layers, - pad_token_id=args.pad_token_id, paged_kv_cache=args.paged_kv_cache, parallel_build=args.parallel_build, quant_mode=args.quant_mode, @@ -523,12 +674,41 @@ def build(rank, args): vocab_size=args.vocab_size, ) - engine_name = get_engine_name(args.model_name, args.dtype, - args.world_size, cur_rank) - engine = build_rank_engine(builder, builder_config, engine_name, - cur_rank, args) + engine_name = get_engine_name( + args.model_name, + args.dtype, + args.world_size, + cur_rank, + ) + engine = build_rank_engine( + builder, + builder_config, + engine_name, + cur_rank, + args, + ) assert engine is not None, f'Failed to build engine for rank {cur_rank}' + local_num_kv_heads = (args.num_kv_heads + args.world_size - + 1) // args.world_size + kv_dtype = str_dtype_to_trt(args.dtype) + if args.quant_mode.has_int8_kv_cache(): + kv_dtype = str_dtype_to_trt('int8') + elif args.quant_mode.has_fp8_kv_cache(): + kv_dtype = str_dtype_to_trt('fp8') + check_gpt_mem_usage( + engine=engine, + kv_dtype=kv_dtype, + use_gpt_attention_plugin=args.use_gpt_attention_plugin, + paged_kv_cache=args.paged_kv_cache, + max_batch_size=args.max_batch_size, + max_beam_width=args.max_beam_width, + max_input_len=args.max_input_len, + max_output_len=args.max_output_len, + local_num_kv_heads=local_num_kv_heads, + head_size=args.hidden_size / args.num_heads, + num_layers=args.num_layers) + if cur_rank == 0: # Use in-memory timing cache for multiple builder passes. if not args.parallel_build: @@ -554,8 +734,8 @@ def run_build(args=None): if args.parallel_build and args.world_size > 1 and \ torch.cuda.device_count() >= args.world_size: logger.warning( - f'Parallelly build TensorRT engines. Please make sure that all of the {args.world_size} GPUs are totally free.' - ) + f'Parallelly build TensorRT engines. Please make sure that all ' + f'of the {args.world_size} GPUs are totally free.') mp.spawn(build, nprocs=args.world_size, args=(args, )) else: args.parallel_build = False diff --git a/examples/chatglm/process.py b/examples/chatglm/process.py new file mode 100644 index 000000000000..8f6034bdf712 --- /dev/null +++ b/examples/chatglm/process.py @@ -0,0 +1,40 @@ +# 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 + + +def process_response_chatglm_6b(responseList): + # from chatglm-6b/modeling_chatflm.py + for i, response in enumerate(responseList): + response = response.strip() + punkts = [ + [",", "，"], + ["!", "！"], + [":", "："], + [";", "；"], + ["\?", "？"], + ] + for item in punkts: + response = re.sub(r"([\u4e00-\u9fff])%s" % item[0], + r"\1%s" % item[1], response) + response = re.sub(r"%s([\u4e00-\u9fff])" % item[0], + r"%s\1" % item[1], response) + + responseList[i] = response + return responseList + + +def process_response(responseList): + return responseList diff --git a/examples/chatglm/quantize.py b/examples/chatglm/quantize.py new file mode 100644 index 000000000000..d78601dcb104 --- /dev/null +++ b/examples/chatglm/quantize.py @@ -0,0 +1,147 @@ +# 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. +""" +Adapted from examples/quantization/hf_ptq.py +""" + +import argparse +import random + +import numpy as np +import torch +from datasets import load_dataset +from torch.utils.data import DataLoader +from transformers import AutoModelForCausalLM, AutoTokenizer + +from tensorrt_llm._utils import str_dtype_to_torch +from tensorrt_llm.logger import logger +from tensorrt_llm.models.quantized.ammo import quantize_and_export + + +def get_calib_dataloader(data="cnn_dailymail", + tokenizer=None, + batch_size=1, + calib_size=512, + block_size=512, + cache_dir=None): + print("Loading calibration dataset") + if data == "pileval": + dataset = load_dataset( + "json", + data_files="https://the-eye.eu/public/AI/pile/val.jsonl.zst", + split="train", + cache_dir=cache_dir) + dataset = dataset["text"][:calib_size] + elif data == "cnn_dailymail": + dataset = load_dataset("cnn_dailymail", + name="3.0.0", + split="train", + cache_dir=cache_dir) + dataset = dataset["article"][:calib_size] + else: + raise NotImplementedError + + batch_encoded = tokenizer.batch_encode_plus(dataset, + return_tensors="pt", + padding=True, + max_length=block_size) + batch_encoded = batch_encoded["input_ids"] + batch_encoded = batch_encoded.cuda() + + calib_dataloader = DataLoader(batch_encoded, + batch_size=batch_size, + shuffle=False) + + return calib_dataloader + + +def get_tokenizer(ckpt_path, **kwargs): + logger.info(f"Loading tokenizer from {ckpt_path}") + tokenizer = AutoTokenizer.from_pretrained(ckpt_path, + trust_remote_code=True, + padding_side="left", + **kwargs) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + return tokenizer + + +def get_model(ckpt_path, dtype="float16", cache_dir=None): + logger.info(f"Loading model from {ckpt_path}") + torch_dtype = str_dtype_to_torch(dtype) + model = AutoModelForCausalLM.from_pretrained( + ckpt_path, + device_map="auto", + cache_dir=cache_dir, + trust_remote_code=True, + torch_dtype=torch_dtype, + ) + model.eval() + model = model.to(memory_format=torch.channels_last) + return model + + +def get_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model_dir", + type=str, + required=True, + help="Directory of a HF model checkpoint") + parser.add_argument("--dtype", help="Model data type.", default="float16") + parser.add_argument( + "--qformat", + type=str, + choices=['fp8', 'int4_awq'], + default='fp8', + help='Quantization format. Currently only fp8 is supported. ' + 'For int8 smoothquant, use smoothquant.py instead. ') + parser.add_argument("--calib_size", + type=int, + default=512, + help="Number of samples for calibration.") + parser.add_argument("--export_path", default="exported_model") + parser.add_argument("--cache_dir", + type=str, + default=None, + help="Directory of dataset cache.") + parser.add_argument('--seed', type=int, default=None, help='Random seed') + args = parser.parse_args() + return args + + +def main(): + if not torch.cuda.is_available(): + raise EnvironmentError("GPU is required for inference.") + + args = get_args() + + if args.seed is not None: + random.seed(args.seed) + np.random.seed(args.seed) + + tokenizer = get_tokenizer(args.model_dir, cache_dir=args.cache_dir) + model = get_model(args.model_dir, args.dtype, cache_dir=args.cache_dir) + + calib_dataloader = get_calib_dataloader(tokenizer=tokenizer, + calib_size=args.calib_size, + cache_dir=args.cache_dir) + model = quantize_and_export(model, + qformat=args.qformat, + calib_dataloader=calib_dataloader, + export_path=args.export_path) + + +if __name__ == "__main__": + main() diff --git a/examples/chatglm/requirements.txt b/examples/chatglm/requirements.txt index 140929584e0f..dd2c60a1545b 100644 --- a/examples/chatglm/requirements.txt +++ b/examples/chatglm/requirements.txt @@ -1,5 +1,5 @@ datasets~=2.14.5 -evaluate +evaluate~=0.4.1 protobuf rouge_score~=0.1.2 sentencepiece diff --git a/examples/chatglm/run.py b/examples/chatglm/run.py index 6eaf71b4b48e..24559b3de56d 100644 --- a/examples/chatglm/run.py +++ b/examples/chatglm/run.py @@ -14,8 +14,6 @@ # limitations under the License. import argparse import json -import os -import re from pathlib import Path import torch @@ -32,13 +30,16 @@ def parse_arguments(args=None): parser = argparse.ArgumentParser() parser.add_argument( - '--model_version', + '--model_name', '-m', type=str, - default="3", - choices=["1", "2", "3", "2-32k", "3-32k"], + required=True, + choices=[ + "chatglm_6b", "chatglm2_6b", "chatglm2_6b_32k", "chatglm3_6b", + "chatglm3_6b_base", "chatglm3_6b_32k", "glm_10b" + ], help= - '1, 2, 3, 2-32k, 3-32k for ChatGLM-6B, ChatGLM2-6B, ChatGLM3-6B, ChatGLM2-32k and ChatGLM3-32k respectively' + 'the name of the model, use "_" rather than "-" to connect the name parts' ) parser.add_argument('--max_output_len', type=int, default=1024) parser.add_argument('--log_level', type=str, default='error') @@ -73,44 +74,15 @@ def parse_arguments(args=None): return parser.parse_args(args) -def process_response(responseList): - for i, response in enumerate(responseList): - response = response.strip() - punkts = [ - [",", "，"], - ["!", "！"], - [":", "："], - [";", "；"], - ["\?", "？"], - ] - for item in punkts: - response = re.sub(r"([\u4e00-\u9fff])%s" % item[0], - r"\1%s" % item[1], response) - response = re.sub(r"%s([\u4e00-\u9fff])" % item[0], - r"%s\1" % item[1], response) - - responseList[i] = response - return responseList - - if __name__ == '__main__': args = parse_arguments() tensorrt_llm.logger.set_level(args.log_level) - if args.model_version == "1": - model_name = "chatglm-6b" - elif args.model_version in ["2", "3"]: - model_name = "chatglm%s-6b" % args.model_version - else: - model_name = "chatglm%s-6b-32k" % args.model_version.split("-")[0] - - config_path = os.path.join(args.engine_dir, model_name + '-config.json') + config_path = Path(args.engine_dir) / (args.model_name + '-config.json') with open(config_path, 'r') as f: config = json.load(f) dtype = config['builder_config']['precision'] - end_id = config['builder_config']['eos_token_id'] - pad_id = config['builder_config']['pad_token_id'] max_batch_size = config['builder_config']['max_batch_size'] max_input_len = config['builder_config']['max_input_len'] max_output_len = config['builder_config']['max_output_len'] @@ -138,16 +110,21 @@ def process_response(responseList): serialize_path = find_engines( Path(args.engine_dir), - model_name=model_name, + model_name=args.model_name, dtype=dtype, tp_size=world_size, rank=runtime_rank, )[0] if args.tokenizer_dir is None: - args.tokenizer_dir = model_name + args.tokenizer_dir = args.model_name tokenizer = transformers.AutoTokenizer.from_pretrained( args.tokenizer_dir, trust_remote_code=True) + end_id = tokenizer.eos_token_id + pad_id = tokenizer.pad_token_id + if args.model_name in ["glm_10b"]: + sop_id = tokenizer.sop_token_id + eop_id = tokenizer.eop_token_id input_ids = None input_text = None if args.input_tokens is None: @@ -171,6 +148,13 @@ def process_response(responseList): max_input_len, input_lengths) else: max_input_len = max_input_len_real + if args.model_name in ["glm_10b"]: + input_ids = torch.cat( + (input_ids, input_ids.new_full((batch_size, 1), sop_id)), + dim=-1, + ) + input_lengths += 1 + max_input_len_real += 1 else: input_ids = [] @@ -182,7 +166,6 @@ def process_response(responseList): input_ids = torch.tensor(input_ids, dtype=torch.int32).cuda().unsqueeze(0) - input_ids_padding = input_ids.clone() if remove_input_padding: input_ids_no_padding = torch.zeros(1, torch.sum(input_lengths), @@ -221,14 +204,14 @@ def process_response(responseList): hidden_size=config['builder_config']['hidden_size'] // world_size, gpt_attention_plugin=use_gpt_attention_plugin, remove_input_padding=config['builder_config']['remove_input_padding'], - model_name=model_name, + model_name=args.model_name, paged_kv_cache=config['builder_config']['paged_kv_cache'], quant_mode=QuantMode(config['builder_config']['quant_mode']), dtype=dtype, ) sampling_config = SamplingConfig( - end_id=end_id, + end_id=eop_id if args.model_name in ["glm_10b"] else end_id, pad_id=pad_id, num_beams=beam_width, temperature=args.temperature, @@ -240,18 +223,22 @@ def process_response(responseList): with open(serialize_path, 'rb') as f: engine_buffer = f.read() - if model_name == "chatglm-6b": - decoder = ChatGLMGenerationSession( - model_config, - engine_buffer, - runtime_mapping, - ) - else: - decoder = GenerationSession( - model_config, - engine_buffer, - runtime_mapping, - ) + if args.model_name in ["chatglm_6b", "glm_10b"]: + session = ChatGLMGenerationSession + elif args.model_name in [ + "chatglm2_6b", + "chatglm2_6b_32k", + "chatglm3_6b", + "chatglm3_6b_base", + "chatglm3_6b_32k", + ]: + session = GenerationSession + decoder = session( + model_config, + engine_buffer, + runtime_mapping, + ) + decoder.setup( len(input_text), max_input_len, @@ -266,17 +253,31 @@ def process_response(responseList): return_dict=True, ) torch.cuda.synchronize() + output_ids = output["output_ids"] output_lengths = output["sequence_lengths"] if runtime_rank == 0: + + if args.model_name in ["chatglm_6b"]: + from process import process_response_chatglm_6b as process_response + elif args.model_name in [ + "chatglm2_6b", + "chatglm2_6b_32k", + "chatglm3_6b", + "chatglm3_6b_base", + "chatglm3_6b_32k", + "glm_10b", + ]: + from process import process_response + for i in range(batch_size): print("\nInput %2d ---> len=%d\n%s" % (i, input_lengths[i], input_text[i])) print("\nOutput %2d --->" % i) - output_ids__one_batch = output_ids[i, :, input_lengths[i]:] + output_ids_one_batch = output_ids[i, :, input_lengths[i]:] output_lengths_one_batch = output_lengths[i] - output_token_list = tokenizer.batch_decode(output_ids__one_batch, + output_token_list = tokenizer.batch_decode(output_ids_one_batch, skip_special_tokens=True) output_token_list = process_response(output_token_list) for j, (length, simple_output) in enumerate( diff --git a/examples/chatglm/summarize.py b/examples/chatglm/summarize.py deleted file mode 100644 index daf759c7929f..000000000000 --- a/examples/chatglm/summarize.py +++ /dev/null @@ -1,473 +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 copy -import json -from pathlib import Path - -import evaluate -import numpy as np -import torch -from datasets import load_dataset -from transformers import AutoModel, AutoTokenizer - -import tensorrt_llm -import tensorrt_llm.profiler as profiler -from tensorrt_llm.logger import logger -from tensorrt_llm.runtime import (ChatGLMGenerationSession, GenerationSession, - ModelConfig, SamplingConfig) - -from build import find_engines # isort:skip - -model_name = "" - - -def TRT(args, config): - - model_name = config['builder_config']['name'] - dtype = config['builder_config']['precision'] - world_size = config['builder_config']['tensor_parallel'] - assert world_size == tensorrt_llm.mpi_world_size(), \ - f'Engine world size ({world_size}) != Runtime world size ({tensorrt_llm.mpi_world_size()})' - - world_size = config['builder_config']['tensor_parallel'] - remove_input_padding = config['plugin_config']['remove_input_padding'] - - model_config = ModelConfig( - model_name=model_name, - vocab_size=config['builder_config']['vocab_size'], - num_layers=config['builder_config']['num_layers'], - num_heads=config['builder_config']['num_heads'] // world_size, - num_kv_heads=max(config['builder_config']['num_kv_heads'] // world_size, - 1), - hidden_size=config['builder_config']['hidden_size'] // world_size, - gpt_attention_plugin=bool( - config['plugin_config']['gpt_attention_plugin']), - remove_input_padding=remove_input_padding, - tokens_per_block=config['plugin_config']['tokens_per_block'], - paged_kv_cache=config['plugin_config']['paged_kv_cache'], - dtype=dtype, - use_custom_all_reduce=config['plugin_config']['use_custom_all_reduce'], - ) - - runtime_rank = tensorrt_llm.mpi_rank() - runtime_mapping = tensorrt_llm.Mapping(world_size, - runtime_rank, - tp_size=world_size) - torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node) - - serialize_path = find_engines( - args.engine_dir, - model_name=model_name, - dtype=dtype, - tp_size=world_size, - rank=runtime_rank, - )[0] - - tensorrt_llm.logger.set_level(args.log_level) - - with open(serialize_path, 'rb') as f: - engine_buffer = f.read() - - if model_name == 'chatglm-6b': - decoder = ChatGLMGenerationSession( - model_config, - engine_buffer, - runtime_mapping, - ) - else: - decoder = GenerationSession( - model_config, - engine_buffer, - runtime_mapping, - ) - - return decoder - - -def main(args): - 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 - tokenizer = AutoTokenizer.from_pretrained( - args.tokenizer, - padding_side='left', - trust_remote_code=True, - ) - - if args.eval_type == 'code_completion': - dataset_name = "openai_humaneval" - dataset_revision = None - dataset_input_key = 'prompt' - dataset_output_key = 'canonical_solution' - elif args.eval_type == 'summarize': - dataset_name = "ccdv/cnn_dailymail" - dataset_revision = "3.0.0" - dataset_input_key = 'article' - dataset_output_key = 'highlights' - args.dataset_path.mkdir(parents=True, exist_ok=True) - dataset = load_dataset(dataset_name, - dataset_revision, - cache_dir=args.dataset_path) - - config_path = str(args.engine_dir / 'config.json') - with open(config_path, 'r') as f: - config = json.load(f) - - max_batch_size = args.batch_size - - # runtime parameters - # repetition_penalty = 1 - top_k = args.top_k - output_len = args.output_len - test_token_num = 800 - # top_p = 0.0 - # random_seed = 5 - temperature = 1 - num_beams = args.num_beams - length_penalty = args.length_penalty - - pad_id = tokenizer.encode(tokenizer.pad_token, add_special_tokens=False)[0] - end_id = tokenizer.encode(tokenizer.eos_token, add_special_tokens=False)[0] - - if test_trt_llm: - tensorrt_llm_gpt = TRT(args, config) - - if test_hf: - model = AutoModel.from_pretrained( - args.hf_model_location, - trust_remote_code=True, - ) - model.cuda() - if args.data_type == 'fp16': - model.half() - - def eval_tensorrt_llm(datapoint, eval_type='summarize'): - batch_size = len(datapoint) - append_str = ' TL;DR: ' if eval_type == 'summarize' else '' - line = copy.copy(datapoint) - line_encoded = [] - input_lengths = [] - for i in range(batch_size): - line[i] = line[i] + append_str - - line[i] = line[i].strip() - line[i] = line[i].replace(" n't", "n't") - - input_id = tokenizer.encode( - line[i], - return_tensors='pt', - ).type(torch.int32) - if model_name == 'chatglm-6b': - input_id = input_id[:, -test_token_num:] - else: - input_id = input_id[:, :test_token_num] - - line_encoded.append(input_id) - input_lengths.append(input_id.shape[-1]) - - max_length = max(input_lengths) - - if tensorrt_llm_gpt.remove_input_padding: - line_encoded = [t.to(torch.int32).cuda() for t in line_encoded] - else: - # do padding, should move outside the profiling to prevent the overhead - for i in range(batch_size): - pad_size = max_length - input_lengths[i] - - pad = torch.ones([1, pad_size], dtype=torch.int32) * pad_id - line_encoded[i] = torch.cat( - [line_encoded[i].to(torch.int32), pad], axis=-1) - - line_encoded = torch.cat(line_encoded, axis=0).cuda() - input_lengths = torch.tensor(input_lengths, - dtype=torch.int32).cuda() - - sampling_config = SamplingConfig( - end_id=end_id, - pad_id=pad_id, - top_k=top_k, - num_beams=num_beams, - length_penalty=length_penalty, - ) - - with torch.no_grad(): - tensorrt_llm_gpt.setup(batch_size, - max_context_length=max_length, - max_new_tokens=output_len, - beam_width=num_beams, - max_kv_cache_length=args.max_kv_cache_len) - - if tensorrt_llm_gpt.remove_input_padding: - output_ids = tensorrt_llm_gpt.decode_batch( - line_encoded, sampling_config) - else: - output_ids = tensorrt_llm_gpt.decode( - line_encoded, - input_lengths, - sampling_config, - ) - - torch.cuda.synchronize() - - # Extract a list of tensors of shape beam_width x output_ids. - if tensorrt_llm_gpt.mapping.is_first_pp_rank(): - output_beams_list = [ - tokenizer.batch_decode(output_ids[batch_idx, :, - input_lengths[batch_idx]:], - skip_special_tokens=True) - for batch_idx in range(batch_size) - ] - return output_beams_list, output_ids[:, :, max_length:].tolist() - return [], [] - - def eval_hf(datapoint, eval_type='summarize'): - batch_size = len(datapoint) - append_str = ' TL;DR: ' if eval_type == 'summarize' else '' - if batch_size > 1: - logger.warning( - f"HF does not support batch_size > 1 to verify correctness due to padding and attention mask. Current batch size is {batch_size}" - ) - - line = copy.copy(datapoint) - line_encoded = [] - input_lengths = [] - for i in range(batch_size): - line[i] = line[i] + append_str - - line[i] = line[i].strip() - line[i] = line[i].replace(" n't", "n't") - - input_id = tokenizer.encode( - line[i], - return_tensors='pt', - ).type(torch.int64) - if model_name == 'chatglm-6b': - input_id = input_id[:, -test_token_num:] - else: - input_id = input_id[:, :test_token_num] - - line_encoded.append(input_id) - input_lengths.append(input_id.shape[-1]) - - max_length = max(input_lengths) - - for i in range(batch_size): - pad_size = max_length - input_lengths[i] - - pad = torch.ones([1, pad_size], dtype=torch.int64) * pad_id - line_encoded[i] = torch.cat([pad, line_encoded[i].to(torch.int64)], - axis=-1) - - line_encoded = torch.cat(line_encoded, axis=0).cuda() - - with torch.no_grad(): - output = model.generate(line_encoded, - max_length=len(line_encoded[0]) + - output_len, - top_k=top_k, - temperature=temperature, - eos_token_id=tokenizer.eos_token_id, - pad_token_id=tokenizer.pad_token_id, - num_beams=num_beams, - num_return_sequences=num_beams, - early_stopping=True, - length_penalty=length_penalty) - - tokens_list = output[:, len(line_encoded[0]):].tolist() - output = output.reshape([batch_size, num_beams, -1]) - output_lines_list = [ - tokenizer.batch_decode(output[:, i, len(line_encoded[0]):], - skip_special_tokens=True) - for i in range(num_beams) - ] - - return output_lines_list, tokens_list - - if test_trt_llm: - datapoint = dataset['test'][0:1] - output, _ = eval_tensorrt_llm(datapoint[dataset_input_key], - eval_type=args.eval_type) - if runtime_rank == 0: - 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( - "---------------------------------------------------------") - - if test_hf: - datapoint = dataset['test'][0:1] - output, _ = eval_hf(datapoint[dataset_input_key], - eval_type=args.eval_type) - 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("---------------------------------------------------------") - - metric_tensorrt_llm = [evaluate.load("rouge") for _ in range(num_beams)] - metric_hf = [evaluate.load("rouge") for _ in range(num_beams)] - for i in range(num_beams): - metric_tensorrt_llm[i].seed = 0 - metric_hf[i].seed = 0 - - ite_count = 0 - data_point_idx = 0 - while (data_point_idx < len(dataset['test'])) 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['test'][data_point_idx:(data_point_idx + - max_batch_size)] - - if test_trt_llm: - profiler.start('tensorrt_llm') - output_tensorrt_llm, _ = eval_tensorrt_llm( - datapoint[dataset_input_key]) - profiler.stop('tensorrt_llm') - - if test_hf: - profiler.start('hf') - output_hf, _ = eval_hf(datapoint[dataset_input_key]) - profiler.stop('hf') - - if runtime_rank == 0: - if test_trt_llm: - for batch_idx in range(len(output_tensorrt_llm)): - for beam_idx in range(num_beams): - metric_tensorrt_llm[beam_idx].add_batch( - predictions=[ - output_tensorrt_llm[batch_idx][beam_idx] - ], - references=[ - datapoint[dataset_output_key][batch_idx] - ]) - if test_hf: - for beam_idx in range(num_beams): - 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] - ]) - - logger.debug('-' * 100) - logger.debug(f"Input : {datapoint[dataset_input_key]}") - if test_trt_llm: - logger.debug(f'TensorRT-LLM Output: {output_tensorrt_llm}') - if test_hf: - logger.debug(f'HF Output: {output_hf}') - logger.debug(f"highlights : {datapoint[dataset_output_key]}") - - data_point_idx += max_batch_size - ite_count += 1 - - if runtime_rank == 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)' - ) - for beam_idx in range(num_beams): - logger.info(f"TensorRT-LLM beam {beam_idx} result") - 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: - assert computed_metrics_tensorrt_llm[ - 'rouge1'] * 100 > args.tensorrt_llm_rouge1_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)' - ) - for beam_idx in range(num_beams): - logger.info(f"HF beam {beam_idx} result") - computed_metrics_hf = metric_hf[beam_idx].compute() - for key in computed_metrics_hf.keys(): - logger.info(f' {key} : {computed_metrics_hf[key] * 100}') - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument( - '--model_version', - '-m', - type=str, - required=True, - choices=["1", "2", "3", "2-32k", "3-32k"], - help= - '1, 2, 3, 2-32k, 3-32k for ChatGLM-6B, ChatGLM2-6B, ChatGLM3-6B, ChatGLM2-32k and ChatGLM3-32k respectively' - ) - parser.add_argument('--hf_model_location', type=str, default=None) - parser.add_argument( - '--tokenizer', - default=None, - help='tokenizer path; defaults to hf_model_location if left unspecified' - ) - parser.add_argument('--test_hf', action='store_true', default=True) - parser.add_argument('--test_trt_llm', action='store_true', default=True) - parser.add_argument('--data_type', - type=str, - choices=['fp32', 'fp16'], - default='fp16') - parser.add_argument('--dataset_path', type=Path, default='dataset') - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument('--engine_dir', type=Path, default='trtModel') - 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_kv_cache_len', - type=int, - default=None, - help='The max kv cache length. \ - If the final sequence length exceeds the kv cache length, we will enable cyclic kv cache. \ - If it is set to None, we will use the max sequence length.') - parser.add_argument('--check_accuracy', action='store_true', default=True) - parser.add_argument('--tensorrt_llm_rouge1_threshold', - type=float, - default=15.0) - parser.add_argument('--num_beams', type=int, default=1) - parser.add_argument('--top_k', type=int, default=1) - parser.add_argument('--eval_type', - type=str, - default='summarize', - choices=['summarize', 'code_completion']) - parser.add_argument('--length_penalty', type=float, default=1.0) - - args = parser.parse_args() - - if args.model_version == "1": - args.model_name = "chatglm-6b" - elif args.model_version in ["2", "3"]: - args.model_name = "chatglm%s-6b" % args.model_version - else: - args.model_name = "chatglm%s-6b-32k" % args.model_version.split("-")[0] - - if args.tokenizer == None: - args.tokenizer = args.model_name - - main(args) diff --git a/examples/chatglm/weight.py b/examples/chatglm/weight.py index 4961c499ee65..e47f45ebcbf9 100644 --- a/examples/chatglm/weight.py +++ b/examples/chatglm/weight.py @@ -13,15 +13,26 @@ # See the License for the specific language governing permissions and # limitations under the License. import time +from pathlib import Path +from typing import Dict, List, Optional, Union +import numpy as np import torch import torch.nn.functional as F +import transformers import tensorrt_llm +import tensorrt_llm.logger as logger from tensorrt_llm._utils import str_dtype_to_torch, torch_to_numpy +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.quantized.quant import get_dummy_quant_scales from tensorrt_llm.quantization import QuantMode +def split_matrix(weight: np.ndarray, tp_size: int, rank: int, dim: int): + return np.ascontiguousarray(split(weight, tp_size, rank, dim=dim)) + + def tile_kv_weight_bias(v, kv_num_head, tp_size): head_size = v.shape[0] // kv_num_head reps = tp_size // kv_num_head @@ -69,16 +80,30 @@ def load_quant_weight(src, value_dst, scale_dst, plugin_weight_only_quant_type): def load_from_hf( trt_model, - hf_model, + hf_model_dir, mapping=None, dtype="float32", - model_version="3", + model_name=None, multi_query_mode=False, ): - # [TODO] Merge model_version=="1" and model_version>="2" + + assert model_name is not None, "Model name must be set" + tensorrt_llm.logger.info("Loading weights from HF") + + if not Path(hf_model_dir).exists(): + tensorrt_llm.logger.info( + "No weight file found from %s, use random weights" % hf_model_dir) + return trt_model + tik = time.time() + hf_model = transformers.AutoModel.from_pretrained(hf_model_dir, + trust_remote_code=True) + num_layers = hf_model.config.num_layers + hidden_size = hf_model.config.hidden_size + num_heads = hf_model.config.num_attention_heads + torch_type = str_dtype_to_torch(dtype) quant_mode = getattr(trt_model, 'quant_mode', QuantMode(0)) if quant_mode.is_int8_weight_only(): @@ -87,36 +112,85 @@ def load_from_hf( plugin_weight_only_quant_type = torch.quint4x2 use_weight_only = quant_mode.is_weight_only() - hidden_size = hf_model.config.hidden_size - num_heads = hf_model.config.num_attention_heads - - layers_per_pipeline_stage = trt_model.num_layers // mapping.pp_size + layers_per_pipeline_stage = num_layers // mapping.pp_size layers_range = list( range(mapping.pp_rank * layers_per_pipeline_stage, (mapping.pp_rank + 1) * layers_per_pipeline_stage)) feed_weight_count = 0 - if model_version == "1": + if model_name in ["chatglm_6b", "glm_10b"]: num_kv_heads = hf_model.config.num_attention_heads + elif model_name in [ + "chatglm2_6b", + "chatglm2_6b_32k", + "chatglm3_6b", + "chatglm3_6b_base", + "chatglm3_6b_32k", + ]: + num_kv_heads = hf_model.config.multi_query_group_num - if mapping.is_first_pp_rank(): - # Embedding + if mapping.is_first_pp_rank(): + # Embedding + if model_name in ["chatglm_6b"]: weight = hf_model.transformer.word_embeddings.weight.to( - torch_type).detach().cpu() + torch_type).detach() + trt_model.embedding.weight.value = torch_to_numpy(weight) + feed_weight_count += 1 + elif model_name in [ + "chatglm2_6b", + "chatglm2_6b_32k", + "chatglm3_6b", + "chatglm3_6b_base", + "chatglm3_6b_32k", + ]: + weight = hf_model.transformer.embedding.word_embeddings.weight.to( + torch_type).detach() + trt_model.embedding.weight.value = torch_to_numpy(weight) + feed_weight_count += 1 + elif model_name in ["glm_10b"]: + weight = hf_model.word_embeddings.weight.to(torch_type).detach() trt_model.embedding.weight.value = torch_to_numpy(weight) + weight = hf_model.transformer.position_embeddings.weight.to( + torch_type).detach() + trt_model.position_embeddings.weight.value = torch_to_numpy(weight) + weight = hf_model.transformer.block_position_embeddings.weight.to( + torch_type).detach() + trt_model.block_embeddings.weight.value = torch_to_numpy(weight) + feed_weight_count += 3 + + if mapping.is_last_pp_rank(): + # Final normalization + if model_name in ["chatglm_6b"]: + weight = hf_model.transformer.final_layernorm.weight.to( + torch_type).detach() + trt_model.final_norm.weight.value = torch_to_numpy(weight) + bias = hf_model.transformer.final_layernorm.bias.to( + torch_type).detach() + trt_model.final_norm.bias.value = torch_to_numpy(bias) + feed_weight_count += 2 + elif model_name in [ + "chatglm2_6b", + "chatglm2_6b_32k", + "chatglm3_6b", + "chatglm3_6b_base", + "chatglm3_6b_32k", + ]: + weight = hf_model.transformer.encoder.final_layernorm.weight.to( + torch_type).detach() + trt_model.final_norm.weight.value = torch_to_numpy(weight) feed_weight_count += 1 - if mapping.is_last_pp_rank(): - # Final normalization + elif model_name in ["glm_10b"]: weight = hf_model.transformer.final_layernorm.weight.to( - torch_type).detach().cpu() + torch_type).detach() trt_model.final_norm.weight.value = torch_to_numpy(weight) bias = hf_model.transformer.final_layernorm.bias.to( - torch_type).detach().cpu() + torch_type).detach() trt_model.final_norm.bias.value = torch_to_numpy(bias) feed_weight_count += 2 - # Final LM - weight = hf_model.lm_head.weight.to(torch_type).detach().cpu() + # Final LM + if model_name in ["chatglm_6b"]: + weight = hf_model.lm_head.weight.to(torch_type).detach() if weight.shape[0] % mapping.tp_size != 0: pad_width = trt_model.lm_head.out_features * mapping.tp_size - weight.shape[ 0] @@ -125,128 +199,25 @@ def load_from_hf( dim=0)[mapping.rank] trt_model.lm_head.weight.value = torch_to_numpy(split_weight) feed_weight_count += 1 - - for layer_idx in range(28): - if layer_idx not in layers_range: - continue - i = int(layer_idx) - mapping.pp_rank * layers_per_pipeline_stage - if i >= trt_model.num_layers: - continue - - # Pre normalization - weight = hf_model.transformer.layers[i].input_layernorm.weight.to( - torch_type).detach().cpu() - trt_model.layers[i].pre_norm.weight.value = torch_to_numpy(weight) - bias = hf_model.transformer.layers[i].input_layernorm.bias.to( - torch_type).detach().cpu() - trt_model.layers[i].pre_norm.bias.value = torch_to_numpy(bias) - feed_weight_count += 2 - - # QKV multiplication weight - weight = hf_model.transformer.layers[ - i].attention.query_key_value.weight.to( - torch_type).detach().cpu() - split_weight = split_qkv(weight, mapping.tp_size, mapping.tp_rank, - hidden_size, num_heads, num_kv_heads) - dst = trt_model.layers[i].attention.qkv - if use_weight_only: - load_quant_weight( - src=split_weight, - value_dst=dst.weight, - scale_dst=dst.per_channel_scale, - plugin_weight_only_quant_type=plugin_weight_only_quant_type) - else: - dst.weight.value = torch_to_numpy(split_weight) - feed_weight_count += 1 - - # QKV multiplication bias - bias = hf_model.transformer.layers[ - i].attention.query_key_value.bias.to(torch_type).detach().cpu() - split_bias = split_qkv(bias, mapping.tp_size, mapping.tp_rank, - hidden_size, num_heads, num_kv_heads) - trt_model.layers[i].attention.qkv.bias.value = torch_to_numpy( - split_bias) - feed_weight_count += 1 - - # Dense multiplication weight (no bias) - weight = hf_model.transformer.layers[i].attention.dense.weight.to( - torch_type).detach().cpu() - split_weight = torch.chunk(weight, mapping.tp_size, - dim=1)[mapping.rank] - dst = trt_model.layers[i].attention.dense - if use_weight_only: - load_quant_weight( - src=split_weight, - value_dst=dst.weight, - scale_dst=dst.per_channel_scale, - plugin_weight_only_quant_type=plugin_weight_only_quant_type) - else: - dst.weight.value = torch_to_numpy(split_weight) - feed_weight_count += 1 - - # Post normalization - weight = hf_model.transformer.layers[ - i].post_attention_layernorm.weight.to( - torch_type).detach().cpu() - trt_model.layers[i].post_norm.weight.value = torch_to_numpy(weight) - bias = hf_model.transformer.layers[ - i].post_attention_layernorm.bias.to(torch_type).detach().cpu() - trt_model.layers[i].post_norm.bias.value = torch_to_numpy(bias) - feed_weight_count += 2 - - # Multilayer perceptron h -> 4h (no bias) - weight = hf_model.transformer.layers[i].mlp.dense_h_to_4h.weight.to( - torch_type).detach().cpu() + elif model_name in [ + "chatglm2_6b", + "chatglm2_6b_32k", + "chatglm3_6b", + "chatglm3_6b_base", + "chatglm3_6b_32k", + ]: + weight = hf_model.transformer.output_layer.weight.to( + torch_type).detach() + if weight.shape[0] % mapping.tp_size != 0: + pad_width = trt_model.lm_head.out_features * mapping.tp_size - weight.shape[ + 0] + weight = F.pad(weight, (0, 0, 0, pad_width)) split_weight = torch.chunk(weight, mapping.tp_size, dim=0)[mapping.rank] - dst = trt_model.layers[i].mlp.fc - if use_weight_only: - load_quant_weight( - src=split_weight, - value_dst=dst.weight, - scale_dst=dst.per_channel_scale, - plugin_weight_only_quant_type=plugin_weight_only_quant_type) - else: - dst.weight.value = torch_to_numpy(split_weight) - feed_weight_count += 1 - - # Multilayer perceptron 4h -> h (no bias) - weight = hf_model.transformer.layers[i].mlp.dense_4h_to_h.weight.to( - torch_type).detach().cpu() - split_weight = torch.chunk(weight, mapping.tp_size, - dim=1)[mapping.rank] - dst = trt_model.layers[i].mlp.proj - if use_weight_only: - load_quant_weight( - src=split_weight, - value_dst=dst.weight, - scale_dst=dst.per_channel_scale, - plugin_weight_only_quant_type=plugin_weight_only_quant_type) - else: - dst.weight.value = torch_to_numpy(split_weight) - feed_weight_count += 1 - - assert feed_weight_count == 4 + trt_model.num_layers * 9, "Some weights not loaded from HF" - - else: - num_kv_heads = hf_model.config.multi_query_group_num - - if mapping.is_first_pp_rank(): - # Embedding - weight = hf_model.transformer.embedding.word_embeddings.weight.to( - torch_type).detach().cpu() - trt_model.embedding.weight.value = torch_to_numpy(weight) - feed_weight_count += 1 - if mapping.is_last_pp_rank(): - # Final normalization - weight = hf_model.transformer.encoder.final_layernorm.weight.to( - torch_type).detach().cpu() - trt_model.final_norm.weight.value = torch_to_numpy(weight) + trt_model.lm_head.weight.value = torch_to_numpy(split_weight) feed_weight_count += 1 - - # Final LM - weight = hf_model.transformer.output_layer.weight.to( - torch_type).detach().cpu() + elif model_name in ["glm_10b"]: + weight = hf_model.word_embeddings.weight.to(torch_type).detach() if weight.shape[0] % mapping.tp_size != 0: pad_width = trt_model.lm_head.out_features * mapping.tp_size - weight.shape[ 0] @@ -256,72 +227,180 @@ def load_from_hf( trt_model.lm_head.weight.value = torch_to_numpy(split_weight) feed_weight_count += 1 - for layer_idx in range(28): - if layer_idx not in layers_range: - continue - i = int(layer_idx) - mapping.pp_rank * layers_per_pipeline_stage - if i >= trt_model.num_layers: - continue + # Weight per layer + for layer_idx in range(num_layers): + if layer_idx not in layers_range: + continue + i = int(layer_idx) - mapping.pp_rank * layers_per_pipeline_stage + if i >= num_layers: + continue - # Pre normalization + # Pre normalization + if model_name in ["chatglm_6b"]: + weight = hf_model.transformer.layers[i].input_layernorm.weight.to( + torch_type).detach() + trt_model.layers[i].pre_norm.weight.value = torch_to_numpy(weight) + bias = hf_model.transformer.layers[i].input_layernorm.bias.to( + torch_type).detach() + trt_model.layers[i].pre_norm.bias.value = torch_to_numpy(bias) + feed_weight_count += 2 + elif model_name in [ + "chatglm2_6b", + "chatglm2_6b_32k", + "chatglm3_6b", + "chatglm3_6b_base", + "chatglm3_6b_32k", + ]: weight = hf_model.transformer.encoder.layers[ - i].input_layernorm.weight.to(torch_type).detach().cpu() + i].input_layernorm.weight.to(torch_type).detach() trt_model.layers[i].pre_norm.weight.value = torch_to_numpy(weight) feed_weight_count += 1 + elif model_name in ["glm_10b"]: + weight = hf_model.transformer.layers[i].input_layernorm.weight.to( + torch_type).detach() + trt_model.layers[i].pre_norm.weight.value = torch_to_numpy(weight) + bias = hf_model.transformer.layers[i].input_layernorm.bias.to( + torch_type).detach() + trt_model.layers[i].pre_norm.bias.value = torch_to_numpy(bias) + feed_weight_count += 2 - # QKV multiplication weight + # QKV multiplication weight + if model_name in ["chatglm_6b"]: + weight = hf_model.transformer.layers[ + i].attention.query_key_value.weight.to(torch_type).detach() + elif model_name in [ + "chatglm2_6b", + "chatglm2_6b_32k", + "chatglm3_6b", + "chatglm3_6b_base", + "chatglm3_6b_32k", + ]: weight = hf_model.transformer.encoder.layers[ i].self_attention.query_key_value.weight.to( - torch_type).detach().cpu() - split_weight = split_qkv(weight, mapping.tp_size, mapping.tp_rank, - hidden_size, num_heads, num_kv_heads) - dst = trt_model.layers[i].attention.qkv - if use_weight_only: - load_quant_weight( - src=split_weight, - value_dst=dst.weight, - scale_dst=dst.per_channel_scale, - plugin_weight_only_quant_type=plugin_weight_only_quant_type) - else: - dst.weight.value = torch_to_numpy(split_weight) - feed_weight_count += 1 - - # QKV multiplication bias + torch_type).detach() + elif model_name in ["glm_10b"]: + weight = hf_model.transformer.layers[ + i].attention.query_key_value.weight.to(torch_type).detach() + + split_weight = split_qkv(weight, mapping.tp_size, mapping.tp_rank, + hidden_size, num_heads, num_kv_heads) + dst = trt_model.layers[i].attention.qkv + if use_weight_only: + load_quant_weight( + src=split_weight, + value_dst=dst.weight, + scale_dst=dst.per_channel_scale, + plugin_weight_only_quant_type=plugin_weight_only_quant_type) + else: + dst.weight.value = torch_to_numpy(split_weight) + feed_weight_count += 1 + + # QKV multiplication bias + if model_name in ["chatglm_6b"]: + bias = hf_model.transformer.layers[ + i].attention.query_key_value.bias.to(torch_type).detach() + elif model_name in [ + "chatglm2_6b", + "chatglm2_6b_32k", + "chatglm3_6b", + "chatglm3_6b_base", + "chatglm3_6b_32k", + ]: bias = hf_model.transformer.encoder.layers[ - i].self_attention.query_key_value.bias.to( - torch_type).detach().cpu() + i].self_attention.query_key_value.bias.to(torch_type).detach() + elif model_name in ["glm_10b"]: + bias = hf_model.transformer.layers[ + i].attention.query_key_value.bias.to(torch_type).detach() + + split_bias = split_qkv(bias, mapping.tp_size, mapping.tp_rank, + hidden_size, num_heads, num_kv_heads) + trt_model.layers[i].attention.qkv.bias.value = torch_to_numpy( + split_bias) + feed_weight_count += 1 + + # Dense multiplication weight + if model_name in ["chatglm_6b"]: + weight = hf_model.transformer.layers[i].attention.dense.weight.to( + torch_type).detach() + elif model_name in [ + "chatglm2_6b", + "chatglm2_6b_32k", + "chatglm3_6b", + "chatglm3_6b_base", + "chatglm3_6b_32k", + ]: + weight = hf_model.transformer.encoder.layers[ + i].self_attention.dense.weight.to(torch_type).detach() + elif model_name in ["glm_10b"]: + weight = hf_model.transformer.layers[i].attention.dense.weight.to( + torch_type).detach() + + split_weight = torch.chunk(weight, mapping.tp_size, dim=1)[mapping.rank] + dst = trt_model.layers[i].attention.dense + if use_weight_only: + load_quant_weight( + src=split_weight, + value_dst=dst.weight, + scale_dst=dst.per_channel_scale, + plugin_weight_only_quant_type=plugin_weight_only_quant_type) + else: + dst.weight.value = torch_to_numpy(split_weight) + feed_weight_count += 1 + + # Dense multiplication bias, only GLM-10B + if model_name in ["glm_10b"]: + bias = hf_model.transformer.layers[i].attention.dense.bias.to( + torch_type).detach() split_bias = split_qkv(bias, mapping.tp_size, mapping.tp_rank, hidden_size, num_heads, num_kv_heads) - trt_model.layers[i].attention.qkv.bias.value = torch_to_numpy( + trt_model.layers[i].attention.dense.bias.value = torch_to_numpy( split_bias) feed_weight_count += 1 - # Dense multiplication weight (no bias) - weight = hf_model.transformer.encoder.layers[ - i].self_attention.dense.weight.to(torch_type).detach().cpu() - split_weight = torch.chunk(weight, mapping.tp_size, - dim=1)[mapping.rank] - dst = trt_model.layers[i].attention.dense - if use_weight_only: - load_quant_weight( - src=split_weight, - value_dst=dst.weight, - scale_dst=dst.per_channel_scale, - plugin_weight_only_quant_type=plugin_weight_only_quant_type) - else: - dst.weight.value = torch_to_numpy(split_weight) - feed_weight_count += 1 - - # Post normalization + # Post normalization + if model_name in ["chatglm_6b"]: + weight = hf_model.transformer.layers[ + i].post_attention_layernorm.weight.to(torch_type).detach() + trt_model.layers[i].post_norm.weight.value = torch_to_numpy(weight) + bias = hf_model.transformer.layers[ + i].post_attention_layernorm.bias.to(torch_type).detach() + trt_model.layers[i].post_norm.bias.value = torch_to_numpy(bias) + feed_weight_count += 2 + elif model_name in [ + "chatglm2_6b", + "chatglm2_6b_32k", + "chatglm3_6b", + "chatglm3_6b_base", + "chatglm3_6b_32k", + ]: weight = hf_model.transformer.encoder.layers[ - i].post_attention_layernorm.weight.to( - torch_type).detach().cpu() + i].post_attention_layernorm.weight.to(torch_type).detach() trt_model.layers[i].post_norm.weight.value = torch_to_numpy(weight) feed_weight_count += 1 + elif model_name in ["glm_10b"]: + weight = hf_model.transformer.layers[ + i].post_attention_layernorm.weight.to(torch_type).detach() + trt_model.layers[i].post_norm.weight.value = torch_to_numpy(weight) + bias = hf_model.transformer.layers[ + i].post_attention_layernorm.bias.to(torch_type).detach() + trt_model.layers[i].post_norm.bias.value = torch_to_numpy(bias) + feed_weight_count += 2 - # Multilayer perceptron h -> 4h (no bias) + # Multilayer perceptron h -> 4h weight + if model_name in ["chatglm_6b"]: + weight = hf_model.transformer.layers[i].mlp.dense_h_to_4h.weight.to( + torch_type).detach() + split_weight = torch.chunk(weight, mapping.tp_size, + dim=0)[mapping.rank] + elif model_name in [ + "chatglm2_6b", + "chatglm2_6b_32k", + "chatglm3_6b", + "chatglm3_6b_base", + "chatglm3_6b_32k", + ]: weight = hf_model.transformer.encoder.layers[ - i].mlp.dense_h_to_4h.weight.to(torch_type).detach().cpu() + i].mlp.dense_h_to_4h.weight.to(torch_type).detach() split_weight = torch.chunk(weight, 2 * mapping.tp_size, dim=0) # swap first and second half weight in columns to adapt trt_llm Swiglu split_weight = torch.cat( @@ -331,36 +410,279 @@ def load_from_hf( ], dim=0, ) - dst = trt_model.layers[i].mlp.fc - if use_weight_only: - load_quant_weight( - src=split_weight, - value_dst=dst.weight, - scale_dst=dst.per_channel_scale, - plugin_weight_only_quant_type=plugin_weight_only_quant_type) - else: - dst.weight.value = torch_to_numpy(split_weight) + elif model_name in ["glm_10b"]: + weight = hf_model.transformer.layers[i].mlp.dense_h_to_4h.weight.to( + torch_type).detach() + split_weight = torch.chunk(weight, mapping.tp_size, + dim=0)[mapping.rank] + + dst = trt_model.layers[i].mlp.fc + if use_weight_only: + load_quant_weight( + src=split_weight, + value_dst=dst.weight, + scale_dst=dst.per_channel_scale, + plugin_weight_only_quant_type=plugin_weight_only_quant_type) + else: + dst.weight.value = torch_to_numpy(split_weight) + feed_weight_count += 1 + + # Multilayer perceptron h -> 4h bias, only GLM-10B + if model_name in ["glm_10b"]: + bias = hf_model.transformer.layers[i].mlp.dense_h_to_4h.bias.to( + torch_type).detach() + split_bias = split_qkv(bias, mapping.tp_size, mapping.tp_rank, + hidden_size, num_heads, num_kv_heads) + trt_model.layers[i].mlp.fc.bias.value = torch_to_numpy(split_bias) feed_weight_count += 1 - # Multilayer perceptron 4h -> h (no bias) + # Multilayer perceptron 4h -> h weight + if model_name in ["chatglm_6b"]: + weight = hf_model.transformer.layers[i].mlp.dense_4h_to_h.weight.to( + torch_type).detach() + elif model_name in [ + "chatglm2_6b", + "chatglm2_6b_32k", + "chatglm3_6b", + "chatglm3_6b_base", + "chatglm3_6b_32k", + ]: weight = hf_model.transformer.encoder.layers[ - i].mlp.dense_4h_to_h.weight.to(torch_type).detach().cpu() - split_weight = torch.chunk(weight, mapping.tp_size, - dim=1)[mapping.rank] - dst = trt_model.layers[i].mlp.proj - if use_weight_only: - load_quant_weight( - src=split_weight, - value_dst=dst.weight, - scale_dst=dst.per_channel_scale, - plugin_weight_only_quant_type=plugin_weight_only_quant_type) - else: - dst.weight.value = torch_to_numpy(split_weight) + i].mlp.dense_4h_to_h.weight.to(torch_type).detach() + elif model_name in ["glm_10b"]: + weight = hf_model.transformer.layers[i].mlp.dense_4h_to_h.weight.to( + torch_type).detach() + + split_weight = torch.chunk(weight, mapping.tp_size, dim=1)[mapping.rank] + dst = trt_model.layers[i].mlp.proj + if use_weight_only: + load_quant_weight( + src=split_weight, + value_dst=dst.weight, + scale_dst=dst.per_channel_scale, + plugin_weight_only_quant_type=plugin_weight_only_quant_type) + else: + dst.weight.value = torch_to_numpy(split_weight) + feed_weight_count += 1 + + # Multilayer perceptron 4h -> h bias, only GLM-10B + if model_name in ["glm_10b"]: + bias = hf_model.transformer.layers[i].mlp.dense_4h_to_h.bias.to( + torch_type).detach() + split_bias = split_qkv(bias, mapping.tp_size, mapping.tp_rank, + hidden_size, num_heads, num_kv_heads) + trt_model.layers[i].mlp.proj.bias.value = torch_to_numpy(split_bias) feed_weight_count += 1 - assert feed_weight_count == 3 + trt_model.num_layers * 7, "Some weights not loaded from HF" - + del hf_model tok = time.time() + # Final check + if model_name in ["chatglm_6b"]: + weight_count = 4 + num_layers * 9 + elif model_name in [ + "chatglm2_6b", + "chatglm2_6b_32k", + "chatglm3_6b", + "chatglm3_6b_base", + "chatglm3_6b_32k", + ]: + weight_count = 3 + num_layers * 7 + elif model_name in ["glm_10b"]: + weight_count = 6 + num_layers * 12 + if feed_weight_count < weight_count: + tensorrt_llm.logger.error("%d weights not loaded from HF" % + (weight_count - feed_weight_count)) + return None tensorrt_llm.logger.info("Loading weights finish in %.2fs" % (tok - tik)) return trt_model + + +def load_from_hf_checkpoint( + trtllm_falcon: tensorrt_llm.models.FalconForCausalLM, + model_dir: Union[str, Path], + mapping=Mapping(), + dtype: Union[str, torch.dtype] = torch.float32, +): + logger.info('Loading weights from HF Falcon...') + tik = time.time() + + model_dir = Path(model_dir) + if isinstance(dtype, str): + dtype = tensorrt_llm._utils.str_dtype_to_torch(dtype) + + def is_bias(_name): + return 'bias' in _name + + layers_range = trtllm_falcon.get_transformer_layers( + trtllm_falcon.mapping, trtllm_falcon.num_layers) + for model_file in iterate_shard_files(model_dir, mapping.tp_rank): + logger.debug(f'Loading file {str(model_file)}...') + state_dict = load_state_dict(model_file, dtype) + for name, param in state_dict.items(): + logger.debug(f'Converting weight {name}...') + i = retrieved_layer_index_from_name(name) + if i is None: + layer = None + else: + if i not in layers_range: + continue + layer = trtllm_falcon.layers[i - layers_range[0]] + + if 'self_attention.query_key_value' in name: + if not is_bias(name): + layer.attention.qkv.weight.value = split_qkv_weight( + trtllm_falcon, + param, + mapping.tp_size, + mapping.tp_rank, + is_bias=False, + num_kv_heads=trtllm_falcon.num_kv_heads) + else: + layer.attention.qkv.bias.value = split_qkv_weight( + trtllm_falcon, + param, + mapping.tp_size, + mapping.tp_rank, + is_bias=True, + num_kv_heads=trtllm_falcon.num_kv_heads) + elif 'self_attention.dense' in name: + if not is_bias(name): + layer.attention.dense.weight.value = split_matrix( + param, mapping.tp_size, mapping.tp_rank, dim=1) + else: + layer.attention.dense.bias.value = param + elif 'mlp.dense_h_to_4h' in name: + if not is_bias(name): + layer.mlp.fc.weight.value = split_matrix(param, + mapping.tp_size, + mapping.tp_rank, + dim=0) + else: + layer.mlp.fc.bias.value = split_matrix(param, + mapping.tp_size, + mapping.tp_rank, + dim=0) + elif 'mlp.dense_4h_to_h' in name: + if not is_bias(name): + layer.mlp.proj.weight.value = split_matrix(param, + mapping.tp_size, + mapping.tp_rank, + dim=1) + else: + layer.mlp.proj.bias.value = param + elif 'ln_attn' in name or 'input_layernorm' in name: + if not is_bias(name): + layer.input_layernorm.weight.value = param + else: + layer.input_layernorm.bias.value = param + elif 'ln_mlp' in name: + assert layer.mlp_layernorm is not None + if not is_bias(name): + layer.mlp_layernorm.weight.value = param + else: + layer.mlp_layernorm.bias.value = param + elif 'post_attention_layernorm' in name: + assert layer.post_layernorm is not None + if not is_bias(name): + layer.post_layernorm.weight.value = param + else: + layer.post_layernorm.bias.value = param + elif 'word_embeddings' in name: + if mapping.is_first_pp_rank(): + trtllm_falcon.embedding.weight.value = param.copy() + if mapping.is_last_pp_rank(): + trtllm_falcon.lm_head.weight.value = split_matrix( + param, mapping.tp_size, mapping.tp_rank, dim=0) + elif 'ln_f' in name: + if mapping.is_last_pp_rank(): + if not is_bias(name): + trtllm_falcon.ln_f.weight.value = param + else: + trtllm_falcon.ln_f.bias.value = param + del state_dict + + tok = time.time() + t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) + logger.info(f'Weights loaded. Total time: {t}') + + +def get_scaling_factors( + model_path: Union[str, Path], + num_layers: int, + quant_mode: Optional[QuantMode] = None, +) -> Optional[Dict[str, List[int]]]: + """ Get the scaling factors for Falcon model + + Returns a dictionary of scaling factors for the selected layers of the + Falcon model. + + Args: + model_path (str): Path to the quantized Falcon model + layers (list): List of layers to get the scaling factors for. If None, + all layers are selected. + + Returns: + dict: Dictionary of scaling factors for the selected layers of the + Falcon model. + + example: + + { + 'qkv_act': qkv_act_scale, + 'qkv_weights': qkv_weights_scale, + 'qkv_out' : qkv_outputs_scale, + 'dense_act': dense_act_scale, + 'dense_weights': dense_weights_scale, + 'fc_act': fc_act_scale, + 'fc_weights': fc_weights_scale, + 'proj_act': proj_act_scale, + 'proj_weights': proj_weights_scale, + } + """ + + if model_path is None: + logger.warning(f"--quantized_fp8_model_path not specified. " + f"Initialize quantization scales automatically.") + return get_dummy_quant_scales(num_layers) + weight_dict = np.load(model_path) + + # yapf: disable + scaling_factor = { + 'qkv_act': [], + 'qkv_weights': [], + 'qkv_output': [], + 'dense_act': [], + 'dense_weights': [], + 'fc_act': [], + 'fc_weights': [], + 'proj_act': [], + 'proj_weights': [], + } + + for layer in range(num_layers): + scaling_factor['qkv_act'].append(max( + weight_dict[f'_np:layers:{layer}:attention:qkv:q:activation_scaling_factor'].item(), + weight_dict[f'_np:layers:{layer}:attention:qkv:k:activation_scaling_factor'].item(), + weight_dict[f'_np:layers:{layer}:attention:qkv:v:activation_scaling_factor'].item() + )) + scaling_factor['qkv_weights'].append(max( + weight_dict[f'_np:layers:{layer}:attention:qkv:q:weights_scaling_factor'].item(), + weight_dict[f'_np:layers:{layer}:attention:qkv:k:weights_scaling_factor'].item(), + weight_dict[f'_np:layers:{layer}:attention:qkv:v:weights_scaling_factor'].item() + )) + if quant_mode is not None and quant_mode.has_fp8_kv_cache(): + # Not calibrarting KV cache. + scaling_factor['qkv_output'].append(1.0) + scaling_factor['dense_act'].append(weight_dict[f'_np:layers:{layer}:attention:dense:activation_scaling_factor'].item()) + scaling_factor['dense_weights'].append(weight_dict[f'_np:layers:{layer}:attention:dense:weights_scaling_factor'].item()) + scaling_factor['fc_act'].append(weight_dict[f'_np:layers:{layer}:mlp:fc:activation_scaling_factor'].item()) + scaling_factor['fc_weights'].append(weight_dict[f'_np:layers:{layer}:mlp:fc:weights_scaling_factor'].item()) + scaling_factor['proj_act'].append(weight_dict[f'_np:layers:{layer}:mlp:proj:activation_scaling_factor'].item()) + scaling_factor['proj_weights'].append(weight_dict[f'_np:layers:{layer}:mlp:proj:weights_scaling_factor'].item()) + # yapf: enable + for k, v in scaling_factor.items(): + assert len(v) == num_layers, \ + f'Expect scaling factor {k} of length {num_layers}, got {len(v)}' + + return scaling_factor diff --git a/examples/common/utils.py b/examples/common/utils.py new file mode 100644 index 000000000000..5db966dacc84 --- /dev/null +++ b/examples/common/utils.py @@ -0,0 +1,80 @@ +import re +from pathlib import Path +from typing import Dict, Optional, Union + +import torch + + +def load_state_dict( + file_path: Union[str, Path], + dtype: torch.dtype, + device: Optional[Union[str, torch.device]] = None, +) -> Dict[str, torch.Tensor]: + """ Load weights from model file + + `safetensors` or `pytorch binary` is supported + + # Args. + file_path: model file path, ends with .bin or .safetensors. + dtype: torch.dtype, data type. + device: torch device like, optional. If None, load to cpu. + # Returns. + Dict[str, torch.Tensor] + """ + file_path = Path(file_path) + assert isinstance(dtype, torch.dtype) + + if device is None: + device = 'cpu' + + model_params = {} + if file_path.suffix == '.safetensors': + # load from safetensors file + from safetensors import safe_open + with safe_open(file_path, framework='pt', device=device) as f: + for name in f.keys(): + model_params[name] = f.get_tensor(name).to(dtype).clone() + elif file_path.suffix == '.bin': + # load from pytorch bin file + state_dict = torch.load(file_path, map_location=device) + for name in state_dict: + model_params[name] = state_dict[name].to(dtype) + else: + raise NotImplementedError( + f'Support .safetensors or .bin files, but got {str(file_path)}') + return model_params + + +def retrieved_layer_index_from_name(name: str) -> Optional[int]: + # This method is a hacky function to retrieve the layer index from + # HF model. Most of HF models have similar naming convention but + # please check carefully before applying if this method works well + # on your target model. + res = re.search(r'\d+', name) + return int(res.group()) if res is not None else res + + +def iterate_shard_files(model_dir: Union[Path, str], + rank: int, + progress_bar: bool = True): + model_dir = Path(model_dir) + + # '.bin' or '.safetensors'. In case that both exist, '.safetensor' + # files will be loaded first. + shard_files = list(model_dir.glob('*.safetensors')) + if not shard_files: + # The model checkpoint is stored in .bin file. + shard_files = list(model_dir.glob('*.bin')) + + try: + import tqdm + if progress_bar: + # Show a progress bar per rank. + desc = f'Rank [{rank}] Loading weights' + shard_files = tqdm.tqdm(shard_files, desc=desc, position=rank) + + except ImportError: + pass + + for shard_file in shard_files: + yield shard_file diff --git a/examples/enc_dec/README.md b/examples/enc_dec/README.md index 2ea21a29bb09..da1e3126e009 100644 --- a/examples/enc_dec/README.md +++ b/examples/enc_dec/README.md @@ -4,44 +4,122 @@ This document shows how to build and run an Encoder-Decoder (Enc-Dec) model in T ## Overview -The TensorRT-LLM Enc-Dec implementation can be found in [tensorrt_llm/models/enc_dec/model.py](../../tensorrt_llm/models/enc_dec/model.py). The TensorRT-LLM Enc-Dec example code is located in [`examples/enc_dec`](./). There are two main files in that folder: +The TensorRT-LLM Enc-Dec implementation can be found in [tensorrt_llm/models/enc_dec/model.py](../../tensorrt_llm/models/enc_dec/model.py). The TensorRT-LLM Enc-Dec example code is located in [`examples/enc_dec`](./): * [`build.py`](./build.py) to build the [TensorRT](https://developer.nvidia.com/tensorrt) engine(s) needed to run the Enc-Dec model, - * [`run.py`](./run.py) to run the inference on an input text. + * [`run.py`](./run.py) to run the inference on an example input text. + * Enc-Dec models can have specific implementations, such as the popular T5 family (T5, mT5, Flan-T5) and BART family (BART, mBART). They are located under subfolders `/t5` and `/bart`, each containing: + * [`/hf_convert.py`](./t5/hf_convert.py) to convert weights from HuggingFace PyTorch format to TRT-LLM format, and split weights for multi-GPU inference, + * [`/weight.py`](./t5/weight.py) to map the converted & split weights to TRT-LLM model. ## Usage -The TensorRT-LLM Enc-Dec example code locates at [examples/enc_dec](./). It takes HF weights as input, and builds the corresponding TensorRT engines. For single GPU, there will be two TensorRT engines, one for Encoder and one for Decoder. +The TensorRT-LLM Enc-Dec example code locates at [examples/enc_dec](./). It takes HuggingFace model name as input, and builds the corresponding TensorRT engines. On each GPU, there will be two TensorRT engines, one for Encoder and one for Decoder. ## Encoder-Decoder Model Support + +The implementation is designed to support generic encoder-decoder models by abstracting the common and derivative components of different model architectures, such as: - [T5](https://huggingface.co/docs/transformers/main/en/model_doc/t5) +- [T5v1.1](https://huggingface.co/docs/transformers/model_doc/t5v1.1) and [Flan-T5](https://huggingface.co/docs/transformers/model_doc/flan-t5) +- [mT5 (coming)](https://huggingface.co/docs/transformers/model_doc/mt5) +- [UL2 (coming)](https://huggingface.co/docs/transformers/model_doc/ul2) and [Flan-UL2 (coming)](https://huggingface.co/docs/transformers/model_doc/flan-ul2) +- [BART (coming)](https://huggingface.co/docs/transformers/model_doc/bart) +- [mBART (coming)](https://huggingface.co/docs/transformers/model_doc/mbart) + +It also supports full Tensor Parallelism (TP), Pipeline Parallelism (PP), and a hybrid of the two. Currently, Fused Multi-Head Attention (FMHA) is not yet enabled for T5 family due to its relative attention design. + +In this example, we use T5 (`t5-small`) and Flan-T5 (`google/flan-t5-small`) to showcase TRT-LLM support on Enc-Dec models. + +### Download weights from HuggingFace Transformers +```bash +git clone https://huggingface.co/t5-small tmp/hf_models/t5-small +git clone https://huggingface.co/google/flan-t5-small tmp/hf_models/flan-t5-small +``` + +### Convert and Split Weights +The `/hf_convert.py` script converts weights from HuggingFace format to TRT-LLM format, and splits weights for multi-GPU inference. `--inference_tensor_para_size` specifies the number of GPUs for tensor parallelism during inference. + +It is fine to save one copy of converted weights at high precision, e.g. float32, if disk space allows. During the following engine building phase, engines of any inference precision can be built by weight dtype casting on the fly. Therefore, you can just keep one set of saved weights and build engines freely at different precisions, instead of saving weights for each inference precision. -In this example, we use T5 to showcase TRT-LLM support on Enc-Dec models. +After weight conversion, TensorRT-LLM converted weights and model configuration will be saved under `/` directory, which is the `--weight_dir` input path you should give to the **next** engine building phase. `X` is Tensor Parallelim size for distributed inference. + +```bash +python t5/hf_convert.py -i tmp/hf_models/t5-small -o tmp/trt_models/t5-small --weight_data_type float32 --inference_tensor_para_size +``` ### Build TensorRT engine(s) -Need to prepare the HuggingFace T5 checkpoint first by following the guides here https://huggingface.co/docs/transformers/main/en/model_doc/t5. +TensorRT-LLM builds TensorRT engine(s) with flexible controls on different types of optimizations. Note that these are just examples to demonstrate multi-GPU inference. For small models like T5-small, single GPU is usually sufficient. + +After engine building, TensorRT engines will be saved under `//` directory, which is the `--engine_dir` path you should give to the next engine running phase. It is recommended to have `/` in the output path where `Y` is number of total GPU ranks in a multi-node, multi-GPU setup, because the same `Y` number GPUs could be executed with different TP (Tensor Parallelism) and PP (Pipeline Parallelism) combinations. -TensorRT-LLM Enc-Dec builds TensorRT engine(s) from HF checkpoint. For the first time of running this example, user needs to download the T5 model ckpt from HF. After obtaining the HF T5 ckpt, user can build the TensorRT engines. +We should distinguish between `X` - TP size and `Y` - total number of GPU ranks: +* When `X = Y`, only TP is enabled +* When `X < Y`, both TP and PP are enabled. In such case, please make sure you have completed weight conversion step for `TP=X`. ```bash -# download t5-small ckpt to ./models (one-time) -python download.py +# Example 1: build t5-small using a single GPU, FP32, running gready search +# use_gpt_attention_plugin is necessary in Enc-Dec. +# Try use_gemm_plugin to prevent accuracy issue. +# It is recommend to use --remove_input_padding along with --use_gpt_attention_plugin for better performance +python build.py --model_type t5 \ + --weight_dir tmp/trt_models/t5-small/tp1 \ + -o tmp/trt_engines/t5-small/1-gpu \ + --engine_name t5-small \ + --remove_input_padding \ + --use_bert_attention_plugin \ + --use_gpt_attention_plugin \ + --use_gemm_plugin \ + --use_rmsnorm_plugin \ + --dtype float32 \ + --max_beam_width 1 + +# Example 2: build flan-t5-small using 4-way tensor parallelism on a node with 8 GPUs (but only use 4 of them, for demonstration purpose), BF16, enabling beam search up to width=3 +python build.py --model_type t5 \ + --world_size 4 \ + --tp_size 4 \ + --gpus_per_node 4 \ + --weight_dir tmp/trt_models/flan-t5-small/tp4 \ + -o tmp/trt_engines/flan-t5-small/4-gpu \ + --engine_name flan-t5-small \ + --remove_input_padding \ + --use_bert_attention_plugin \ + --use_gpt_attention_plugin \ + --use_gemm_plugin \ + --use_rmsnorm_plugin \ + --dtype bfloat16 \ + --max_beam_width 3 -# Build t5-small using a single GPU and FP16, supporting beam search up to 3 beam_width -python build.py --model_dir ./models/ \ +# Example 3: build flan-t5-small using 2-way tensor parallelism and 2-way pipeline parallelism on a node with 8 GPUs, FP16, enabling beam search up to width=3 +python build.py --model_type t5 \ + --world_size 4 \ + --tp_size 2 \ + --pp_size 2 \ + --gpus_per_node 8 \ + --weight_dir tmp/trt_models/flan-t5-small/tp2 \ + -o tmp/trt_engines/flan-t5-small/4-gpu \ + --engine_name flan-t5-small \ + --remove_input_padding \ --use_bert_attention_plugin \ --use_gpt_attention_plugin \ + --use_gemm_plugin \ + --use_rmsnorm_plugin \ --dtype float16 \ --max_beam_width 3 -# build.py will by default save the TRT engines into ./trt_engines ``` ### Run -To run a TensorRT-LLM Enc-Dec model using the engines generated by build.py +Run a TensorRT-LLM Enc-Dec model using the engines generated by build.py. +Note that during model deployment, only the TensorRT engine files are needed. Previously downloaded model checkpoints and converted weights can be removed. ```bash -# Run inference with beam search -python3 run.py --max_new_token=64 --num_beams=3 +# Example 1: inference w/ single GPU, FP32, greedy search, compare results with HuggingFace FP32 +python3 run.py --engine_dir tmp/trt_engines/t5-small/1-gpu/float32/tp1 --engine_name t5-small --model_name t5-small --max_new_token=64 --num_beams=1 --compare_hf_fp32 + +# Example 2: inference w/ 4 GPUs (4-way TP, as configured during the engine building step), BF16, greedy search +mpirun --allow-run-as-root -np 4 python3 run.py --engine_dir tmp/trt_engines/flan-t5-small/4-gpu/bfloat16/tp4 --engine_name flan-t5-small --model_name google/flan-t5-small --max_new_token=64 --num_beams=1 + +# Example 3: inference w/ 4 GPUs (2-way TP and 2-way PP, as configured during the engine building step), FP16, beam search +mpirun --allow-run-as-root -np 4 python3 run.py --engine_dir tmp/trt_engines/flan-t5-small/4-gpu/float16/tp2 --engine_name flan-t5-small --model_name google/flan-t5-small --max_new_token=64 --num_beams=3 ``` diff --git a/examples/enc_dec/build.py b/examples/enc_dec/build.py index 4f0faeb8b14f..cb43e7f4d268 100644 --- a/examples/enc_dec/build.py +++ b/examples/enc_dec/build.py @@ -1,22 +1,26 @@ import argparse +import configparser import time from pathlib import Path import torch +import torch.multiprocessing as mp import tensorrt_llm from tensorrt_llm._utils import str_dtype_to_trt from tensorrt_llm.builder import Builder from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping from tensorrt_llm.network import net_guard -from weight import load_t5_from_pytorch, parse_config # isort:skip +from t5.weight import parse_t5_config, load_from_hf_t5, load_from_binary_t5 # isort:skip -MODEL_NAME = "enc_dec" - -def get_engine_name(model, dtype, tp_size, rank): - return '{}_{}_tp{}_rank{}.engine'.format(model, dtype, tp_size, rank) +def get_engine_name(model, dtype, tp_size, pp_size, rank): + if pp_size == 1: + return '{}_{}_tp{}_rank{}.engine'.format(model, dtype, tp_size, rank) + return '{}_{}_tp{}_pp{}_rank{}.engine'.format(model, dtype, tp_size, + pp_size, rank) def serialize_engine(engine, path): @@ -29,17 +33,61 @@ def serialize_engine(engine, path): logger.info(f'Engine serialized. Total time: {t}') -def parse_arguments(args, component): +def parse_config(ini_file, component, args): + config = configparser.ConfigParser() + config.read(ini_file) + model_type = config.get('structure', 'model_type') + args = globals()[f'parse_{model_type}_config'](config, component, args) + return args + + +def parse_arguments(component): parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument('--dtype', + parser.add_argument('--world_size', + type=int, + default=1, + help='MPI world size (must equal TP * PP)') + 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( + '--gpus_per_node', + type=int, + default=8, + help= + 'Number of GPUs each node has in a multi-node setup. This is a cluster spec and can be greater/smaller than world size' + ) + parser.add_argument('--parallel_build', default=False, action='store_true') + parser.add_argument('--weight_dir', + '-i', type=str, - default='float16', - choices=['float16', 'float32', 'bfloat16']) - parser.add_argument('--logits_dtype', + default=None, + help='Path to the converted weight file') + parser.add_argument( + '--output_dir', + '-o', + type=Path, + default='trt_engines', + help= + 'The path to save the serialized engine files, timing cache file and model configs' + ) + parser.add_argument( + '--weight_from_pytorch_ckpt', + default=False, + action='store_true', + help= + 'Load weight from PyTorch checkpoint. model_dir must point to ckpt directory' + ) + parser.add_argument('--engine_name', + '-n', type=str, - default='float32', - choices=['float16', 'float32']) + default='enc_dec', + help='TensorRT engine name prefix') parser.add_argument( '--timing_cache', type=str, @@ -47,18 +95,34 @@ def parse_arguments(args, component): help= 'The path of to read timing cache from, will be ignored if the file does not exist' ) + + parser.add_argument('--model_type', + type=str, + choices=['t5', 'bart'], + default='t5') + parser.add_argument( + '--dtype', + type=str, + default='float16', + choices=['float16', 'float32', 'bfloat16'], + help= + 'Target inference dtype. Weights and Computation will be in this dtype, no matter what original dtype the weight checkpoint has.' + ) + parser.add_argument('--logits_dtype', + type=str, + default='float32', + choices=['float16', 'float32']) + parser.add_argument('--log_level', type=str, default='info') - parser.add_argument('--vocab_size', type=int, default=32128) - parser.add_argument('--n_layer', type=int, default=6) - parser.add_argument('--n_positions', type=int, default=1024) - parser.add_argument('--n_embd', type=int, default=1024) - parser.add_argument('--n_head', type=int, default=8) - parser.add_argument('--hidden_act', type=str, default='gelu') - parser.add_argument('--inter_size', type=int, default=None) - parser.add_argument('--no_bias', action="store_false") parser.add_argument('--max_batch_size', type=int, default=8) parser.add_argument('--max_encoder_input_len', type=int, default=1024) - parser.add_argument('--max_input_len', type=int, default=200) + parser.add_argument( + '--max_decoder_input_len', + type=int, + default=1, + help= + 'If you want deocder_forced_input_ids feature, set to value greater than 1. Otherwise, encoder-decoder model start from decoder_start_token_id of length 1' + ) parser.add_argument('--max_output_len', type=int, default=200) parser.add_argument('--max_beam_width', type=int, default=1) parser.add_argument( @@ -101,18 +165,27 @@ def parse_arguments(args, component): help= "Activates layernorm plugin. You can specify the plugin dtype or leave blank to use the model dtype." ) + parser.add_argument( + '--use_rmsnorm_plugin', + nargs='?', + const=None, + type=str, + default=False, + choices=['float16', 'float32', 'bfloat16'], + help= + "Activates rmsnorm plugin. You can specify the plugin dtype or leave blank to use the model dtype." + ) + parser.add_argument( + '--use_lookup_plugin', + nargs='?', + const=None, + default=False, + choices=['float16', 'float32', 'bfloat16'], + help="Activates the lookup plugin which enables embedding sharding.") parser.add_argument('--enable_qk_half_accum', default=False, action='store_true') - parser.add_argument('--gpus_per_node', type=int, default=8) parser.add_argument('--builder_opt', type=int, default=None) - parser.add_argument( - '--output_dir', - type=Path, - default='trt_engines', - help= - 'The path to save the serialized engine files, timing cache file and model configs' - ) parser.add_argument('--remove_input_padding', default=False, action='store_true') @@ -123,13 +196,27 @@ def parse_arguments(args, component): help= 'Seed to use when initializing the random number generator for torch.') parser.add_argument( - '--use_lookup_plugin', - nargs='?', - const=None, + '--use_parallel_embedding', + action="store_true", default=False, - choices=['float16', 'float32', 'bfloat16'], - help="Activates the lookup plugin which enables embedding sharing.") - + 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 sharding is only enabled when embedding_sharding_dim = 0' + ) + parser.add_argument( + '--use_custom_all_reduce', + action='store_true', + help= + 'Activates latency-optimized algorithm for all-reduce instead of NCCL.') parser.add_argument( '--strongly_typed', default=False, @@ -138,20 +225,22 @@ def parse_arguments(args, component): 'This option is introduced with trt 9.1.0.1+ and will reduce the building time significantly for fp8.' ) - args = parser.parse_args(args) + # parse cmdline args + args = parser.parse_args() logger.set_level(args.log_level) - args.bias = not args.no_bias - if args.inter_size is None: - args.inter_size = 4 * args.n_embd - - if args.model_dir is not None: - logger.info(f"Setting model configuration from {args.model_dir}.") + # parse model config and add to args + if args.weight_dir is not None: + logger.info(f"Setting model configuration from {args.weight_dir}.") args = parse_config( - Path(args.model_dir) / "config.ini", component, args) + Path(args.weight_dir) / "config.ini", component, args) + + assert args.pp_size * args.tp_size == args.world_size + plugins_args = [ 'use_bert_attention_plugin', 'use_gpt_attention_plugin', - 'use_gemm_plugin', 'use_layernorm_plugin', 'use_lookup_plugin' + 'use_gemm_plugin', 'use_layernorm_plugin', 'use_rmsnorm_plugin', + 'use_lookup_plugin' ] for plugin_arg in plugins_args: if getattr(args, plugin_arg) is None: @@ -160,6 +249,9 @@ def parse_arguments(args, component): ) setattr(args, plugin_arg, args.dtype) + if args.dtype == 'bfloat16': + assert args.use_gemm_plugin, "Please use gemm plugin when dtype is bfloat16" + return args @@ -172,13 +264,23 @@ def build_rank_engine(builder: Builder, @param args: The cmd line arguments. @return: The built engine. ''' - kv_dtype = str_dtype_to_trt(args.dtype) + dtype = str_dtype_to_trt(args.dtype) + + mapping = Mapping(world_size=args.world_size, + rank=rank, + tp_size=args.tp_size, + pp_size=args.pp_size) + + assert args.n_layer % args.pp_size == 0, \ + f"num_layers {args.n_layer} must be a multiple of pipeline parallelism size {args.pp_size}" # Initialize Module if args.component == 'encoder': tllm_model = tensorrt_llm.models.EncoderModel( num_layers=args.n_layer, num_heads=args.n_head, + num_kv_heads=args.n_head, + head_size=args.head_size, hidden_size=args.hidden_size, ffn_hidden_size=args.ffn_hidden_size, vocab_size=args.vocab_size, @@ -197,15 +299,22 @@ def build_rank_engine(builder: Builder, layernorm_position=args.layernorm_position, layernorm_type=args.layernorm_type, hidden_act=args.hidden_act, - dtype=kv_dtype) + mlp_type=args.mlp_type, + dtype=dtype, + use_parallel_embedding=args.use_parallel_embedding, + embedding_sharding_dim=args.embedding_sharding_dim, + mapping=mapping) elif args.component == 'decoder': tllm_model = tensorrt_llm.models.DecoderModel( num_layers=args.n_layer, num_heads=args.n_head, + num_kv_heads=args.n_head, + head_size=args.head_size, hidden_size=args.hidden_size, ffn_hidden_size=args.ffn_hidden_size, encoder_hidden_size=args.encoder_hidden_size, encoder_num_heads=args.encoder_num_heads, + encoder_head_size=args.encoder_head_size, vocab_size=args.vocab_size, max_position_embeddings=args.n_positions, has_position_embedding=args.has_position_embedding, @@ -222,21 +331,32 @@ def build_rank_engine(builder: Builder, layernorm_position=args.layernorm_position, layernorm_type=args.layernorm_type, hidden_act=args.hidden_act, - dtype=kv_dtype, - logits_dtype=args.logits_dtype) - - # No support for relative attention bias in plain TRT mode + mlp_type=args.mlp_type, + dtype=dtype, + logits_dtype=args.logits_dtype, + use_parallel_embedding=args.use_parallel_embedding, + embedding_sharding_dim=args.embedding_sharding_dim, + mapping=mapping) + + # No support for relative attention bias in plain TRT mode. Please use attention plugin # (If to add such support, need to add into # Attention and BertAttention at tensorrt_llm/layers/attention.py) if args.relative_attention: assert args.use_bert_attention_plugin, "Relative attention bias is only supported when using BertAttention Plugin" assert args.use_gpt_attention_plugin, "Relative attention bias is only supported when using GPTAttention Plugin" - if args.model_dir is not None: - load_t5_from_pytorch(tllm_model, - args.model_dir, - args.component, - dtype=args.dtype) + if args.weight_from_pytorch_ckpt: + assert args.tp_size == 1, "Loading from framework model via memory is for demonstration purpose. For multi-GPU inference, please use loading from binary for better performance." + globals()[f'load_from_hf_{args.model_type}'](tllm_model, + args.weight_dir, + args.component, + dtype=args.dtype) + else: + globals()[f'load_from_binary_{args.model_type}'](tllm_model, + args.weight_dir, + args, + mapping=mapping, + dtype=args.dtype) # Module -> Network network = builder.create_network() @@ -252,14 +372,18 @@ def build_rank_engine(builder: Builder, if args.use_layernorm_plugin: network.plugin_config.set_layernorm_plugin( dtype=args.use_layernorm_plugin) + if args.use_rmsnorm_plugin: + network.plugin_config.set_rmsnorm_plugin(dtype=args.use_rmsnorm_plugin) if args.enable_qk_half_accum: network.plugin_config.enable_qk_half_accum() if args.remove_input_padding: network.plugin_config.enable_remove_input_padding() - if args.use_lookup_plugin: - # Use the plugin for the embedding parallelism and sharing + # Use the plugin for the embedding parallelism and sharding network.plugin_config.set_lookup_plugin(dtype=args.dtype) + if args.world_size > 1: + network.plugin_config.set_nccl_plugin(args.dtype, + args.use_custom_all_reduce) with net_guard(network): # Prepare @@ -269,14 +393,13 @@ def build_rank_engine(builder: Builder, if args.component == 'encoder': inputs = tllm_model.prepare_inputs( args.max_batch_size, - args.max_input_len, + args.max_encoder_input_len, ) elif args.component == 'decoder': inputs = tllm_model.prepare_inputs( - args.n_layer, args.max_batch_size, args.max_beam_width, - args.max_input_len, + args.max_decoder_input_len, args.max_output_len, args.max_encoder_input_len, ) @@ -291,9 +414,6 @@ def build_rank_engine(builder: Builder, # Network -> Engine engine = builder.build_engine(network, builder_config) - if rank == 0: - config_path = args.output_dir / args.component / 'config.json' - builder.save_config(builder_config, config_path) tensorrt_llm.tools.cleanup(network, tllm_model) @@ -303,58 +423,72 @@ def build_rank_engine(builder: Builder, def build(rank, args): torch.cuda.set_device(rank % args.gpus_per_node) tensorrt_llm.logger.set_level(args.log_level) - component_dir = args.output_dir / args.component + component_dir = args.output_dir / args.dtype / f"tp{args.tp_size}" / args.component component_dir.mkdir(parents=True, exist_ok=True) - timing_cache_file = args.timing_cache if args.timing_cache else component_dir / "model.cache" - timing_cache = timing_cache_file builder = Builder() apply_query_key_layer_scaling = False - # Currently only support single GPU - world_size = 1 - for cur_rank in range(world_size): + cache = None + for cur_rank in range(args.world_size): + # skip other ranks if parallel_build is enabled + if args.parallel_build and cur_rank != rank: + continue builder_config = builder.create_builder_config( - name=MODEL_NAME, + name=args.engine_name, precision=args.dtype, - timing_cache=timing_cache, - tensor_parallel=world_size, # TP only + timing_cache=component_dir / + args.timing_cache if cache is None else cache, + tensor_parallel=args.tp_size, + pipeline_parallel=args.pp_size, + gpus_per_node=args.gpus_per_node, + parallel_build=args.parallel_build, num_layers=args.n_layer, num_heads=args.n_head, hidden_size=args.hidden_size, + head_size=args.head_size, vocab_size=args.vocab_size, hidden_act=args.hidden_act, max_position_embeddings=args.n_positions, apply_query_key_layer_scaling=apply_query_key_layer_scaling, max_batch_size=args.max_batch_size, - max_input_len=args.max_input_len, + max_beam_width=args.max_beam_width, + max_decoder_input_len=args.max_decoder_input_len, max_output_len=args.max_output_len, + max_encoder_input_len=args.max_encoder_input_len, opt_level=args.builder_opt, cross_attention=(args.component == 'decoder'), has_position_embedding=args.has_position_embedding, has_token_type_embedding=args.has_token_type_embedding, strongly_typed=args.strongly_typed) - engine_name = get_engine_name(MODEL_NAME, args.dtype, world_size, - cur_rank) + engine_name = get_engine_name(args.engine_name, args.dtype, + args.tp_size, args.pp_size, cur_rank) engine = build_rank_engine(builder, builder_config, engine_name, cur_rank, args) assert engine is not None, f'Failed to build engine for rank {cur_rank}' if cur_rank == 0: + # save build config + config_path = component_dir / 'config.json' + builder.save_config(builder_config, config_path) + # Use in-memory timing cache for multiple builder passes. - timing_cache = builder_config.trt_builder_config.get_timing_cache() + if not args.parallel_build: + cache = builder_config.trt_builder_config.get_timing_cache() serialize_engine(engine, component_dir / engine_name) if rank == 0: - ok = builder.save_timing_cache(builder_config, timing_cache_file) + # save timing cache to speedup future use + ok = builder.save_timing_cache(builder_config, + component_dir / args.timing_cache) assert ok, "Failed to save timing cache." -def run_build(component, args=None): +def run_build(component): assert component == 'encoder' or component == 'decoder', 'Unsupported component!' - args = parse_arguments(args, component) + args = parse_arguments(component) args.component = component if args.random_seed is not None: @@ -363,9 +497,16 @@ def run_build(component, args=None): logger.set_level(args.log_level) tik = time.time() - # Currently only support single GPU serial build - logger.info('Serially build TensorRT engines.') - build(0, args) + if args.parallel_build and args.world_size > 1 and \ + torch.cuda.device_count() >= args.world_size: + logger.warning( + f'Parallelly build TensorRT engines. Please make sure that all of the {args.world_size} GPUs are totally free.' + ) + mp.spawn(build, nprocs=args.world_size, args=(args, )) + else: + args.parallel_build = False + logger.info('Serially build TensorRT engines.') + build(0, args) tok = time.time() t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) diff --git a/examples/enc_dec/download.py b/examples/enc_dec/download.py deleted file mode 100644 index 44a3d44abb2b..000000000000 --- a/examples/enc_dec/download.py +++ /dev/null @@ -1,18 +0,0 @@ -import torch -from transformers import T5ForConditionalGeneration, T5Tokenizer - -tokenizer = T5Tokenizer.from_pretrained("t5-small") -model = T5ForConditionalGeneration.from_pretrained("t5-small") - -input_ids = tokenizer("translate English to German: The house is wonderful.", - return_tensors="pt").input_ids -outputs = model.generate(input_ids, decoder_input_ids=torch.IntTensor([[ - 0, -]])) -print("input", input_ids, "\noutput", outputs) -print(tokenizer.decode(outputs[0], skip_special_tokens=True)) - -torch.save(model.state_dict(), './models/t5_small.ckpt') - -for k, v in model.state_dict().items(): - print(k) diff --git a/examples/enc_dec/models/config.ini b/examples/enc_dec/models/config.ini deleted file mode 100644 index c9ac53e7d2a4..000000000000 --- a/examples/enc_dec/models/config.ini +++ /dev/null @@ -1,48 +0,0 @@ -[encoder] -n_layer = 6 -n_head = 8 -hidden_size = 512 -ffn_hidden_size = 2048 -vocab_size = 32128 -n_positions = 1024 -has_position_embedding = False -has_token_type_embedding = False -has_embedding_layernorm = False -has_embedding_scale = False -q_scaling = 0.125 -has_attention_qkvo_bias = False -has_mlp_bias = False -has_model_final_layernorm = True -layernorm_eps = 1e-6 -layernorm_position = pre_layernorm -layernorm_type = RmsNorm -hidden_act = relu -relative_attention = True -num_buckets = 32 -max_distance = 128 -storage_dtype = float32 - -[decoder] -n_layer = 6 -n_head = 8 -hidden_size = 512 -ffn_hidden_size = 2048 -vocab_size = 32128 -n_positions = 1024 -has_position_embedding = False -has_token_type_embedding = False -has_embedding_layernorm = False -has_embedding_scale = False -q_scaling = 0.125 -has_attention_qkvo_bias = False -has_mlp_bias = False -has_model_final_layernorm = True -layernorm_eps = 1e-6 -layernorm_position = pre_layernorm -layernorm_type = RmsNorm -hidden_act = relu -has_lm_head_bias = False -relative_attention = True -num_buckets = 32 -max_distance = 128 -storage_dtype = float32 diff --git a/examples/enc_dec/run.py b/examples/enc_dec/run.py index f5e4c1d814e9..93be31dbb943 100644 --- a/examples/enc_dec/run.py +++ b/examples/enc_dec/run.py @@ -1,12 +1,14 @@ import argparse import json +import time from pathlib import Path import tensorrt as trt import torch -from transformers import AutoTokenizer, T5ForConditionalGeneration +from transformers import AutoConfig, AutoTokenizer, T5ForConditionalGeneration import tensorrt_llm +from tensorrt_llm import logger from tensorrt_llm._utils import trt_dtype_to_torch from tensorrt_llm.runtime import ModelConfig, SamplingConfig @@ -30,24 +32,37 @@ def read_config(config_path: Path): config = json.load(f) use_gpt_attention_plugin = config["plugin_config"]["gpt_attention_plugin"] remove_input_padding = config["plugin_config"]["remove_input_padding"] - world_size = config["builder_config"]["tensor_parallel"] - 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["builder_config"]["num_heads"] // world_size - hidden_size = config["builder_config"]["hidden_size"] // world_size + tp_size = config['builder_config']['tensor_parallel'] + pp_size = config['builder_config']['pipeline_parallel'] + gpus_per_node = config['builder_config']['gpus_per_node'] + 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["builder_config"]["num_heads"] + hidden_size = config["builder_config"]["hidden_size"] + head_size = config["builder_config"]["head_size"] vocab_size = config["builder_config"]["vocab_size"] num_layers = config["builder_config"]["num_layers"] + num_kv_heads = config['builder_config'].get('num_kv_heads', num_heads) + + assert (num_heads % tp_size) == 0 + num_heads = num_heads // tp_size + hidden_size = hidden_size // tp_size + num_kv_heads = (num_kv_heads + tp_size - 1) // tp_size + cross_attention = config["builder_config"]["cross_attention"] has_position_embedding = config["builder_config"]["has_position_embedding"] has_token_type_embedding = config["builder_config"][ "has_token_type_embedding"] - num_kv_heads = num_heads + use_custom_all_reduce = config['plugin_config'].get('use_custom_all_reduce', + False) + dtype = config["builder_config"]["precision"] model_config = ModelConfig( num_heads=num_heads, num_kv_heads=num_kv_heads, hidden_size=hidden_size, + head_size=head_size, vocab_size=vocab_size, num_layers=num_layers, gpt_attention_plugin=use_gpt_attention_plugin, @@ -55,69 +70,97 @@ def read_config(config_path: Path): cross_attention=cross_attention, has_position_embedding=has_position_embedding, has_token_type_embedding=has_token_type_embedding, - ) + use_custom_all_reduce=use_custom_all_reduce, + dtype=dtype) - dtype = config["builder_config"]["precision"] - max_input_len = config["builder_config"]["max_input_len"] - - return model_config, world_size, dtype, max_input_len + return model_config, tp_size, pp_size, gpus_per_node, dtype def parse_arguments(): parser = argparse.ArgumentParser() - parser.add_argument("--max_new_tokens", type=int, required=True) + parser.add_argument("--max_new_tokens", type=int, default=64) + parser.add_argument('--max_kv_cache_len', + type=int, + default=None, + help='The max kv cache length. \ + If the final sequence length exceeds the kv cache length, we will enable cyclic kv cache. \ + If it is set to None, we will use the max sequence length.') parser.add_argument("--log_level", type=str, default="error") - parser.add_argument("--engine_dir", type=str, default="trt_engines") + parser.add_argument("--engine_dir", "-i", type=str, default="trt_engines") parser.add_argument("--engine_name", type=str, default="enc_dec") - parser.add_argument("--tokenizer", + parser.add_argument("--model_name", type=str, - help="HF tokenizer config path", + help="HuggingFace model name", default="t5-small") parser.add_argument("--num_beams", type=int, help="Use beam search if num_beams >1", default=1) parser.add_argument("--debug_mode", - type=bool, help="Whether or not to turn on the debug mode", - default=False) + action='store_true') + parser.add_argument("--compare_hf_fp32", + help="Compare results with HuggingFace FP32", + action='store_true') return parser.parse_args() class TRTLLMEncDecModel: def __init__(self, engine_name, engine_dir, debug_mode=False): + # in multi-node setup, it's important to set_device at the very beginning so .to('cuda') refers to current device + # accordingly, all input & output tensors should be moved to current device + # otherwise, it's default to 'cuda:0' + self.runtime_rank = tensorrt_llm.mpi_rank() + device_id = self.runtime_rank % torch.cuda.device_count() + torch.cuda.set_device(device_id) + self.device = torch.cuda.current_device() + engine_dir = Path(engine_dir) - # model config - encoder_config_path = engine_dir / "encoder" / "config.json" - encoder_model_config, world_size, dtype, max_input_len = read_config( - encoder_config_path) - decoder_config_path = engine_dir / "decoder" / "config.json" - decoder_model_config, _, _, _ = read_config(decoder_config_path) - self.encoder_model_config = encoder_model_config - self.decoder_model_config = decoder_model_config - - # MGMN config - runtime_rank = tensorrt_llm.mpi_rank() - runtime_mapping = tensorrt_llm.Mapping(world_size, runtime_rank) - torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node) - - # load engine - engine_name = get_engine_name(engine_name, dtype, world_size, - runtime_rank) - with open(engine_dir / "encoder" / engine_name, "rb") as f: - encoder_engine_buffer = f.read() - with open(engine_dir / "decoder" / engine_name, "rb") as f: - decoder_engine_buffer = f.read() + def engine_setup(component): + # model config + config_path = engine_dir / component / "config.json" + model_config, tp_size, pp_size, gpus_per_node, dtype = read_config( + config_path) + + # MGMN config + world_size = tp_size * pp_size + runtime_rank = tensorrt_llm.mpi_rank() + assert runtime_rank < world_size, "Runtime GPU rank exceeds MPI world size. Did you launch more MPI processes than required?" + runtime_mapping = tensorrt_llm.Mapping(world_size, + runtime_rank, + tp_size=tp_size, + pp_size=pp_size, + gpus_per_node=gpus_per_node) + + # load engine + engine_fname = get_engine_name(engine_name, dtype, tp_size, pp_size, + runtime_rank) + with open(engine_dir / component / engine_fname, "rb") as f: + engine_buffer = f.read() + + return model_config, runtime_mapping, engine_buffer + + # Note: encoder and decoder doesn't necessarily have the same TP & PP config + self.encoder_model_config, self.encoder_runtime_mapping, encoder_engine_buffer = engine_setup( + component='encoder') + self.decoder_model_config, self.decoder_runtime_mapping, decoder_engine_buffer = engine_setup( + component='decoder') + + # for Pipeline Parallelism in encoder + self.nccl_comm = torch.classes.FasterTransformer.NcclCommunicatorOp( + self.encoder_runtime_mapping.tp_size, + self.encoder_runtime_mapping.pp_size, + self.encoder_runtime_mapping.rank) # session setup self.encoder_session = tensorrt_llm.runtime.Session.from_serialized_engine( encoder_engine_buffer) self.decoder_session = tensorrt_llm.runtime.GenerationSession( - decoder_model_config, + self.decoder_model_config, decoder_engine_buffer, - runtime_mapping, + self.decoder_runtime_mapping, debug_mode=debug_mode) self.stream = torch.cuda.current_stream().cuda_stream @@ -125,39 +168,84 @@ def __init__(self, engine_name, engine_dir, debug_mode=False): def from_engine(cls, engine_name, engine_dir, debug_mode=False): return cls(engine_name, engine_dir, debug_mode=debug_mode) + def process_input(self, + input_ids, + remove_input_padding=False, + pad_token_id=0): + if remove_input_padding: + # in remove padding mode --> flatten input, calculate actual length and max length + # Note: 1st token should never be removed, even if it is pad_token_id + first_ids = input_ids[:, 0] + input_ids = input_ids[:, 1:] + input_lengths = 1 + (input_ids != pad_token_id).sum(dim=1).type( + torch.IntTensor).to(self.device) # [batch_size] + new_ids = [] + for i in range(len(input_ids)): + row = input_ids[i, :] + row = row[row != pad_token_id] + new_ids.append( + torch.cat( + (torch.IntTensor([first_ids[i]]).to(self.device), row))) + input_ids = torch.cat(new_ids).unsqueeze(dim=0) # [1, num_tokens] + else: + # in padding mode --> keep input, just calculate actual length and max length + # Note: 1st token should always count, even if it is pad_token_id. e.g., decoder start id in enc-dec models could be a single pad_token_id, we should count + input_lengths = torch.tensor( + 1 + (input_ids[:, 1:] != pad_token_id).sum(dim=1).type( + torch.IntTensor).to(self.device), + dtype=torch.int32, + device=self.device) + max_input_length = torch.max(input_lengths).item() + return input_ids, input_lengths, max_input_length + def encoder_run(self, input_ids, + input_lengths, + max_input_length, position_ids=None, token_type_ids=None, debug_mode=False): - batch_size = input_ids.shape[0] - input_lengths = torch.tensor([len(x) for x in input_ids], - dtype=torch.int32, - device='cuda') - max_input_length = torch.max(input_lengths).item() - # set input tensors and shapes - inputs = { - 'input_ids': input_ids, - 'input_lengths': input_lengths, - } - if self.encoder_model_config.has_position_embedding: - inputs['position_ids'] = position_ids - if self.encoder_model_config.has_token_type_embedding: - inputs['token_type_ids'] = token_type_ids - for k, v in inputs.items(): - self.encoder_session.context.set_input_shape(k, v.shape) - - # set output tensors and shapes - outputs = { - 'encoder_output': - torch.empty((batch_size, max_input_length, - self.encoder_model_config.hidden_size), - dtype=trt_dtype_to_torch( - self.encoder_session.engine.get_tensor_dtype( - 'encoder_output')), - device='cuda') - } + # each engine has hidden_dim/TP, don't forget to multiply TP + hidden_size = self.encoder_model_config.hidden_size * self.encoder_runtime_mapping.tp_size + hidden_states_shape = (input_ids.shape[0], input_ids.shape[1], + hidden_size) # [1,num_tokens,D] or [BS,seqlen,D] + hidden_states_dtype = lambda name: trt_dtype_to_torch( + self.encoder_session.engine.get_tensor_dtype(name)) + + # input tensors. only first PP rank has id input, others are hidden_states input + inputs = {} + if self.encoder_runtime_mapping.is_first_pp_rank(): + inputs['input_ids'] = input_ids.contiguous() + if self.encoder_model_config.has_position_embedding: + inputs['position_ids'] = position_ids.contiguous() + if self.encoder_model_config.has_token_type_embedding: + inputs['token_type_ids'] = token_type_ids.contiguous() + else: + # just need a placeholder, engine will call NCCL to recv and fill data from previous rank + inputs['hidden_states_input'] = torch.empty( + hidden_states_shape, + dtype=hidden_states_dtype('hidden_states_input'), + device=self.device).contiguous() + inputs['input_lengths'] = input_lengths + # use shape info to pass max length info in remove padding mode + inputs['max_input_length'] = torch.empty( + (max_input_length, ), + dtype=hidden_states_dtype('max_input_length'), + device=self.device).contiguous() + + # output tensors. only last PP rank final encoder output, others are intermediate hidden_states output. Need broadcast later + outputs = {} + if self.encoder_runtime_mapping.is_last_pp_rank(): + outputs['encoder_output'] = torch.empty( + hidden_states_shape, + dtype=hidden_states_dtype('encoder_output'), + device=self.device).contiguous() + else: + outputs['hidden_states_output'] = torch.empty( + hidden_states_shape, + dtype=hidden_states_dtype('hidden_states_output'), + device=self.device).contiguous() # ------------------------------------------- if debug_mode: @@ -173,15 +261,40 @@ def encoder_run(self, shape = context.get_tensor_shape(name) outputs[name] = torch.zeros(tuple(shape), dtype=trt_dtype_to_torch(dtype), - device='cuda') + device=self.device) context.set_tensor_address(name, outputs[name].data_ptr()) # ------------------------------------------- # TRT session run + # Note: runtime.Session's run() method will set input/output tensor address, here we only need to provide tensor shape + self.encoder_session.set_shapes(inputs) ok = self.encoder_session.run(inputs, outputs, self.stream) assert ok, "Runtime execution failed" torch.cuda.synchronize() + # Tensor Parallelism is handled by model/engine definition + # But we need to broadcast among PP group at the end of encoder's Pipeline Parallelism + # After this, all ranks should recv the encoder output, and world might be re-configured using decoder's TP-PP config + def pp_communicate_encoder_output(encoder_output): + if self.encoder_runtime_mapping.is_last_pp_rank(): + for pp_rank in self.encoder_runtime_mapping.pp_group: + if pp_rank != self.encoder_runtime_mapping.rank: + self.nccl_comm.send(encoder_output, pp_rank) + return encoder_output + else: + self.nccl_comm.recv(encoder_output, + self.encoder_runtime_mapping.pp_group[-1]) + return encoder_output + + if self.encoder_runtime_mapping.has_pp(): + # use hidden_states output buffer to receive output as the shapes are same + encoder_output_buf = outputs[ + 'encoder_output'] if self.encoder_runtime_mapping.is_last_pp_rank( + ) else outputs['hidden_states_output'] + encoder_output = pp_communicate_encoder_output(encoder_output_buf) + else: + encoder_output = outputs['encoder_output'] + # ------------------------------------------- if debug_mode: torch.cuda.synchronize() @@ -190,11 +303,11 @@ def encoder_run(self, print("Debug output for Encoder") print("--------------------------------------") print("Registered output tensors are: ", outputs.keys()) - print_tensor('encoder_output', outputs['encoder_output']) + print_tensor('encoder_output', encoder_output) print("--------------------------------------") # ------------------------------------------- - return outputs + return encoder_output def generate( self, @@ -207,27 +320,31 @@ def generate( bos_token_id=None, debug_mode=False, ): - # encoder run - encoder_input_lengths = torch.tensor( - [len(x) for x in encoder_input_ids], - dtype=torch.int32, - device='cuda') - encoder_outputs = self.encoder_run(encoder_input_ids, - debug_mode=debug_mode) - encoder_output = encoder_outputs['encoder_output'] - torch.cuda.synchronize() - - # decoder_batch_size = decoder_input_ids.shape[0] - decoder_input_lengths = torch.tensor( - [len(x) for x in decoder_input_ids], - dtype=torch.int32, - device='cuda') - decoder_max_input_length = torch.max(decoder_input_lengths).item() + ## ensure all externally provided tensors are on the correct device. + encoder_input_ids = encoder_input_ids.to(self.device) + decoder_input_ids = decoder_input_ids.to(self.device) + + ## encoder run + logger.info(f"Rank {self.runtime_rank} Running encoder engine ...") + encoder_input_ids, encoder_input_lengths, encoder_max_input_length = self.process_input( + encoder_input_ids, self.encoder_model_config.remove_input_padding, + pad_token_id) + encoder_output = self.encoder_run(encoder_input_ids, + encoder_input_lengths, + encoder_max_input_length, + debug_mode=debug_mode) + + ## decoder run + logger.info(f"Rank {self.runtime_rank} Running decoder engine ...") + decoder_input_ids, decoder_input_lengths, decoder_max_input_length = self.process_input( + decoder_input_ids, self.decoder_model_config.remove_input_padding, + pad_token_id) # generation config sampling_config = SamplingConfig(end_id=eos_token_id, pad_id=pad_token_id, - num_beams=num_beams) + num_beams=num_beams, + min_length=1) # decoder autoregressive generation self.decoder_session.setup( @@ -235,9 +352,10 @@ def generate( decoder_max_input_length, max_new_tokens, num_beams, - encoder_max_input_length=encoder_output.shape[1]) - + max_kv_cache_length=None, + encoder_max_input_length=encoder_max_input_length) torch.cuda.synchronize() + output_ids = self.decoder_session.decode( decoder_input_ids, decoder_input_lengths, @@ -255,35 +373,83 @@ def generate( os.environ["TOKENIZERS_PARALLELISM"] = "false" args = parse_arguments() - tensorrt_llm.logger.set_level(args.log_level) - - input_text = "translate English to German: The house is wonderful, radiating timeless charm and offering a warm, inviting interior with beautiful details and a serene backyard." - tokenizer = AutoTokenizer.from_pretrained(args.tokenizer) - input_ids = tokenizer(input_text, return_tensors="pt").input_ids.type( - torch.IntTensor - ).cuda( - ) # by default int64, must cast to int32! otherwise TRT OOTB or lookup plugin will interpret as [a, 0, b, 0, c, 0, ...] + logger.set_level(args.log_level) + + test_remove_padding = True + if not test_remove_padding: + input_text = "translate English to German: The house is wonderful, radiating timeless charm and offering a warm, inviting interior with beautiful details and a serene backyard." + else: + input_text = [ + "translate English to German: The house is wonderful.", + "summarize: I am a high-performance inference optimizer and runtime.", + ] + + tokenizer = AutoTokenizer.from_pretrained(args.model_name) + tokenized_inputs = tokenizer(input_text, return_tensors="pt", padding=True) + max_new_tokens = args.max_new_tokens + input_ids = tokenized_inputs.input_ids.type(torch.IntTensor).to( + 'cuda') # [batch_size, padded_length] + # by default int64, must cast to int32! otherwise C++ kernel will interpret as [a, 0, b, 0, c, 0, ...] + + if tensorrt_llm.mpi_rank() == 0: + print("--------------------------------------") + print( + f"BOS={tokenizer.bos_token_id}, PAD={tokenizer.pad_token_id}, EOS={tokenizer.eos_token_id}" + ) + print("input text: ", input_text) + print("input ids: ", input_ids) + print("input lengths: ", tokenized_inputs.attention_mask.sum(dim=1)) + print("--------------------------------------") - print("--------------------------------------") - print( - f"BOS={tokenizer.bos_token_id}, PAD={tokenizer.pad_token_id}, EOS={tokenizer.eos_token_id}" - ) - print("input ids: ", input_ids) - print("input length: ", input_ids.shape[1]) - print("input tokens: ", input_text) - print("--------------------------------------") + model_config = AutoConfig.from_pretrained(args.model_name) - hf_model = T5ForConditionalGeneration.from_pretrained( - args.tokenizer).cuda().eval() # start_id for decoder (could add more input_ids as forced_decoder_ids) - decoder_input_ids = torch.IntTensor( - [[hf_model.config.decoder_start_token_id]]).cuda() + decoder_input_ids = torch.IntTensor([[model_config.decoder_start_token_id] + ]).to('cuda') + decoder_input_ids = decoder_input_ids.repeat((input_ids.shape[0], 1)) + + # simple comparison with HF on FP32 + if args.compare_hf_fp32: + if tensorrt_llm.mpi_rank() == 0: + if "t5" in args.model_name: + hf_model = T5ForConditionalGeneration.from_pretrained( + args.model_name).to('cuda') + else: + pass + + tik = time.time() + hf_output_ids = hf_model.generate( + input_ids=input_ids, + decoder_input_ids=decoder_input_ids, + max_new_tokens=max_new_tokens, + num_beams=args.num_beams, + bos_token_id=tokenizer.bos_token_id, + pad_token_id=tokenizer.pad_token_id, + eos_token_id=tokenizer.eos_token_id, + use_cache=True) + torch.cuda.synchronize() + tok = time.time() + + output_ids = hf_output_ids.squeeze(dim=1) + hf_output_text = tokenizer.batch_decode(output_ids, + skip_special_tokens=True) + decoder_input_lengths = (decoder_input_ids != + tokenizer.pad_token_id).sum(dim=1) + output_gen_lengths = (output_ids != tokenizer.eos_token_id).sum( + dim=1) - decoder_input_lengths + print("--------------------------------------") + print("HF output_ids: ", output_ids) + print("HF output text: ", hf_output_text) + print("HF output generated lengths: ", output_gen_lengths) + print(f"HF E2E time {(tok-tik)*1000}ms") + print("--------------------------------------") # TRT-LLM runtime tllm_model = TRTLLMEncDecModel.from_engine(args.engine_name, args.engine_dir, debug_mode=args.debug_mode) + tik = time.time() tllm_output_ids = tllm_model.generate( encoder_input_ids=input_ids, decoder_input_ids=decoder_input_ids, @@ -294,10 +460,38 @@ def generate( eos_token_id=tokenizer.eos_token_id, debug_mode=args.debug_mode, ) - - print("--------------------------------------") - print("TRTLLM output_ids: ", tllm_output_ids) - print("TRTLLM output length: ", tllm_output_ids[0].shape[1]) - print("TRTLLM tokens: ", - tokenizer.decode(tllm_output_ids[0][0], skip_special_tokens=True)) - print("--------------------------------------") + tok = time.time() + + inference_dtype = tllm_model.encoder_model_config.dtype + + if tensorrt_llm.mpi_rank() == 0: + output_ids = tllm_output_ids[:, 0, :] + output_text = tokenizer.batch_decode(output_ids, + skip_special_tokens=True) + decoder_input_lengths = (decoder_input_ids != + tokenizer.pad_token_id).sum(dim=1) + output_gen_lengths = (output_ids != tokenizer.eos_token_id).sum( + dim=1) - decoder_input_lengths + print("--------------------------------------") + print("TRT-LLM output_ids: ", output_ids) + print("TRT-LLM output text: ", output_text) + print("TRT-LLM output generated lengths: ", output_gen_lengths) + print(f"TRT-LLM E2E time {(tok-tik)*1000}ms") + print("Precision:", inference_dtype) + print("--------------------------------------") + + # simple accuracy check + if args.compare_hf_fp32: + from difflib import SequenceMatcher + match_rate = SequenceMatcher(None, "\n".join(output_text), + "\n".join(hf_output_text)).ratio() + print(output_text) + print(hf_output_text) + assert match_rate > 0.95, f"Incorrect results! Match rate {match_rate}" + print( + f"TRT-LLM results match HF FP32 results with literal match rate {match_rate}" + ) + if inference_dtype != "float32": + print( + f"Caveat: comparing TRT-LLM {inference_dtype} results with HF float32 results. Close match are not expected." + ) diff --git a/examples/enc_dec/t5/hf_convert.py b/examples/enc_dec/t5/hf_convert.py new file mode 100644 index 000000000000..f3890e36885c --- /dev/null +++ b/examples/enc_dec/t5/hf_convert.py @@ -0,0 +1,222 @@ +import argparse +import configparser +import logging +import multiprocessing +import os +from datetime import datetime +from pathlib import Path + +dir_path = os.path.dirname(os.path.realpath(__file__)) + +import numpy as np +import torch # pytype: disable=import-error +from transformers import T5ForConditionalGeneration + +from tensorrt_llm._utils import str_dtype_to_torch, torch_to_numpy + +LOGGER = logging.getLogger(__name__) + +extra_configs = { + "structure": { + "t5_with_bias": "false", + "use_gated_activation": "false", + "position_embedding_type": "relative", + 'model_type': 't5' + } +} + + +def fuse_qkv(model, factor, saved_dir): + + def get_attn_module(component, block, layer, attn_type): + m = getattr(model, component) + m = m.block[int(block)].layer[int(layer)] + m = getattr(m, attn_type) + return m + + for name, param in model.named_parameters(): + if 'Attention.q' in name: + q = param + component, _, block_idx, _, layer_idx, attn_type, *_ = name.split( + '.') + attn_mdl = get_attn_module(component, block_idx, layer_idx, + attn_type) + shape = q.shape # (d_out, d_in) + qkv = torch.cat([q, attn_mdl.k.weight, attn_mdl.v.weight], + dim=0).reshape([3, shape[0], + shape[1]]) # (3, d_out, d_in) + qkv = torch_to_numpy(qkv) + # embed_dim --> hidden_dim qkv projection weights, [3, hidden_dim, embed_dim] split dim=1 + # ColumnLinear projection weights W=[3*d_out/TP, d_in], split dim=0 or dim=1 in [3, d_out/TP, d_in] + split_dim = 1 + split_vals = np.split(qkv, factor, axis=split_dim) + for j in range(factor): + saved_path = saved_dir / f"{component}.block.{block_idx}.layer.{layer_idx}.{attn_type}.qkv.weight.{j}.bin" + split_vals[j].tofile(saved_path.as_posix()) + + +def split_and_convert_process(key, val, factor, saved_dir): + # The split_factor indicates the number of ranks to implement + # distributed GEMMs. For Tensor Parallelism, each rank/GPU works + # on split_hidden_dim // split_factor channels. + + saved_key = key + LOGGER.debug(f"key: {key}, val.shape: {val.shape}") + + if "shared.weight" in key or "layer_norm.weight" in key: + # embedding table / layernorm weight, no split + saved_path = saved_dir / f"{saved_key}.bin" + val.tofile(saved_path.as_posix()) + + elif "relative_attention_bias" in key: + # relative attention table, transpose [num_buckets, num_heads] -> [num_heads, num_buckets] + # and split on num_heads // split_factor dim + split_dim = 0 + val = np.ascontiguousarray(val.transpose(1, 0)) + split_vals = np.split(val, factor, axis=split_dim) + for j in range(factor): + saved_path = saved_dir / f"{saved_key}.{j:d}.bin" + split_vals[j].tofile(saved_path.as_posix()) + + elif ("SelfAttention.o.weight" in key or "EncDecAttention.o.weight" in key + or "DenseReluDense.wo.weight" in key): + # RowLinear projection weight W=[d_out, d_in/TP], split dim=-1 + split_dim = -1 + split_vals = np.split(val, factor, axis=split_dim) + for j in range(factor): + saved_path = saved_dir / f"{saved_key}.{j:d}.bin" + split_vals[j].tofile(saved_path.as_posix()) + + elif ("lm_head.weight" in key or "DenseReluDense.wi.weight" in key + or "DenseReluDense.wi_0.weight" in key + or "DenseReluDense.wi_1.weight" in key): + # ColumnLinear projection weights W=[d_out/TP, d_in], split dim=0 + split_dim = 0 + if "DenseReluDense.wi_0.weight" in key: + saved_key = key.replace("wi_0", "wi") + elif "DenseReluDense.wi_1.weight" in key: + saved_key = key.replace("wi_1", "wi2") + split_vals = np.split(val, factor, axis=split_dim) + for j in range(factor): + saved_path = saved_dir / f"{saved_key}.{j:d}.bin" + split_vals[j].tofile(saved_path.as_posix()) + + elif (("encoder" in key and + ("SelfAttention.q.weight" in key or "SelfAttention.k.weight" in key + or "SelfAttention.v.weight" in key)) or + ("decoder" in key and + ("SelfAttention.q.weight" in key or "SelfAttention.k.weight" in key + or "SelfAttention.v.weight" in key or "EncDecAttention.q.weight" + in key or "EncDecAttention.k.weight" in key + or "EncDecAttention.v.weight" in key))): + # weight needs to be fused, handled by fuse_qkv() + pass + + elif "encoder.embed_tokens.weight" in key or "decoder.embed_tokens.weight" in key: + LOGGER.warning(f"Not save {key}, using shared.weight directly.") + + else: + LOGGER.warning(f"cannot find key '{key}' with shape {val.shape}") + + +def convert_checkpoint(args): + saved_dir = Path(args.output_dir) / f"tp{args.inference_tensor_para_size}" + saved_dir.mkdir(parents=True, exist_ok=True) + + t5_model = T5ForConditionalGeneration.from_pretrained(args.input_dir) + t5_model = t5_model.to(str_dtype_to_torch(args.weight_data_type)) + + config = configparser.ConfigParser() + extra_configs["structure"]["use_gated_activation"] = str( + t5_model.encoder.config.is_gated_act) + + config["encoder"] = {} + for key, val in t5_model.encoder.config.to_dict().items(): + config["encoder"][key] = f"{val}" + config["encoder"]["weight_data_type"] = args.weight_data_type + + # manually set q_scaling to offset attention scaling's effect. + # TODO: modify kernels to control whether to disable attention scaling + def get_offset_q_scaling(config) -> str: + d_model = config.d_model + num_heads = config.num_heads + head_size = d_model / num_heads + scaling = 1 / head_size**.5 + return str(scaling) + + config["encoder"]["q_scaling"] = get_offset_q_scaling( + t5_model.encoder.config) + + config["decoder"] = {} + for key, val in t5_model.decoder.config.to_dict().items(): + config["decoder"][key] = f"{val}" + config["decoder"]["weight_data_type"] = args.weight_data_type + + config["decoder"]["q_scaling"] = get_offset_q_scaling( + t5_model.decoder.config) + + for key, val in extra_configs.items(): + config[key] = {} + for val_key, val_val in val.items(): + config[key][val_key] = val_val + with open((saved_dir / f"config.ini").as_posix(), 'w') as configfile: + config.write(configfile) + + i_gpu_num = args.inference_tensor_para_size + + pool = multiprocessing.Pool(args.processes) + pool.starmap_async(split_and_convert_process, + [(name, torch_to_numpy(param), i_gpu_num, saved_dir) + for name, param in t5_model.state_dict().items()]) + + pool.close() + pool.join() + + fuse_qkv(t5_model, i_gpu_num, saved_dir) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + formatter_class=argparse.RawTextHelpFormatter) + parser.add_argument("--input_dir", + "-i", + type=str, + help="Path to the framework checkpoint file", + required=True) + parser.add_argument("--output_dir", + "-o", + type=str, + help="Path to the converted TRT-LLM model weight file", + required=True) + parser.add_argument("--inference_tensor_para_size", + "-i_g", + type=int, + help="How many gpus for inference", + required=True) + parser.add_argument( + "--processes", + "-p", + type=int, + help="How many processes to spawn for conversion (default: 4)", + default=4) + parser.add_argument("--weight_data_type", + type=str, + default="float32", + choices=["float32", "float16", "bfloat16"]) + parser.add_argument("--verbose", + action="store_true", + help="Provide verbose messages") + args = parser.parse_args() + log_format = "%(asctime)s %(name)s [%(levelname)s] %(message)s" + logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, + format=log_format) + LOGGER.info("\n=============== Argument ===============") + for key in vars(args): + LOGGER.info(f"{key}: {vars(args)[key]}") + LOGGER.info("========================================") + + start_time = datetime.now() + convert_checkpoint(args) + stop_time = datetime.now() + run_time = (stop_time - start_time) + LOGGER.info("Spend {} (h:m:s) to convert the model".format(run_time)) diff --git a/examples/enc_dec/weight.py b/examples/enc_dec/t5/weight.py similarity index 55% rename from examples/enc_dec/weight.py rename to examples/enc_dec/t5/weight.py index b775ad478bc0..0eaae0c57503 100644 --- a/examples/enc_dec/weight.py +++ b/examples/enc_dec/t5/weight.py @@ -1,27 +1,37 @@ -import configparser +import time +from os import path +from pathlib import Path +from typing import Optional, Union import numpy as np import torch -from tensorrt_llm._utils import str_dtype_to_torch, torch_to_numpy -from tensorrt_llm.functional import LayerNormPositionType, LayerNormType +from tensorrt_llm import logger +from tensorrt_llm._utils import (numpy_to_dtype, str_dtype_to_np, + str_dtype_to_torch, torch_to_numpy) +from tensorrt_llm.functional import (LayerNormPositionType, LayerNormType, + MLPType) +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models import ( # TODO: probably need to change model name to distinguish from other models + DecoderModel, EncoderModel) layernorm_type_map = {i.name: i.value for i in LayerNormType} layernorm_position_map = {i.name: i.value for i in LayerNormPositionType} +mlp_type_map = {i.name: i.value for i in MLPType} -def parse_config(ini_file, component, args): - config = configparser.ConfigParser() - config.read(ini_file) +def parse_t5_config(config, component, args): if component == 'encoder': - args.n_layer = config.getint(component, 'n_layer') - args.n_head = config.getint(component, 'n_head') - args.hidden_size = config.getint(component, 'hidden_size') - args.ffn_hidden_size = config.getint(component, 'ffn_hidden_size') + args.n_layer = config.getint(component, 'num_layers') + args.n_head = config.getint(component, 'num_heads') + args.head_size = config.getint(component, 'd_kv') + args.hidden_size = config.getint(component, 'd_model') + args.ffn_hidden_size = config.getint(component, 'd_ff') args.vocab_size = config.getint(component, 'vocab_size') args.n_positions = config.getint(component, 'n_positions') args.has_position_embedding = config.getboolean( - component, 'has_position_embedding', fallback=False) + component, 'has_position_embedding', + fallback=False) # TODO: hardcoded here args.has_token_type_embedding = config.getboolean( component, 'has_token_type_embedding', fallback=False) args.has_embedding_layernorm = config.getboolean( @@ -31,35 +41,43 @@ def parse_config(ini_file, component, args): fallback=False) args.q_scaling = config.getfloat(component, 'q_scaling', fallback=1.0) args.has_attention_qkvo_bias = config.getboolean( - component, 'has_attention_qkvo_bias', fallback=False) + component, 'has_attention_qkvo_bias', + fallback=False) # TODO: hardcoded here args.has_mlp_bias = config.getboolean(component, 'has_mlp_bias', fallback=False) args.has_model_final_layernorm = config.getboolean( - component, 'has_model_final_layernorm', fallback=False) - args.layernorm_eps = config.getfloat(component, - 'layernorm_eps', - fallback=1e-5) + component, 'has_model_final_layernorm', fallback=True) + args.layernorm_eps = config.getfloat(component, 'layer_norm_epsilon') args.layernorm_position = layernorm_position_map[config.get( - component, 'layernorm_position')] + component, 'layernorm_position', + fallback='pre_layernorm')] # TODO: hardcoded here args.layernorm_type = layernorm_type_map[config.get( - component, 'layernorm_type')] - args.hidden_act = config.get(component, 'hidden_act') + component, 'layernorm_type', + fallback='RmsNorm')] # TODO: hardcoded here + args.hidden_act = config.get(component, 'dense_act_fn') + args.gated_act = config.getboolean(component, 'is_gated_act') + args.mlp_type = mlp_type_map['GatedMLP' if args.gated_act else 'MLP'] args.relative_attention = config.getboolean(component, 'relative_attention', - fallback=False) - args.num_buckets = config.getint(component, 'num_buckets') - args.max_distance = config.getint(component, 'max_distance') + fallback=True) + args.num_buckets = config.getint(component, + 'relative_attention_num_buckets') + args.max_distance = config.getint(component, + 'relative_attention_max_distance') + args.ckpt_weight_dtype = config.get(component, 'weight_data_type') elif component == 'decoder': - args.n_layer = config.getint(component, 'n_layer') - args.n_head = config.getint(component, 'n_head') - args.hidden_size = config.getint(component, 'hidden_size') - args.ffn_hidden_size = config.getint(component, 'ffn_hidden_size') + args.n_layer = config.getint(component, 'num_decoder_layers') + args.n_head = config.getint(component, 'num_heads') + args.head_size = config.getint(component, 'd_kv') + args.hidden_size = config.getint(component, 'd_model') + args.ffn_hidden_size = config.getint(component, 'd_ff') args.vocab_size = config.getint(component, 'vocab_size') args.n_positions = config.getint(component, 'n_positions') args.has_position_embedding = config.getboolean( - component, 'has_position_embedding', fallback=False) + component, 'has_position_embedding', + fallback=False) # TODO: hardcoded here args.has_token_type_embedding = config.getboolean( component, 'has_token_type_embedding', fallback=False) args.has_embedding_layernorm = config.getboolean( @@ -74,28 +92,35 @@ def parse_config(ini_file, component, args): 'has_mlp_bias', fallback=False) args.has_model_final_layernorm = config.getboolean( - component, 'has_model_final_layernorm', fallback=False) - args.layernorm_eps = config.getfloat(component, - 'layernorm_eps', - fallback=1e-5) + component, 'has_model_final_layernorm', fallback=True) + args.layernorm_eps = config.getfloat(component, 'layer_norm_epsilon') args.layernorm_position = layernorm_position_map[config.get( - component, 'layernorm_position')] - args.layernorm_type = layernorm_type_map[config.get( - component, 'layernorm_type')] - args.hidden_act = config.get(component, 'hidden_act') - args.has_lm_head_bias = config.getboolean(component, - 'has_lm_head_bias', - fallback=False) + component, 'layernorm_position', + fallback='pre_layernorm')] # TODO: hardcoded here + args.layernorm_type = layernorm_type_map[config.get(component, + 'layernorm_type', + fallback='RmsNorm')] + args.hidden_act = config.get(component, 'dense_act_fn') + args.gated_act = config.getboolean(component, 'is_gated_act') + args.mlp_type = mlp_type_map['GatedMLP' if args.gated_act else 'MLP'] + args.has_lm_head_bias = config.getboolean( + component, # TODO: T5 with bias + 'has_lm_head_bias', + fallback=False) args.relative_attention = config.getboolean(component, 'relative_attention', - fallback=False) - args.num_buckets = config.getint(component, 'num_buckets') - args.max_distance = config.getint(component, 'max_distance') + fallback=True) + args.num_buckets = config.getint(component, + 'relative_attention_num_buckets') + args.max_distance = config.getint(component, + 'relative_attention_max_distance') args.logits_dtype = config.get(component, 'logits_dtype', fallback='float32') - args.encoder_hidden_size = config.getint('encoder', 'hidden_size') - args.encoder_num_heads = config.getint('encoder', 'n_head') + args.encoder_hidden_size = config.getint('encoder', 'd_model') + args.encoder_num_heads = config.getint('encoder', 'num_heads') + args.encoder_head_size = config.getint('encoder', 'd_kv') + args.ckpt_weight_dtype = config.get(component, 'weight_data_type') else: assert False, 'Unsupported component!' @@ -108,13 +133,10 @@ def fuse_qkv(q, k, v): return qkv_weight -def load_t5_from_pytorch(tllm_model, - pytorch_ckpt_path, - component, - dtype="float32"): +def load_from_hf_t5(tllm_model, pytorch_ckpt_path, component, dtype="float32"): torch_dtype = str_dtype_to_torch(dtype) - pytorch_ckpt = torch.load(pytorch_ckpt_path + '/t5_small.ckpt') + pytorch_ckpt = torch.load(path.join(pytorch_ckpt_path, 'pytorch_model.bin')) pytorch_model = { key: torch_to_numpy(value.to(torch_dtype)) for key, value in pytorch_ckpt.items() @@ -294,3 +316,128 @@ def load_t5_from_pytorch(tllm_model, 'decoder.final_layer_norm.bias'] tllm_model.lm_head.weight.value = pytorch_model['lm_head.weight'] + + +# TODO: only support t5, biases are not loaded +def load_from_binary_t5(tllm_model: Union[EncoderModel, DecoderModel], + dir_path, + args, + mapping=Mapping(), + dtype='float32', + use_parallel_embedding=False, + sharding_dim=0, + share_embedding_table=False, + scaling_factors=None): + logger.info('Loading weights from binary...') + tik = time.time() + + ckpt_np_dtype = str_dtype_to_np(args.ckpt_weight_dtype) + + def fromfile(name, split=True, shape=None) -> Optional[np.ndarray]: + p = path.join( + dir_path, + f'{name}.{str(mapping.tp_rank)}.bin' if split else f'{name}.bin') + if Path(p).exists(): + # load from original dtype and cast to inference dtype + t = np.fromfile(p, dtype=ckpt_np_dtype) + t = numpy_to_dtype(t, dtype) + if shape is not None: + t = t.reshape(shape) + t = np.ascontiguousarray(t) + return t + return None + + component = 'encoder' if isinstance(tllm_model, EncoderModel) else 'decoder' + + if mapping.is_first_pp_rank(): + wte = fromfile('shared.weight', + shape=[args.vocab_size, -1], + split=False) + tllm_model.embedding.vocab_embedding.weight.value = wte + + # T5 special: all layers use 1st layer's attn table + relative_attention_table = fromfile( + f'{component}.block.0.layer.0.SelfAttention.relative_attention_bias.weight', + shape=[args.n_head // mapping.tp_size, args.num_buckets]) + + # TP is by loading different split weights. PP is by loading different layer weights + # TODO: fix llama's wrong def of .num_layers field in PP. enc_dec is the correct way + layers_range = list( + range(mapping.pp_rank * tllm_model.num_layers, + (mapping.pp_rank + 1) * tllm_model.num_layers, 1)) + + for layer_idx in layers_range: + pp_offset_layer_idx = layer_idx - mapping.pp_rank * tllm_model.num_layers + layer = getattr(tllm_model, f'{component}_layers')[pp_offset_layer_idx] + layer_prefix = f'{component}.block.{layer_idx}' + + self_attention_layer = getattr( + layer, 'attention' if component == 'encoder' else 'self_attention') + + # attention table for all layers + self_attention_layer.rel_attn_table.value = relative_attention_table + + # self attention + attention_hidden_size = args.n_head * args.head_size # head size * num_heads not necessarily equals hidden_dim, such as Flan-T5 + self_attention_layer.qkv.weight.value = fromfile( + f'{layer_prefix}.layer.0.SelfAttention.qkv.weight', + shape=[ + 3 * attention_hidden_size // mapping.tp_size, args.hidden_size + ]) + self_attention_layer.dense.weight.value = fromfile( + f'{layer_prefix}.layer.0.SelfAttention.o.weight', + shape=[args.hidden_size, attention_hidden_size // mapping.tp_size]) + self_attention_layernorm = getattr( + layer, 'self_attention_layernorm' + if component == 'decoder' else 'attention_layernorm') + self_attention_layernorm.weight.value = fromfile( + f'{layer_prefix}.layer.0.layer_norm.weight', split=False) + + # cross attention + if component == 'decoder': + attention_hidden_size = args.n_head * args.head_size + layer.cross_attention.qkv.weight.value = fromfile( + f'{layer_prefix}.layer.1.EncDecAttention.qkv.weight', + shape=[ + 3 * attention_hidden_size // mapping.tp_size, + args.hidden_size + ]) + layer.cross_attention.dense.weight.value = fromfile( + f'{layer_prefix}.layer.1.EncDecAttention.o.weight', + shape=[ + args.hidden_size, attention_hidden_size // mapping.tp_size + ]) + layer.cross_attention_layernorm.weight.value = fromfile( + f'{layer_prefix}.layer.1.layer_norm.weight', split=False) + + # MLP + hf_component_idx = 1 if component == 'encoder' else 2 + layer.mlp.fc.weight.value = fromfile( + f'{layer_prefix}.layer.{hf_component_idx}.DenseReluDense.wi.weight', + shape=[args.ffn_hidden_size // mapping.tp_size, args.hidden_size]) + if args.gated_act: + layer.mlp.gate.weight.value = fromfile( + f'{layer_prefix}.layer.{hf_component_idx}.DenseReluDense.wi2.weight', + shape=[ + args.ffn_hidden_size // mapping.tp_size, args.hidden_size + ]) + layer.mlp.proj.weight.value = fromfile( + f'{layer_prefix}.layer.{hf_component_idx}.DenseReluDense.wo.weight', + shape=[args.hidden_size, args.ffn_hidden_size // mapping.tp_size]) + layer.mlp_layernorm.weight.value = fromfile( + f'{layer_prefix}.layer.{hf_component_idx}.layer_norm.weight', + split=False) + + if mapping.is_last_pp_rank(): + if tllm_model.has_model_final_layernorm: + tllm_model.final_layernorm.weight.value = fromfile( + f'{component}.final_layer_norm.weight', split=False) + + if component == 'decoder': + tllm_model.lm_head.weight.value = fromfile( + 'lm_head.weight', + shape=[args.vocab_size // mapping.tp_size, args.hidden_size]) + + tok = time.time() + t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) + logger.info(f'Weights loaded. Total time: {t}') diff --git a/examples/falcon/README.md b/examples/falcon/README.md index e57a60b7ffa8..96940c9d5890 100644 --- a/examples/falcon/README.md +++ b/examples/falcon/README.md @@ -4,16 +4,17 @@ This document shows how to build and run a Falcon model in TensorRT-LLM on singl ## Overview -The TensorRT-LLM Falcon implementation can be found in [tensorrt_llm/models/falcon/model.py](../../tensorrt_llm/models/falcon/model.py). The TensorRT-LLM Falcon example code is located in [`examples/falcon`](./). There are three main files in that folder: +The TensorRT-LLM Falcon implementation can be found in [tensorrt_llm/models/falcon/model.py](../../tensorrt_llm/models/falcon/model.py). The TensorRT-LLM Falcon example code is located in [`examples/falcon`](./). There are three main files: * [`build.py`](./build.py) to build the [TensorRT](https://developer.nvidia.com/tensorrt) engine(s) needed to run the Falcon model, * [`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/cnn_dailymail) dataset using the model. + * and a shared [`../summarize.py`](../summarize.py)to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/cnn_dailymail) dataset using the model. ## Support Matrix * FP16 * BF16 * FP8 + * Groupwise quantization (AWQ) * STRONGLY TYPED * FP8 KV CACHE * Tensor Parallel @@ -155,6 +156,46 @@ python build.py --model_dir falcon/180b \ --parallel_build ``` +#### Groupwise quantization (AWQ) +One can enable AWQ INT4 weight only quantization with these options when building engine with `build.py`: + +- `--use_weight_only` enables weight only GEMMs in the network. +- `--per_group` enable groupwise weight only quantization, for Falcon example, we support AWQ with the group size default as 128. +- `--weight_only_precision` should specify the weight only quantization format. Supported formats are `int4_awq` or `int4_gptq`. +- `--group_size` passes the group size for AWQ with default as 128. +- `--quant_ckpt_path` passes the quantized checkpoint to build the engine. + +AWQ example below involves 2 steps: +1. Weight quantization: + + NVIDIA AMMO toolkit is used for AWQ weight quantization. Please see [examples/quantization/README.md](/examples/quantization/README.md#preparation) for AMMO installation instructions. + + ```bash + # Quantize HF Falcon 180B checkpoint into INT4 AWQ format + python quantize.py --model_dir falcon/180B/ \ + --dtype float16 \ + --qformat int4_awq \ + --export_path ./quantized_int4_awq \ + --calib_size 32 + ``` + The quantized model checkpoint is saved to path `./quantized_int4_awq/falcon_tp1_rank0.npz` for future TRT-LLM engine build. + +2. Build TRT-LLM engine: + + ```bash + python build.py --model_dir falcon/180B/ \ + --quant_ckpt_path ./quantized_int4_awq/falcon_tp1_rank0.npz \ + --dtype float16 \ + --remove_input_padding \ + --use_gpt_attention_plugin float16 \ + --enable_context_fmha \ + --use_gemm_plugin float16 \ + --use_weight_only \ + --weight_only_precision int4_awq \ + --per_group \ + --output_dir ./tmp/falcon/180B/trt_engines/int4_AWQ/1-gpu/ + ``` + ### 4. Run ```bash @@ -162,26 +203,26 @@ pip install -r requirements.txt ``` ```bash -python summarize.py --test_trt_llm \ - --hf_model_location falcon/rw-1b \ - --data_type float16 \ - --engine_dir falcon/rw-1b/trt_engines/fp16/1-gpu/ +python ../summarize.py --test_trt_llm \ + --hf_model_dir falcon/rw-1b \ + --data_type float16 \ + --engine_dir falcon/rw-1b/trt_engines/fp16/1-gpu/ -python summarize.py --test_trt_llm \ - --hf_model_location falcon/7b-instruct \ - --data_type bfloat16 \ - --engine_dir falcon/7b-instruct/trt_engines/bf16/1-gpu +python ../summarize.py --test_trt_llm \ + --hf_model_dir falcon/7b-instruct \ + --data_type bfloat16 \ + --engine_dir falcon/7b-instruct/trt_engines/bf16/1-gpu mpirun -n 2 --allow-run-as-root --oversubscribe \ - python summarize.py --test_trt_llm \ - --hf_model_location falcon/40b-instruct \ - --data_type bfloat16 \ - --engine_dir falcon/40b-instruct/trt_engines/bf16/2-gpu + python ../summarize.py --test_trt_llm \ + --hf_model_dir falcon/40b-instruct \ + --data_type bfloat16 \ + --engine_dir falcon/40b-instruct/trt_engines/bf16/2-gpu mpirun -n 8 --allow-run-as-root --oversubscribe \ - python summarize.py --test_trt_llm \ - --hf_model_location falcon/180b \ - --data_type bfloat16 \ - --engine_dir falcon/180b/trt_engines/bf16/8-gpu + python ../summarize.py --test_trt_llm \ + --hf_model_dir falcon/180b \ + --data_type bfloat16 \ + --engine_dir falcon/180b/trt_engines/bf16/8-gpu ``` ## Troubleshooting diff --git a/examples/falcon/build.py b/examples/falcon/build.py index 498239eb68fb..e8984559b293 100644 --- a/examples/falcon/build.py +++ b/examples/falcon/build.py @@ -24,6 +24,8 @@ import torch.multiprocessing as mp from onnx import TensorProto, helper from transformers import AutoModelForCausalLM, FalconConfig +from weight import (get_scaling_factors, load_from_awq_falcon, + load_from_hf_checkpoint, load_from_hf_falcon) import tensorrt_llm from tensorrt_llm._utils import str_dtype_to_trt @@ -36,10 +38,6 @@ from tensorrt_llm.profiler import check_gpt_mem_usage from tensorrt_llm.quantization import QuantMode -from weight import get_scaling_factors # isort:skip -from weight import load_from_hf_falcon # isort:skip -from weight import load_from_hf_checkpoint # isort:skip - MODEL_NAME = 'falcon' @@ -154,6 +152,7 @@ def parse_arguments(): parser.add_argument('--tp_size', type=int, default=1) parser.add_argument('--pp_size', type=int, default=1) parser.add_argument('--model_dir', type=str, default=None) + parser.add_argument('--quant_ckpt_path', type=str, default=None) parser.add_argument('--dtype', type=str, default='float16', @@ -297,7 +296,33 @@ def parse_arguments(): action='store_true', help= 'Activates latency-optimized algorithm for all-reduce instead of NCCL.') - + parser.add_argument( + '--per_group', + default=False, + action="store_true", + help= + 'By default, we use a single static scaling factor to scale weights in the int4 range. ' + 'per_group chooses at run time, and for each group, a custom scaling factor. ' + 'The flag is built for GPTQ/AWQ quantization.') + parser.add_argument('--group_size', + type=int, + default=128, + help='Group size used in GPTQ/AWQ quantization.') + parser.add_argument( + '--use_weight_only', + default=False, + action="store_true", + help='Quantize weights for the various GEMMs to INT4/INT8.' + 'See --weight_only_precision to set the precision') + parser.add_argument( + '--weight_only_precision', + type=str, + default='int4_awq', + choices=['int4_awq'], + help= + 'Define the precision for the weights when using weight-only quantization.' + 'You must also use --use_weight_only for that argument to have an impact.' + ) args = parser.parse_args() logger.set_level(args.log_level) @@ -326,6 +351,15 @@ def parse_arguments(): assert args.enable_context_fmha args.quant_mode = QuantMode(0) + if args.use_weight_only: + assert args.enable_fp8, "FP8 and Weight-only cannot be activated simultaneously!" + if args.weight_only_precision == 'int4_awq': + args.quant_mode = QuantMode.from_description( + quantize_weights=True, + quantize_activations=False, + per_token=False, + per_channel=False, + per_group=args.per_group) if args.fp8_kv_cache: args.quant_mode = args.quant_mode.set_fp8_kv_cache() if args.enable_fp8: @@ -415,16 +449,30 @@ def build_rank_engine(builder: Builder, parallel_attention=args.parallel_attention, new_decoder_architecture=args.new_decoder_architecture) - if args.enable_fp8 or args.fp8_kv_cache: + quantize_kwargs = {} + if args.use_weight_only and args.weight_only_precision == 'int4_awq': + quantize_kwargs = { + "group_size": args.group_size, + "zero": False, + "pre_quant_scale": True, + "exclude_modules": [], + } + elif args.enable_fp8 or args.fp8_kv_cache: logger.info(f'Loading scaling factors from ' f'{args.quantized_fp8_model_path}') quant_scales = get_scaling_factors(args.quantized_fp8_model_path, num_layers=args.n_layer, quant_mode=args.quant_mode) - tensorrt_llm_falcon = quantize_model(tensorrt_llm_falcon, - quant_mode=args.quant_mode, - quant_scales=quant_scales) - if args.model_dir is not None: + quantize_kwargs = {"quant_scales": quant_scales} + tensorrt_llm_falcon = quantize_model(tensorrt_llm_falcon, args.quant_mode, + **quantize_kwargs) + + if args.per_group: + load_from_awq_falcon(tensorrt_llm_falcon=tensorrt_llm_falcon, + quant_ckpt_path=args.quant_ckpt_path, + mapping=mapping, + dtype=args.dtype) + elif args.model_dir is not None: logger.info(f'Loading HF Falcon ... from {args.model_dir}') tik = time.time() if not args.load_by_shard: @@ -470,6 +518,10 @@ def build_rank_engine(builder: Builder, if args.multi_block_mode: network.plugin_config.enable_mmha_multi_block_mode() + if args.per_group: + network.plugin_config.set_weight_only_groupwise_quant_matmul_plugin( + dtype=args.dtype) + if args.world_size > 1: network.plugin_config.set_nccl_plugin(args.dtype, args.use_custom_all_reduce) @@ -527,6 +579,9 @@ def build(rank, args): # skip other ranks if parallel_build is enabled if args.parallel_build and cur_rank != rank: continue + # NOTE: when only int8 kv cache is used together with paged kv cache no int8 tensors are exposed to TRT + int8_trt_flag = args.quant_mode.has_act_or_weight_quant() or ( + not args.paged_kv_cache and args.quant_mode.has_int8_kv_cache()) builder_config = builder.create_builder_config( name=MODEL_NAME, precision=args.dtype, @@ -543,11 +598,13 @@ def build(rank, args): new_decoder_architecture=args.new_decoder_architecture, max_position_embeddings=args.n_positions, max_batch_size=args.max_batch_size, + max_beam_width=args.max_beam_width, max_input_len=args.max_input_len, max_output_len=args.max_output_len, max_num_tokens=args.max_num_tokens, quant_mode=args.quant_mode, strongly_typed=args.strongly_typed, + int8=int8_trt_flag, opt_level=args.builder_opt) engine_name = get_engine_name(MODEL_NAME, args.dtype, args.tp_size, args.pp_size, cur_rank) diff --git a/examples/falcon/quantize.py b/examples/falcon/quantize.py index 443c4061f5a3..fd0b86e7b270 100644 --- a/examples/falcon/quantize.py +++ b/examples/falcon/quantize.py @@ -100,7 +100,7 @@ def get_args(): parser.add_argument("--dtype", help="Model data type.", default="float16") parser.add_argument("--qformat", type=str, - choices=['fp8'], + choices=['fp8', 'int4_awq'], default='fp8', help='Quantization format.') parser.add_argument("--calib_size", diff --git a/examples/falcon/requirements.txt b/examples/falcon/requirements.txt index edaad3be36cc..7888a6752130 100644 --- a/examples/falcon/requirements.txt +++ b/examples/falcon/requirements.txt @@ -1,6 +1,6 @@ transformers>=4.31.0 datasets~=2.14.5 +evaluate~=0.4.1 rouge_score~=0.1.2 sentencepiece~=0.1.99 -typing-extensions==4.5.0 tqdm diff --git a/examples/falcon/summarize.py b/examples/falcon/summarize.py deleted file mode 100644 index dbd3cf70df28..000000000000 --- a/examples/falcon/summarize.py +++ /dev/null @@ -1,446 +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 copy -import json -import os -from pathlib import Path - -import numpy as np -import torch -from datasets import load_dataset, load_metric -from transformers import AutoModelForCausalLM, AutoTokenizer - -import tensorrt_llm -import tensorrt_llm.profiler as profiler -from tensorrt_llm.logger import logger -from tensorrt_llm.quantization import QuantMode - -from build import get_engine_name # isort:skip - - -def TRTFalcon(args, config): - builder_config = config['builder_config'] - plugin_config = config['plugin_config'] - - dtype = builder_config['precision'] - tp_size = builder_config['tensor_parallel'] - pp_size = builder_config['pipeline_parallel'] - world_size = tp_size * pp_size - assert world_size == tensorrt_llm.mpi_world_size(), \ - f'Engine world size ({world_size}) != Runtime world size '\ - f'({tensorrt_llm.mpi_world_size()})' - - num_heads = builder_config['num_heads'] // tp_size - hidden_size = builder_config['hidden_size'] // tp_size - vocab_size = builder_config['vocab_size'] - num_layers = builder_config['num_layers'] - num_kv_heads = builder_config.get('num_kv_heads', num_heads) - num_kv_heads = (num_kv_heads + tp_size - 1) // tp_size - quant_mode = QuantMode(builder_config['quant_mode']) - - use_gpt_attention_plugin = bool(plugin_config['gpt_attention_plugin']) - paged_kv_cache = plugin_config['paged_kv_cache'] - tokens_per_block = plugin_config['tokens_per_block'] - remove_input_padding = plugin_config['remove_input_padding'] - use_custom_all_reduce = plugin_config.get('use_custom_all_reduce', False) - - model_config = tensorrt_llm.runtime.ModelConfig( - vocab_size=vocab_size, - num_layers=num_layers, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - hidden_size=hidden_size, - gpt_attention_plugin=use_gpt_attention_plugin, - paged_kv_cache=paged_kv_cache, - tokens_per_block=tokens_per_block, - remove_input_padding=remove_input_padding, - quant_mode=quant_mode, - dtype=dtype, - use_custom_all_reduce=use_custom_all_reduce) - - runtime_rank = tensorrt_llm.mpi_rank() - runtime_mapping = tensorrt_llm.Mapping(world_size, - runtime_rank, - tp_size=tp_size, - pp_size=pp_size) - torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node) - - engine_name = get_engine_name('falcon', dtype, tp_size, pp_size, - runtime_rank) - serialize_path = os.path.join(args.engine_dir, engine_name) - - profiler.start('load tensorrt_llm engine') - with open(serialize_path, 'rb') as f: - engine_buffer = f.read() - decoder = tensorrt_llm.runtime.GenerationSession(model_config, - engine_buffer, - runtime_mapping, - debug_mode=args.debug) - profiler.stop('load tensorrt_llm engine') - loading_time = profiler.elapsed_time_in_sec("load tensorrt_llm engine") - logger.info(f'Load engine takes: {loading_time} sec') - return decoder - - -def main(args): - 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 - hf_model_location = args.hf_model_location - profiler.start('load tokenizer') - tokenizer = AutoTokenizer.from_pretrained(hf_model_location, - padding_side='left') - profiler.stop('load tokenizer') - logger.info( - f'Load tokenizer takes: {profiler.elapsed_time_in_sec("load tokenizer")} sec' - ) - if tokenizer.pad_token_id is None: - tokenizer.pad_token_id = tokenizer.eos_token_id - - dataset_cnn = load_dataset("ccdv/cnn_dailymail", - '3.0.0', - cache_dir=args.dataset_path) - - max_batch_size = args.batch_size - - # runtime parameters - top_k = args.top_k - output_len = args.output_len - test_token_num = 923 - temperature = 1 - repetition_penalty = 1 - num_beams = args.num_beams - - pad_id = tokenizer.pad_token_id - end_id = tokenizer.eos_token_id - - if test_trt_llm: - config_path = os.path.join(args.engine_dir, 'config.json') - with open(config_path, 'r') as f: - config = json.load(f) - tensorrt_llm_falcon = TRTFalcon(args, config) - - if test_hf: - profiler.start('load HF model') - torch_dtype = tensorrt_llm._utils.str_dtype_to_torch(args.data_type) - model = AutoModelForCausalLM.from_pretrained( - hf_model_location, - trust_remote_code=True, - torch_dtype=torch_dtype, - device_map='auto' if args.hf_device_map_auto else None) - if not args.hf_device_map_auto: - model.cuda() - profiler.stop('load HF model') - hf_loading_time = profiler.elapsed_time_in_sec('load HF model') - logger.info(f'Load HF model takes: {hf_loading_time} sec') - - 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.hf_model_location}\n') - if test_hf: - with (output_dir / 'hf.out').open('w') as f: - f.write(f'Model path: {args.hf_model_location}\n') - - def summarize_tensorrt_llm(datapoint): - batch_size = len(datapoint['article']) - - line = copy.copy(datapoint['article']) - line_encoded = [] - input_lengths = [] - for i in range(batch_size): - line[i] = line[i] + ' TL;DR: ' - - line[i] = line[i].strip() - line[i] = line[i].replace(" n't", "n't") - - input_id = tokenizer.encode(line[i], - return_tensors='pt').type(torch.int32) - input_id = input_id[:, -test_token_num:] - - line_encoded.append(input_id) - input_lengths.append(input_id.shape[-1]) - - # do padding, should move outside the profiling to prevent the overhead - max_length = max(input_lengths) - if tensorrt_llm_falcon.remove_input_padding: - line_encoded = [ - torch.tensor(t, dtype=torch.int32).cuda() for t in line_encoded - ] - else: - for i in range(batch_size): - pad_size = max_length - input_lengths[i] - - pad = torch.ones([1, pad_size]).type(torch.int32) * pad_id - line_encoded[i] = torch.cat([line_encoded[i], pad], axis=-1) - - line_encoded = torch.cat(line_encoded, axis=0).cuda() - input_lengths = torch.tensor(input_lengths, - dtype=torch.int32).cuda() - - sampling_config = tensorrt_llm.runtime.SamplingConfig( - end_id=end_id, - pad_id=pad_id, - top_k=top_k, - num_beams=num_beams, - temperature=temperature, - repetition_penalty=repetition_penalty) - - with torch.no_grad(): - tensorrt_llm_falcon.setup(batch_size, - max_context_length=max_length, - max_new_tokens=output_len, - beam_width=num_beams, - max_kv_cache_length=args.max_kv_cache_len) - - if tensorrt_llm_falcon.remove_input_padding: - output_ids = tensorrt_llm_falcon.decode_batch( - line_encoded, sampling_config) - else: - output_ids = tensorrt_llm_falcon.decode( - line_encoded, - input_lengths, - sampling_config, - ) - torch.cuda.synchronize() - - output_lines_list, tokens_list = [], [] - # output_ids = [batch_size, num_beams, output_len] - if tensorrt_llm_falcon.mapping.is_first_pp_rank(): - tokens_list = output_ids[:, :, max_length:].tolist() - output_lines_list = [ - tokenizer.batch_decode(output_ids[:, i, max_length:], - skip_special_tokens=True) - for i in range(num_beams) - ] - return output_lines_list, tokens_list - - def summarize_hf(datapoint): - batch_size = len(datapoint['article']) - if batch_size > 1: - logger.warning( - f"HF does not support batch_size > 1 to verify correctness " - f"due to padding. Current batch size is {batch_size}") - - line = copy.copy(datapoint['article']) - for i in range(batch_size): - line[i] = line[i] + ' TL;DR: ' - - line[i] = line[i].strip() - line[i] = line[i].replace(" n't", "n't") - - line_encoded = tokenizer(line, return_tensors='pt', - padding=True)["input_ids"].long() - - line_encoded = line_encoded[:, -test_token_num:] - line_encoded = line_encoded.cuda() - - with torch.no_grad(): - output = model.generate(line_encoded, - max_length=len(line_encoded[0]) + - output_len, - top_k=top_k, - temperature=temperature, - eos_token_id=tokenizer.eos_token_id, - pad_token_id=tokenizer.pad_token_id, - num_beams=num_beams, - num_return_sequences=num_beams, - early_stopping=True) - - tokens_list = output[:, len(line_encoded[0]):].tolist() - output = output.reshape([batch_size, num_beams, -1]) - output_lines_list = [ - tokenizer.batch_decode(output[:, i, len(line_encoded[0]):], - skip_special_tokens=True) - for i in range(num_beams) - ] - - return output_lines_list, tokens_list - - if test_trt_llm: - datapoint = dataset_cnn['test'][0:1] - summary, _ = summarize_tensorrt_llm(datapoint) - if runtime_rank == 0: - logger.info( - "---------------------------------------------------------") - logger.info("TensorRT-LLM Generated : ") - logger.info(f" Article : {datapoint['article']}") - logger.info(f"\n Highlights : {datapoint['highlights']}") - logger.info(f"\n Summary : {summary}") - logger.info( - "---------------------------------------------------------") - - if test_hf: - datapoint = dataset_cnn['test'][0:1] - summary, _ = summarize_hf(datapoint) - logger.info("---------------------------------------------------------") - logger.info("HF Generated : ") - logger.info(f" Article : {datapoint['article']}") - logger.info(f"\n Highlights : {datapoint['highlights']}") - logger.info(f"\n Summary : {summary}") - logger.info("---------------------------------------------------------") - - if args.max_ite == 0: - return - - metric_tensorrt_llm = [load_metric("rouge") for _ in range(num_beams)] - metric_hf = [load_metric("rouge") for _ in range(num_beams)] - for i in range(num_beams): - metric_tensorrt_llm[i].seed = 0 - metric_hf[i].seed = 0 - - ite_count = 0 - data_point_idx = 0 - while (data_point_idx < len(dataset_cnn['test'])) 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_cnn['test'][data_point_idx:(data_point_idx + - max_batch_size)] - - if test_trt_llm: - profiler.start('tensorrt_llm') - summary_tensorrt_llm, _ = summarize_tensorrt_llm(datapoint) - profiler.stop('tensorrt_llm') - - if test_hf: - profiler.start('hf') - summary_hf, _ = summarize_hf(datapoint) - profiler.stop('hf') - - if runtime_rank == 0: - if test_trt_llm: - for beam_idx in range(num_beams): - for i in range(len(summary_tensorrt_llm[beam_idx])): - metric_tensorrt_llm[beam_idx].add_batch( - predictions=[summary_tensorrt_llm[beam_idx][i]], - references=[datapoint['highlights'][i]]) - if output_dir is not None: - # yapf: disable - for i in range(len(summary_tensorrt_llm[0])): - for beam_idx in range(num_beams): - with (output_dir / 'trtllm.out').open('a') as f: - f.write(f'[{data_point_idx + i}] [Beam {beam_idx}] {summary_tensorrt_llm[beam_idx][i]}\n') - # yapf: enable - if test_hf: - for beam_idx in range(num_beams): - for i in range(len(summary_hf[beam_idx])): - metric_hf[beam_idx].add_batch( - predictions=[summary_hf[beam_idx][i]], - references=[datapoint['highlights'][i]]) - if output_dir is not None: - # yapf: disable - for i in range(len(summary_hf[0])): - for beam_idx in range(num_beams): - with (output_dir / 'hf.out').open('a') as f: - f.write(f'[{data_point_idx + i}] [Beam {beam_idx}] {summary_hf[beam_idx][i]}\n') - # yapf: enable - - logger.debug('-' * 100) - logger.debug(f"Article : {datapoint['article']}") - if test_trt_llm: - logger.debug(f'TensorRT-LLM Summary: {summary_tensorrt_llm}') - if test_hf: - logger.debug(f'HF Summary: {summary_hf}') - logger.debug(f"highlights : {datapoint['highlights']}") - - data_point_idx += max_batch_size - ite_count += 1 - - if runtime_rank == 0: - if test_trt_llm: - np.random.seed(0) # rouge score use sampling to compute the score - latency = profiler.elapsed_time_in_sec("tensorrt_llm") - logger.info(f'TensorRT-LLM (total latency: {latency} sec)') - for beam_idx in range(num_beams): - logger.info(f"TensorRT-LLM beam {beam_idx} result") - 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].mid[2]*100}' - ) - - if args.check_accuracy and beam_idx == 0: - assert computed_metrics_tensorrt_llm['rouge1'].mid[ - 2] * 100 > args.tensorrt_llm_rouge1_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)' - ) - for beam_idx in range(num_beams): - logger.info(f"HF beam {beam_idx} result") - computed_metrics_hf = metric_hf[beam_idx].compute() - for key in computed_metrics_hf.keys(): - logger.info( - f' {key} : {computed_metrics_hf[key].mid[2]*100}') - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--hf_model_location', - type=str, - default='falcon/rw-1b', - help='Directory where a HF model checkpoint locates.') - parser.add_argument('--test_hf', action='store_true') - parser.add_argument('--test_trt_llm', action='store_true') - parser.add_argument('--data_type', - type=str, - choices=['float32', 'float16', 'bfloat16'], - default='float16') - parser.add_argument('--dataset_path', type=str, default='') - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument('--engine_dir', type=str, default='falcon_outputs') - parser.add_argument('--batch_size', type=int, default=1) - parser.add_argument('--max_ite', type=int, default=20) - parser.add_argument('--max_kv_cache_len', - type=int, - default=None, - help='The max kv cache length. \ - If the final sequence length exceeds the kv cache length, we will enable cyclic kv cache. \ - If it is set to None, we will use the max sequence length.') - parser.add_argument('--check_accuracy', action='store_true') - parser.add_argument('--tensorrt_llm_rouge1_threshold', - type=float, - default=15.0) - parser.add_argument('--num_beams', type=int, default=1) - parser.add_argument('--top_k', type=int, default=1) - parser.add_argument('--output_len', type=int, default=100) - parser.add_argument('--debug', action='store_true') - 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( - '--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.") - - args = parser.parse_args() - - main(args) diff --git a/examples/falcon/weight.py b/examples/falcon/weight.py index 8aadf2660568..60b46270e317 100644 --- a/examples/falcon/weight.py +++ b/examples/falcon/weight.py @@ -520,3 +520,176 @@ def get_scaling_factors( f'Expect scaling factor {k} of length {num_layers}, got {len(v)}' return scaling_factor + + +def load_from_awq_falcon( + tensorrt_llm_falcon: tensorrt_llm.models.FalconForCausalLM, + quant_ckpt_path, + mapping=Mapping(), + dtype="float16"): + tensorrt_llm.logger.info( + 'Loading weights from groupwise AWQ Falcon checkpoint...') + tik = time.time() + + packer = torch.ops.fastertransformer.pack_int8_tensor_to_packed_int4 + preprocessor = torch.ops.fastertransformer.preprocess_weights_for_mixed_gemm + torch_dtype = tensorrt_llm._utils.str_dtype_to_torch(dtype) + + if quant_ckpt_path.endswith(".npz"): + awq_falcon = np.load(quant_ckpt_path) + awq_prefix = "_np:" + awq_suffix_list = [ + ":weight", + ":weights_scaling_factor", + ":prequant_scaling_factor", + ] + awq_key_list = [ + "vocab_embedding:weight", # embedding + "lm_head", # lm_head + "final_layernorm", # ln_f + "attention:qkv:", # attention.qkv + "attention:dense", # attention.dense + "mlp:proj", # mlp.proj + "mlp:fc", # mlp.fc + "input_layernorm", # input_layernorm.weight + "mlp_layernorm", # mlp_layernorm.weight + ] + split_sym = ":" + AMMO_WEIGHT_SCALING_FACTOR_COEFF = 7 + + def load(key): + v = torch.from_numpy(awq_falcon[awq_prefix + key]).to(torch_dtype) + if "weights_scaling_factor" in key: + v *= AMMO_WEIGHT_SCALING_FACTOR_COEFF # For AMMO *.npz checkpoints + return v + + group_size = load("layers:0:attention:dense:weight").numel() // load( + "layers:0:attention:dense:weights_scaling_factor").numel() + else: + raise ValueError("Unsupported AWQ quantized checkpoint format") + + def torch_split(v, dim): + if v.shape[dim] % mapping.tp_size != 0: + tensorrt_llm.logger.error( + "Current weight shape is invalid for mapping.tp_size=" + + str(mapping.tp_size)) + raise ValueError("Invalid TP size") + return v.split(v.shape[dim] // mapping.tp_size, + dim=dim)[mapping.tp_rank] + + def AWQ_quantize_pack_preprocess(weight, scale): + weight /= scale.repeat_interleave(group_size, dim=0) + qweight_int8 = torch.clamp(torch.round(weight.cuda()).char(), -8, 7) + int4_weight = preprocessor(packer(qweight_int8.cpu()), torch.quint4x2) + return int4_weight.view(torch.int8) + + def process_and_assign_weight(mOp, v, tp_dim=0): + weight = v[0].T.contiguous() + [k, n] = weight.shape + weight = torch_split(weight, tp_dim) + amax = v[1].reshape((n, k // group_size)).T.contiguous() + amax = torch_split(amax, tp_dim) + pre_quant_scale = v[2].reshape((1, k)) + if tp_dim == 0: + pre_quant_scale = torch_split(pre_quant_scale, 1) + scale = amax / 8.0 + mOp.qweight.value = AWQ_quantize_pack_preprocess(weight, scale) + mOp.scale.value = scale.to(torch_dtype) + mOp.pre_quant_scale.value = pre_quant_scale.to(torch_dtype) + + def get_scale(weight): + [k, n] = weight.shape + weight_t = weight.T.contiguous() + weight_t = weight_t.reshape(n, k // group_size, group_size) + weight_t = torch.abs(weight_t.reshape(-1, group_size)) + amax, idx = weight_t.max(1) + amax = amax.reshape(n, k // group_size).T.contiguous() + scale = amax / 8 + return scale + + def process_and_assign_qkv_weight(prefix, mOp): + q_weight = load(prefix + "q" + awq_suffix_list[0]) + k_weight = load(prefix + "k" + awq_suffix_list[0]) + v_weight = load(prefix + "v" + awq_suffix_list[0]) + dim_k = q_weight.shape[0] + q_weight = torch_split(q_weight, 1) + k_weight = torch_split(k_weight, 1) + v_weight = torch_split(v_weight, 1) + qkv_pre_quant_scale = load(prefix + "q" + awq_suffix_list[2]).reshape( + (1, dim_k)) + qkv_weights = torch.cat((q_weight, k_weight, v_weight), dim=1) + qkv_scale = get_scale(qkv_weights) + + mOp.pre_quant_scale.value = qkv_pre_quant_scale.to(torch_dtype) + mOp.qweight.value = AWQ_quantize_pack_preprocess(qkv_weights, qkv_scale) + mOp.scale.value = qkv_scale.to(torch_dtype) + + # Load weights from AWQ checkpoint into TRT-LLM module + # 1. embedding + v = load(awq_key_list[0]) + # TRT-LLM requires vocab_size to be multiple of 64 for successful GEMM + if v.shape[0] % 64 != 0: + v = torch.nn.functional.pad(v, [0, 0, 0, 64 - v.shape[0] % 64]) + if mapping.is_first_pp_rank(): + tensorrt_llm_falcon.embedding.weight.value = v.to(torch_dtype) + + # 2. lm_head + v = [load(awq_key_list[1] + suf) for suf in awq_suffix_list] + if v[0].shape[0] % 64 != 0: + v[0] = torch.nn.functional.pad(v[0], [0, 0, 0, 64 - v[0].shape[0] % 64]) + v[1] = torch.nn.functional.pad(v[1], [0, 0, 0, 64 - v[1].shape[0] % 64], + value=1) + if mapping.is_last_pp_rank(): + process_and_assign_weight(tensorrt_llm_falcon.lm_head, v, 1) + + # 3. ln_f + v_weight = load(awq_key_list[2] + split_sym + "weight") + v_bias = load(awq_key_list[2] + split_sym + "bias") + if mapping.is_last_pp_rank(): + tensorrt_llm_falcon.ln_f.weight.value = v_weight.to(torch_dtype) + tensorrt_llm_falcon.ln_f.bias.value = v_bias.to(torch_dtype) + + # 4. Weights inside each layer + num_hidden_layers = tensorrt_llm_falcon.num_layers + layers_per_pipeline_stage = num_hidden_layers // mapping.pp_size + layers_range = list( + range(mapping.pp_rank * layers_per_pipeline_stage, + (mapping.pp_rank + 1) * layers_per_pipeline_stage, 1)) + + for l in layers_range: + layer_idx = l - mapping.pp_rank * layers_per_pipeline_stage + prefix = "layers" + split_sym + str(layer_idx) + split_sym + tensorrt_llm.logger.info(f'Process weights in layer: {layer_idx}') + layer = tensorrt_llm_falcon.layers[layer_idx] + + # 4.1 attention.qkv + process_and_assign_qkv_weight(prefix + awq_key_list[3], + layer.attention.qkv) + + # 4.2 attention.dense + v = [load(prefix + awq_key_list[4] + suf) for suf in awq_suffix_list] + process_and_assign_weight(layer.attention.dense, v, 0) + + # 4.3 mlp.proj + v = [load(prefix + awq_key_list[5] + suf) for suf in awq_suffix_list] + process_and_assign_weight(layer.mlp.proj, v, 0) + + # 4.4 mlp.fc + v = [load(prefix + awq_key_list[6] + suf) for suf in awq_suffix_list] + process_and_assign_weight(layer.mlp.fc, v, 1) + + # 4.5 input_layernorm + v = load(prefix + awq_key_list[7] + split_sym + "weight") + layer.input_layernorm.weight.value = v.to(torch_dtype) + v = load(prefix + awq_key_list[7] + split_sym + "bias") + layer.input_layernorm.bias.value = v.to(torch_dtype) + + # 4.6 mlp_layernorm + v = load(prefix + awq_key_list[8] + split_sym + "weight") + layer.mlp_layernorm.weight.value = v.to(torch_dtype) + v = load(prefix + awq_key_list[8] + split_sym + "bias") + layer.mlp_layernorm.bias.value = v.to(torch_dtype) + + tok = time.time() + t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) + tensorrt_llm.logger.info(f'Weights loaded. Elapsed time: {t}') diff --git a/examples/gpt/README.md b/examples/gpt/README.md index 1de29e696d97..366404583225 100644 --- a/examples/gpt/README.md +++ b/examples/gpt/README.md @@ -5,14 +5,13 @@ multiple GPUs or multiple nodes with multiple GPUs. ## Overview -The TensorRT-LLM GPT implementation can be found in [`tensorrt_llm/models/gpt/model.py`](../../tensorrt_llm/models/gpt/model.py). The TensorRT-LLM GPT example -code is located in [`examples/gpt`](./). There are four main files in that folder: +The TensorRT-LLM GPT implementation can be found in [`tensorrt_llm/models/gpt/model.py`](../../tensorrt_llm/models/gpt/model.py). The TensorRT-LLM GPT example code is located in [`examples/gpt`](./). There are four main files: * [`hf_gpt_convert.py`](./hf_gpt_convert.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the [FasterTransformer (FT)](https://github.com/NVIDIA/FasterTransformer) format, * [`build.py`](./build.py) to build the [TensorRT](https://developer.nvidia.com/tensorrt) engine(s) needed to run the GPT model, * [`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/cnn_dailymail) dataset using the model. + * and a shared [`../summarize.py`](../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/cnn_dailymail) dataset using the model. ## Support Matrix * FP16 @@ -257,17 +256,18 @@ python3 build.py --model_dir=./c-model/gpt2/fp16/1-gpu \ --hidden_act gelu ``` -The summarization can be done using the [`summarize.py`](./summarize.py) script as follows: +The summarization can be done using the [`../summarize.py`](../summarize.py) script as follows: ```bash # Run the summarization task. -python3 summarize.py --engine_dir trt_engine/gpt2/fp16/1-gpu \ - --test_hf \ - --batch_size 1 \ - --test_trt_llm \ - --hf_model_location=gpt2 \ - --check_accuracy \ - --tensorrt_llm_rouge1_threshold=14 +python3 ../summarize.py --engine_dir trt_engine/gpt2/fp16/1-gpu \ + --hf_model_dir gpt2 \ + --test_trt_llm \ + --test_hf \ + --batch_size 1 \ + --check_accuracy \ + --tensorrt_llm_rouge1_threshold=14 \ + --no_add_special_tokens ``` ## SmoothQuant @@ -534,5 +534,5 @@ python3 hf_gpt_convert.py -i gpt2 -o ./c-model/gpt2 --tensor-parallelism 2 --sto python3 build.py --model_dir=./c-model/gpt2/2-gpu --dtype bfloat16 --world_size=2 --remove_input_padding --use_gpt_attention_plugin --use_gemm_plugin --parallel_build --max_input_len 1000 --use_parallel_embedding --embedding_sharding_dim 0 --use_lookup_plugin --use_embedding_sharing --output_dir=trt_engine/gpt2/bfloat16/2-gpu -mpirun -np 2 python3 summarize.py --engine_dir trt_engine/gpt2/bfloat16/2-gpu --batch_size 10 --test_trt_llm --check_accuracy --tensorrt_llm_rouge1_threshold=14 --dataset_path ./dataset +mpirun -np 2 python3 ../summarize.py --engine_dir trt_engine/gpt2/bfloat16/2-gpu --hf_model_dir gpt2 --batch_size 10 --test_trt_llm --check_accuracy --tensorrt_llm_rouge1_threshold=14 --dataset_path ./dataset --no_add_special_tokens ``` diff --git a/examples/gpt/build.py b/examples/gpt/build.py index 7e07cbb0c73e..820e0b1b8068 100644 --- a/examples/gpt/build.py +++ b/examples/gpt/build.py @@ -89,6 +89,8 @@ def parse_arguments(args): parser.add_argument('--n_embd', type=int, default=1024) parser.add_argument('--n_head', type=int, default=16) parser.add_argument('--hidden_act', type=str, default='gelu') + parser.add_argument('--rotary_base', type=float, default=10000.0) + parser.add_argument('--rotary_scaling', nargs=2, type=str, default=None) parser.add_argument( '--rotary_pct', type=float, @@ -302,6 +304,14 @@ def parse_arguments(args): action='store_true', help= 'Activates latency-optimized algorithm for all-reduce instead of NCCL.') + parser.add_argument( + '--use_lora_plugin', + nargs='?', + const=None, + default=False, + choices=['float16', 'float32', 'bfloat16'], + help="Activates the lora plugin which enables embedding sharing.") + args = parser.parse_args(args) logger.set_level(args.log_level) @@ -332,7 +342,7 @@ def parse_arguments(args): args.multi_query_mode = multi_query_mode plugins_args = [ 'use_gpt_attention_plugin', 'use_gemm_plugin', 'use_layernorm_plugin', - 'use_lookup_plugin' + 'use_lookup_plugin', 'use_lora_plugin' ] for plugin_arg in plugins_args: if getattr(args, plugin_arg) is None: @@ -379,6 +389,16 @@ def parse_arguments(args): if args.enable_fp8: args.quant_mode = args.quant_mode.set_fp8_qdq() + if args.rotary_scaling is not None: + assert args.use_gpt_attention_plugin, "RoPE scaling is only supported through GPT attention plugin." + rotary_scaling = { + "type": args.rotary_scaling[0], + "factor": float(args.rotary_scaling[1]) + } + assert rotary_scaling["type"] in ["linear", "dynamic"] + assert rotary_scaling["factor"] > 1.0 + args.rotary_scaling = rotary_scaling + if args.max_num_tokens is not None: assert args.enable_context_fmha @@ -428,6 +448,8 @@ def build_rank_engine(builder: Builder, position_embedding_type=PositionEmbeddingType.learned_absolute if args.rotary_pct == 0.0 else PositionEmbeddingType.rope_gpt_neox, rotary_embedding_percentage=args.rotary_pct, + rotary_base=args.rotary_base, + rotary_scaling=args.rotary_scaling, dtype=kv_dtype, logits_dtype=args.logits_dtype, mapping=Mapping(world_size=args.world_size, @@ -437,7 +459,7 @@ def build_rank_engine(builder: Builder, apply_query_key_layer_scaling, quant_mode=args.quant_mode, bias=args.bias, - multi_query_mode=args.multi_query_mode, + num_kv_heads=1 if args.multi_query_mode else args.n_head, use_prompt_tuning=args.max_prompt_embedding_table_size > 0, use_parallel_embedding=args.use_parallel_embedding, embedding_sharding_dim=args.embedding_sharding_dim, @@ -497,6 +519,8 @@ def build_rank_engine(builder: Builder, network.plugin_config.enable_remove_input_padding() if args.paged_kv_cache: network.plugin_config.enable_paged_kv_cache(args.tokens_per_block) + if args.use_lora_plugin: + network.plugin_config.set_lora_plugin(dtype=args.use_lora_plugin) # Quantization plugins. if args.use_smooth_quant: @@ -582,12 +606,12 @@ def build(rank, args): max_position_embeddings=args.n_positions, apply_query_key_layer_scaling=apply_query_key_layer_scaling, max_batch_size=args.max_batch_size, + max_beam_width=args.max_beam_width, max_input_len=args.max_input_len, max_output_len=args.max_output_len, max_num_tokens=args.max_num_tokens, int8=int8_trt_flag, opt_level=args.builder_opt, - multi_query_mode=args.multi_query_mode, strongly_typed=args.strongly_typed, max_prompt_embedding_table_size=args. max_prompt_embedding_table_size, diff --git a/examples/gpt/requirements.txt b/examples/gpt/requirements.txt index f46bff310071..02895a917bb8 100644 --- a/examples/gpt/requirements.txt +++ b/examples/gpt/requirements.txt @@ -1,2 +1,4 @@ datasets~=2.14.5 +evaluate~=0.4.1 rouge_score~=0.1.2 +SentencePiece~=0.1.99 diff --git a/examples/gpt/run.py b/examples/gpt/run.py index 5cf4ebab399f..67426437cf93 100644 --- a/examples/gpt/run.py +++ b/examples/gpt/run.py @@ -23,7 +23,7 @@ import tensorrt_llm from tensorrt_llm.quantization import QuantMode -from tensorrt_llm.runtime import ModelConfig, SamplingConfig +from tensorrt_llm.runtime import LoraManager, ModelConfig, SamplingConfig from build import get_engine_name # isort:skip @@ -55,6 +55,7 @@ def read_config(config_path: Path): 'gather_all_token_logits'] use_custom_all_reduce = config['plugin_config']['use_custom_all_reduce'] quant_mode = QuantMode(config['builder_config']['quant_mode']) + lora_plugin = config['plugin_config']['lora_plugin'] model_config = ModelConfig( num_heads=num_heads, @@ -70,7 +71,8 @@ def read_config(config_path: Path): dtype=dtype, quant_mode=quant_mode, gather_all_token_logits=gather_all_token_logits, - use_custom_all_reduce=use_custom_all_reduce) + use_custom_all_reduce=use_custom_all_reduce, + lora_plugin=lora_plugin) dtype = config['builder_config']['precision'] max_input_len = config['builder_config']['max_input_len'] @@ -294,11 +296,26 @@ def generate( model_config.remove_input_padding) max_input_length = torch.max(input_lengths).item() + + if model_config.lora_plugin: + lora_manager = LoraManager(model_dir=engine_dir, + model_config=model_config) + # an example under batch size 3 + lora_uids = [ + None, "de05b696-711c-4d7c-983d-e1fb9c1a618a", + "5857e9da-c7cb-4541-a419-fb364e6151a2" + ] + else: + lora_manager = None + lora_uids = None + decoder.setup(input_lengths.size(0), max_input_length, max_output_len, beam_width=num_beams, - max_kv_cache_length=max_kv_cache_len) + max_kv_cache_length=max_kv_cache_len, + lora_manager=lora_manager, + lora_uids=lora_uids) ptuning_args = [] if model_config.max_prompt_embedding_table_size == 0 else ptuning_setup( prompt_table, dtype, model_config.hidden_size, tasks, input_ids, diff --git a/examples/gpt/summarize.py b/examples/gpt/summarize.py deleted file mode 100644 index 75528adcac99..000000000000 --- a/examples/gpt/summarize.py +++ /dev/null @@ -1,539 +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 copy -import json -from pathlib import Path - -import numpy as np -import torch -from datasets import load_dataset, load_metric -from transformers import AutoModelForCausalLM, AutoTokenizer, T5Tokenizer - -import tensorrt_llm -import tensorrt_llm.profiler as profiler -from tensorrt_llm.logger import logger -from tensorrt_llm.quantization import QuantMode -from tensorrt_llm.tools.ppl import ppl - -from build import find_engines # isort:skip - - -def TRTGPT(args, config): - dtype = config['builder_config']['precision'] - world_size = config['builder_config']['tensor_parallel'] - assert world_size == tensorrt_llm.mpi_world_size(), \ - f'Engine world size ({world_size}) != Runtime world size ({tensorrt_llm.mpi_world_size()})' - - world_size = config['builder_config']['tensor_parallel'] - num_heads = config['builder_config']['num_heads'] // world_size - hidden_size = config['builder_config']['hidden_size'] // world_size - vocab_size = config['builder_config']['vocab_size'] - num_layers = config['builder_config']['num_layers'] - use_gpt_attention_plugin = bool( - config['plugin_config']['gpt_attention_plugin']) - remove_input_padding = config['plugin_config']['remove_input_padding'] - multi_query_mode = config['builder_config']['multi_query_mode'] - num_kv_heads = 1 if multi_query_mode else num_heads - paged_kv_cache = config['plugin_config']['paged_kv_cache'] - tokens_per_block = config['plugin_config']['tokens_per_block'] - gather_all_token_logits = config['builder_config'].get( - 'gather_all_token_logits', False) - assert not (args.eval_ppl and not gather_all_token_logits), \ - "PPL evaluation requires engine built with gather_all_token_logits enabled" - - use_custom_all_reduce = config['plugin_config']['use_custom_all_reduce'] - quant_mode = QuantMode(config['builder_config'].get('quant_mode', 0)) - - model_config = tensorrt_llm.runtime.ModelConfig( - vocab_size=vocab_size, - num_layers=num_layers, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - hidden_size=hidden_size, - gpt_attention_plugin=use_gpt_attention_plugin, - remove_input_padding=remove_input_padding, - tokens_per_block=tokens_per_block, - paged_kv_cache=paged_kv_cache, - dtype=dtype, - quant_mode=quant_mode, - gather_all_token_logits=gather_all_token_logits, - use_custom_all_reduce=use_custom_all_reduce, - ) - - runtime_rank = tensorrt_llm.mpi_rank() - runtime_mapping = tensorrt_llm.Mapping(world_size, - runtime_rank, - tp_size=world_size) - torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node) - - serialize_path = find_engines(args.engine_dir, - dtype=dtype, - tp_size=world_size, - rank=runtime_rank)[0] - - tensorrt_llm.logger.set_level(args.log_level) - - with open(serialize_path, 'rb') as f: - engine_buffer = f.read() - decoder = tensorrt_llm.runtime.GenerationSession(model_config, - engine_buffer, - runtime_mapping) - - return decoder - - -def main(args): - 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 - hf_model_location = args.hf_model_location - - if args.vocab_file is not None: - tokenizer = T5Tokenizer(vocab_file=args.vocab_file, padding_side='left') - else: - tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, - padding_side='left') - - tokenizer.pad_token = tokenizer.eos_token - - if args.eval_type == 'code_completion': - dataset_name = "openai_humaneval" - dataset_revision = None - dataset_input_key = 'prompt' - dataset_output_key = 'canonical_solution' - elif args.eval_type == 'summarize': - dataset_name = "ccdv/cnn_dailymail" - dataset_revision = "3.0.0" - dataset_input_key = 'article' - dataset_output_key = 'highlights' - dataset = load_dataset(dataset_name, - dataset_revision, - cache_dir=args.dataset_path) - - config_path = str(args.engine_dir / 'config.json') - with open(config_path, 'r') as f: - config = json.load(f) - - max_batch_size = args.batch_size - - # runtime parameters - # repetition_penalty = 1 - top_k = args.top_k - output_len = args.output_len - test_token_num = 923 - # top_p = 0.0 - # random_seed = 5 - temperature = 1 - num_beams = args.num_beams - length_penalty = args.length_penalty - - pad_id = tokenizer.encode(tokenizer.pad_token, add_special_tokens=False)[0] - end_id = tokenizer.encode(tokenizer.eos_token, add_special_tokens=False)[0] - - if test_trt_llm: - tensorrt_llm_gpt = TRTGPT(args, config) - - if test_hf: - model = AutoModelForCausalLM.from_pretrained(hf_model_location, - trust_remote_code=True) - model.cuda() - if args.data_type == 'fp16': - model.half() - elif args.data_type == 'bf16': - model.bfloat16() - - def eval_tensorrt_llm(datapoint, eval_type='summarize'): - batch_size = len(datapoint) - append_str = ' TL;DR: ' if eval_type == 'summarize' else '' - line = copy.copy(datapoint) - line_encoded = [] - input_lengths = [] - for i in range(batch_size): - line[i] = line[i] + append_str - - line[i] = line[i].strip() - line[i] = line[i].replace(" n't", "n't") - - input_id = tokenizer.encode(line[i], - return_tensors='pt', - add_special_tokens=False).type( - torch.int32) - input_id = input_id[:, -test_token_num:] - - line_encoded.append(input_id) - input_lengths.append(input_id.shape[-1]) - - max_length = max(input_lengths) - - if tensorrt_llm_gpt.remove_input_padding: - line_encoded = [ - torch.tensor(t, dtype=torch.int32).cuda() for t in line_encoded - ] - else: - # do padding, should move outside the profiling to prevent the overhead - for i in range(batch_size): - pad_size = max_length - input_lengths[i] - - pad = torch.ones([1, pad_size]).type(torch.int32) * pad_id - line_encoded[i] = torch.cat( - [torch.tensor(line_encoded[i], dtype=torch.int32), pad], - axis=-1) - - line_encoded = torch.cat(line_encoded, axis=0).cuda() - input_lengths = torch.tensor(input_lengths, - dtype=torch.int32).cuda() - - sampling_config = tensorrt_llm.runtime.SamplingConfig( - end_id=end_id, - pad_id=pad_id, - top_k=top_k, - num_beams=num_beams, - length_penalty=length_penalty) - - with torch.no_grad(): - tensorrt_llm_gpt.setup(batch_size, - max_context_length=max_length, - max_new_tokens=output_len, - beam_width=num_beams, - max_kv_cache_length=args.max_kv_cache_len) - - if tensorrt_llm_gpt.remove_input_padding: - outputs = tensorrt_llm_gpt.decode_batch( - line_encoded, - sampling_config, - output_sequence_lengths=True, - return_dict=True) - else: - outputs = tensorrt_llm_gpt.decode(line_encoded, - input_lengths, - sampling_config, - output_sequence_lengths=True, - return_dict=True) - torch.cuda.synchronize() - - # Extract a list of tensors of shape beam_width x output_ids. - if tensorrt_llm_gpt.mapping.is_first_pp_rank(): - 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(batch_size) - ] - - ppls = [] - if args.eval_ppl: - seq_lens = outputs['sequence_lengths'] - context_logits = outputs['context_logits'] - if tensorrt_llm_gpt.remove_input_padding: - context_logits = context_logits.flatten(end_dim=1) - seg_points = [0] + np.cumsum(input_lengths).tolist() - context_logits = [ - context_logits[s:e] - for s, e in zip(seg_points[:-1], seg_points[1:]) - ] - else: - context_logits = [ - context_logits[bidx, :input_lengths[bidx]] - for bidx in range(batch_size) - ] - - # Remove the first generation logits which are same to last context logits - # Step dim at 1 - generation_logits = torch.stack( - outputs['generation_logits'][1:], dim=1) - for bidx in range(batch_size): - # [batch, beam, step] - curr_len = seq_lens[bidx, 0] - curr_ctx_len = input_lengths[bidx] - curr_gen_len = curr_len - curr_ctx_len - - curr_ids = output_ids[bidx, 0, 1:curr_len] - curr_logits = torch.cat([ - context_logits[bidx], - generation_logits[bidx, :curr_gen_len - 1] - ], - dim=0) - curr_ppl = ppl(curr_logits, curr_ids) - ppls.append(curr_ppl) - logger.debug( - f"TensorRT-LLM PPL: {curr_ppl:.3f} | Generation length: {curr_gen_len}" - ) - - return output_beams_list, output_ids[:, :, - max_length:].tolist(), ppls - return [], [], [] - - def eval_hf(datapoint, eval_type='summarize'): - batch_size = len(datapoint) - append_str = ' TL;DR: ' if eval_type == 'summarize' else '' - if batch_size > 1: - logger.warning( - f"HF does not support batch_size > 1 to verify correctness due to padding and attention mask. Current batch size is {batch_size}" - ) - - line = copy.copy(datapoint) - line_encoded = [] - input_lengths = [] - for i in range(batch_size): - line[i] = line[i] + append_str - - line[i] = line[i].strip() - line[i] = line[i].replace(" n't", "n't") - - input_id = tokenizer.encode(line[i], - return_tensors='pt', - add_special_tokens=False).type( - torch.int64) - input_id = input_id[:, -test_token_num:] - - line_encoded.append(input_id) - input_lengths.append(input_id.shape[-1]) - - max_length = max(input_lengths) - - for i in range(batch_size): - pad_size = max_length - input_lengths[i] - - pad = torch.ones([1, pad_size]).type(torch.int64) * pad_id - line_encoded[i] = torch.cat( - [pad, torch.tensor(line_encoded[i], dtype=torch.int64)], - axis=-1) - - line_encoded = torch.cat(line_encoded, axis=0).cuda() - - with torch.no_grad(): - outputs = model.generate(line_encoded, - max_length=len(line_encoded[0]) + - output_len, - top_k=top_k, - temperature=temperature, - eos_token_id=tokenizer.eos_token_id, - pad_token_id=tokenizer.pad_token_id, - num_beams=num_beams, - num_return_sequences=num_beams, - early_stopping=True, - length_penalty=length_penalty, - output_scores=True, - return_dict_in_generate=True) - # model.generate cannot return context logits? - context_outputs = model(line_encoded) - - output_ids = outputs['sequences'] - tokens_list = output_ids[:, len(line_encoded[0]):].tolist() - output_ids = output_ids.reshape([batch_size, num_beams, -1]) - output_lines_list = [ - tokenizer.batch_decode(output_ids[:, i, len(line_encoded[0]):], - skip_special_tokens=True) - for i in range(num_beams) - ] - - ppls = [] - if args.eval_ppl and batch_size == 1: - # Only for batch size of 1 - seq_lens = [output_ids.size(-1) for _ in range(batch_size)] - context_logits = context_outputs['logits'] - # Remove the first generation logits which are same to last context logits - generation_logits = torch.stack(outputs['scores'][1:], dim=1) - - ppls = [] - for bidx in range(batch_size): - curr_len = seq_lens[bidx] - curr_ctx_len = input_lengths[bidx] - curr_gen_len = curr_len - curr_ctx_len - - curr_ids = output_ids[bidx, 0, 1:curr_len] - curr_logits = torch.cat([ - context_logits[bidx], - generation_logits[bidx, :curr_gen_len - 1] - ], - dim=0) - curr_ppl = ppl(curr_logits, curr_ids) - ppls.append(curr_ppl) - logger.debug( - f"HF PPL: {curr_ppl:.3f} | Generation length: {curr_gen_len}" - ) - - return output_lines_list, tokens_list, ppls - - if test_trt_llm: - datapoint = dataset['test'][0:1] - output, *_ = eval_tensorrt_llm(datapoint[dataset_input_key], - eval_type=args.eval_type) - if runtime_rank == 0: - 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( - "---------------------------------------------------------") - - if test_hf: - datapoint = dataset['test'][0:1] - output, *_ = eval_hf(datapoint[dataset_input_key], - eval_type=args.eval_type) - 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("---------------------------------------------------------") - - metric_tensorrt_llm = [load_metric("rouge") for _ in range(num_beams)] - metric_hf = [load_metric("rouge") for _ in range(num_beams)] - for i in range(num_beams): - metric_tensorrt_llm[i].seed = 0 - metric_hf[i].seed = 0 - ppls_trt_llm, ppls_hf = [], [] - - ite_count = 0 - data_point_idx = 0 - while (data_point_idx < len(dataset['test'])) 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['test'][data_point_idx:(data_point_idx + - max_batch_size)] - - if test_trt_llm: - profiler.start('tensorrt_llm') - output_tensorrt_llm, _, curr_ppls_trt_llm = eval_tensorrt_llm( - datapoint[dataset_input_key]) - profiler.stop('tensorrt_llm') - - if test_hf: - profiler.start('hf') - output_hf, _, curr_ppls_hf = eval_hf(datapoint[dataset_input_key]) - profiler.stop('hf') - - if runtime_rank == 0: - if test_trt_llm: - for batch_idx in range(len(output_tensorrt_llm)): - for beam_idx in range(num_beams): - metric_tensorrt_llm[beam_idx].add_batch( - predictions=[ - output_tensorrt_llm[batch_idx][beam_idx] - ], - references=[ - datapoint[dataset_output_key][batch_idx] - ]) - ppls_trt_llm.extend(curr_ppls_trt_llm) - if test_hf: - for beam_idx in range(num_beams): - 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] - ]) - ppls_hf.extend(curr_ppls_hf) - - logger.debug('-' * 100) - logger.debug(f"Input : {datapoint[dataset_input_key]}") - if test_trt_llm: - logger.debug(f'TensorRT-LLM Output: {output_tensorrt_llm}') - if test_hf: - logger.debug(f'HF Output: {output_hf}') - logger.debug(f"highlights : {datapoint[dataset_output_key]}") - - data_point_idx += max_batch_size - ite_count += 1 - - if runtime_rank == 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)' - ) - for beam_idx in range(num_beams): - logger.info(f"TensorRT-LLM beam {beam_idx} result") - 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].mid[2]*100}' - ) - - if args.check_accuracy and beam_idx == 0: - assert computed_metrics_tensorrt_llm['rouge1'].mid[ - 2] * 100 > args.tensorrt_llm_rouge1_threshold - if args.eval_ppl: - logger.info(f" Per-token perplexity: {np.mean(ppls_trt_llm)}") - 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)' - ) - for beam_idx in range(num_beams): - logger.info(f"HF beam {beam_idx} result") - computed_metrics_hf = metric_hf[beam_idx].compute() - for key in computed_metrics_hf.keys(): - logger.info( - f' {key} : {computed_metrics_hf[key].mid[2]*100}') - if args.eval_ppl and args.batch_size == 1: - logger.info(f" Per-token perplexity: {np.mean(ppls_hf)}") - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--hf_model_location', type=str, default='gpt2') - parser.add_argument( - '--tokenizer', - default=None, - help='tokenizer path; defaults to hf_model_location if left unspecified' - ) - parser.add_argument('--vocab_file') - parser.add_argument('--test_hf', action='store_true') - parser.add_argument('--test_trt_llm', action='store_true') - parser.add_argument('--data_type', - type=str, - choices=['fp32', 'fp16', 'bf16'], - default='fp32') - parser.add_argument('--dataset_path', type=str, default='') - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument('--engine_dir', type=Path, default='gpt_outputs') - 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_kv_cache_len', - type=int, - default=None, - help='The max kv cache length. \ - If the final sequence length exceeds the kv cache length, we will enable cyclic kv cache. \ - If it is set to None, we will use the max sequence length.') - parser.add_argument('--check_accuracy', action='store_true') - parser.add_argument('--tensorrt_llm_rouge1_threshold', - type=float, - default=15.0) - parser.add_argument('--num_beams', type=int, default=1) - parser.add_argument('--top_k', type=int, default=1) - parser.add_argument('--eval_type', - type=str, - default='summarize', - choices=['summarize', 'code_completion']) - parser.add_argument('--length_penalty', type=float, default=1.0) - parser.add_argument('--eval_ppl', action='store_true') - - args = parser.parse_args() - if args.tokenizer == None: - args.tokenizer = args.hf_model_location - main(args) diff --git a/examples/gpt/weight.py b/examples/gpt/weight.py index 3d7c9a41bf3b..ec9344c2ab09 100644 --- a/examples/gpt/weight.py +++ b/examples/gpt/weight.py @@ -236,21 +236,22 @@ def set_smoothquant_scale_factors(module, tensor_parallel) if not multi_query_mode else ( n_embd // tensor_parallel + (n_embd // n_head) * 2) - tensorrt_llm_gpt.layers[i].input_layernorm.weight.value = (fromfile( + gpt_layer = tensorrt_llm_gpt.layers[i] + gpt_layer.input_layernorm.weight.value = (fromfile( dir_path, 'model.layers.' + str(i) + '.input_layernorm.weight.bin')) - tensorrt_llm_gpt.layers[i].input_layernorm.bias.value = (fromfile( + gpt_layer.input_layernorm.bias.value = (fromfile( dir_path, 'model.layers.' + str(i) + '.input_layernorm.bias.bin')) t = fromfile( dir_path, 'model.layers.' + str(i) + '.attention.query_key_value.weight.' + suffix, [n_embd, c_attn_out_dim], w_type) if t is not None: - dst = tensorrt_llm_gpt.layers[i].attention.qkv.weight + dst = gpt_layer.attention.qkv.weight if use_smooth_quant: dst.value = np.ascontiguousarray(np.transpose(t, [1, 0])) set_smoothquant_scale_factors( - tensorrt_llm_gpt.layers[i].attention.qkv, - tensorrt_llm_gpt.layers[i].input_layernorm.scale_to_int, + gpt_layer.attention.qkv, + gpt_layer.input_layernorm.scale_to_int, dir_path, 'model.layers.' + str(i) + '.attention.query_key_value.', [1, c_attn_out_dim], @@ -272,7 +273,7 @@ def set_smoothquant_scale_factors(module, dir_path, 'model.layers.' + str(i) + '.attention.query_key_value.bias.' + str(rank) + '.bin') if t is not None: - dst = tensorrt_llm_gpt.layers[i].attention.qkv.bias + dst = gpt_layer.attention.qkv.bias dst.value = np.ascontiguousarray(t) if enable_fp8_qdq: tensorrt_llm_gpt.layers[ @@ -288,21 +289,21 @@ def set_smoothquant_scale_factors(module, i].attention.kv_quant_orig_scale.value = np.array( [1.0 / scaling_factors['qkv_output'][i]], dtype=np.float32) - dst = tensorrt_llm_gpt.layers[i].attention.dense.weight + dst = gpt_layer.attention.dense.weight t = fromfile( dir_path, 'model.layers.' + str(i) + '.attention.dense.weight.' + suffix, [n_embd // tensor_parallel, n_embd], w_type) if use_smooth_quant: dst.value = np.ascontiguousarray(np.transpose(t, [1, 0])) - dense_scale = getattr(tensorrt_llm_gpt.layers[i].attention, + dense_scale = getattr(gpt_layer.attention, "quantization_scaling_factor", None) set_smoothquant_scale_factors( - tensorrt_llm_gpt.layers[i].attention.dense, dense_scale, - dir_path, 'model.layers.' + str(i) + '.attention.dense.', - [1, n_embd], quant_per_token_dyn, quant_per_channel) + gpt_layer.attention.dense, dense_scale, dir_path, + 'model.layers.' + str(i) + '.attention.dense.', [1, n_embd], + quant_per_token_dyn, quant_per_channel) # change it to the real smoother if dense layer is applied smooth quant - tensorrt_llm_gpt.layers[i].attention.dense.smoother.value = np.ones( + gpt_layer.attention.dense.smoother.value = np.ones( [1, n_embd // tensor_parallel], dtype=np.float32) elif use_weight_only: processed_torch_weights, torch_weight_scales = torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( @@ -315,7 +316,7 @@ def set_smoothquant_scale_factors(module, dst.value = np.ascontiguousarray(np.transpose(t, [1, 0])) if bias: - dst = tensorrt_llm_gpt.layers[i].attention.dense.bias + dst = gpt_layer.attention.dense.bias dst.value = fromfile( dir_path, 'model.layers.' + str(i) + '.attention.dense.bias.bin') @@ -327,12 +328,12 @@ def set_smoothquant_scale_factors(module, i].attention.dense.weights_scaling_factor.value = np.array( [scaling_factors['dense_weights'][i]], dtype=fake_fp8_sf_dt) - dst = tensorrt_llm_gpt.layers[i].post_layernorm.weight + dst = gpt_layer.post_layernorm.weight dst.value = fromfile( dir_path, 'model.layers.' + str(i) + '.post_attention_layernorm.weight.bin') - dst = tensorrt_llm_gpt.layers[i].post_layernorm.bias + dst = gpt_layer.post_layernorm.bias dst.value = fromfile( dir_path, 'model.layers.' + str(i) + '.post_attention_layernorm.bias.bin') @@ -344,28 +345,28 @@ def set_smoothquant_scale_factors(module, tensorrt_llm_gpt.layers[ i].mlp.fc.weight.value = np.ascontiguousarray( np.transpose(t, [1, 0])) - set_smoothquant_scale_factors( - tensorrt_llm_gpt.layers[i].mlp.fc, - tensorrt_llm_gpt.layers[i].post_layernorm.scale_to_int, - dir_path, - 'model.layers.' + str(i) + '.mlp.dense_h_to_4h.', - [1, inter_size // tensor_parallel], - quant_per_token_dyn, - quant_per_channel, - rank=rank) + set_smoothquant_scale_factors(gpt_layer.mlp.fc, + gpt_layer.post_layernorm.scale_to_int, + dir_path, + 'model.layers.' + str(i) + + '.mlp.dense_h_to_4h.', + [1, inter_size // tensor_parallel], + quant_per_token_dyn, + quant_per_channel, + rank=rank) elif use_weight_only: - dst = tensorrt_llm_gpt.layers[i].mlp.fc.weight + dst = gpt_layer.mlp.fc.weight processed_torch_weights, torch_weight_scales = torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( numpy_to_torch(t), plugin_weight_only_quant_type) dst.value = torch_to_numpy(processed_torch_weights) - scales = tensorrt_llm_gpt.layers[i].mlp.fc.per_channel_scale + scales = gpt_layer.mlp.fc.per_channel_scale scales.value = torch_to_numpy(torch_weight_scales) else: tensorrt_llm_gpt.layers[ i].mlp.fc.weight.value = np.ascontiguousarray( np.transpose(t, [1, 0])) if bias: - tensorrt_llm_gpt.layers[i].mlp.fc.bias.value = fromfile( + gpt_layer.mlp.fc.bias.value = fromfile( dir_path, 'model.layers.' + str(i) + '.mlp.dense_h_to_4h.bias.' + str(rank) + '.bin') if is_gated_activation(hidden_act): @@ -392,27 +393,27 @@ def set_smoothquant_scale_factors(module, tensorrt_llm_gpt.layers[ i].mlp.proj.weight.value = np.ascontiguousarray( np.transpose(t, [1, 0])) - proj_scale = getattr(tensorrt_llm_gpt.layers[i].mlp, - "quantization_scaling_factor", None) + proj_scale = getattr(gpt_layer.mlp, "quantization_scaling_factor", + None) set_smoothquant_scale_factors( - tensorrt_llm_gpt.layers[i].mlp.proj, proj_scale, dir_path, + gpt_layer.mlp.proj, proj_scale, dir_path, 'model.layers.' + str(i) + '.mlp.dense_4h_to_h.', [1, n_embd], quant_per_token_dyn, quant_per_channel) # change it to the real smoother if proj layer is applied smooth quant - tensorrt_llm_gpt.layers[i].mlp.proj.smoother.value = np.ones( + gpt_layer.mlp.proj.smoother.value = np.ones( [1, inter_size // tensor_parallel], dtype=np.float32) elif use_weight_only: - dst = tensorrt_llm_gpt.layers[i].mlp.proj.weight + dst = gpt_layer.mlp.proj.weight processed_torch_weights, torch_weight_scales = torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( numpy_to_torch(t), plugin_weight_only_quant_type) dst.value = torch_to_numpy(processed_torch_weights) - scales = tensorrt_llm_gpt.layers[i].mlp.proj.per_channel_scale + scales = gpt_layer.mlp.proj.per_channel_scale scales.value = torch_to_numpy(torch_weight_scales) else: - tensorrt_llm_gpt.layers[i].mlp.proj.weight.value = ( - np.ascontiguousarray(np.transpose(t, [1, 0]))) + gpt_layer.mlp.proj.weight.value = (np.ascontiguousarray( + np.transpose(t, [1, 0]))) if bias: - tensorrt_llm_gpt.layers[i].mlp.proj.bias.value = fromfile( + gpt_layer.mlp.proj.bias.value = fromfile( dir_path, 'model.layers.' + str(i) + '.mlp.dense_4h_to_h.bias.bin') @@ -423,7 +424,7 @@ def set_smoothquant_scale_factors(module, np.float32) tensorrt_llm_gpt.layers[ i].attention.kv_orig_quant_scale.value = 1.0 / t - tensorrt_llm_gpt.layers[i].attention.kv_quant_orig_scale.value = t + gpt_layer.attention.kv_quant_orig_scale.value = t if enable_fp8_qdq: tensorrt_llm_gpt.layers[ diff --git a/examples/gptj/README.md b/examples/gptj/README.md index 1eee8d72ab6b..be19559a28d1 100644 --- a/examples/gptj/README.md +++ b/examples/gptj/README.md @@ -5,11 +5,11 @@ This document explains how to build the [GPT-J](https://huggingface.co/EleutherA ## Overview The TensorRT-LLM GPT-J implementation can be found in [`tensorrt_llm/models/gptj/model.py`](../../tensorrt_llm/models/gptj/model.py). The TensorRT-LLM GPT-J example -code is located in [`examples/gptj`](./). There are three main files in that folder: +code is located in [`examples/gptj`](./). There are three main files: * [`build.py`](./build.py) to build the [TensorRT](https://developer.nvidia.com/tensorrt) engine(s) needed to run the GPT-J model, * [`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/cnn_dailymail) dataset using the model. + * and a shared [`../summarize.py`](../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/cnn_dailymail) dataset using the model. ## Support Matrix * FP16 @@ -235,19 +235,18 @@ As previously explained, the first step is to build the TensorRT engine as descr pip install -r requirements.txt ``` -The summarization can be done using the [`summarize.py`](./summarize.py) script as follows: +The summarization can be done using the [`../summarize.py`](../summarize.py) script as follows: ```bash # Run the summarization task. -python3 summarize.py --engine_dir gptj_engine \ - --model_dir gptj_model \ - --test_hf \ - --batch_size 1 \ - --test_trt_llm \ - --tensorrt_llm_rouge1_threshold 14 \ - --data_type fp16 \ - --check_accuracy - +python3 ../summarize.py --engine_dir gptj_engine \ + --hf_model_dir gptj_model \ + --test_hf \ + --batch_size 1 \ + --test_trt_llm \ + --tensorrt_llm_rouge1_threshold 14 \ + --data_type fp16 \ + --check_accuracy ``` ## Known issues diff --git a/examples/gptj/build.py b/examples/gptj/build.py index 61e701270079..fd3747d97a82 100644 --- a/examples/gptj/build.py +++ b/examples/gptj/build.py @@ -462,7 +462,7 @@ def build(rank, args): if args.parallel_build and cur_rank != rank: continue # NOTE(nkorobov): when only int8 kv cache is used together with paged kv cache no int8 tensors are exposed to TRT - int8_trt_flag = args.quant_mode.has_act_and_weight_quant() or ( + int8_trt_flag = args.quant_mode.has_act_or_weight_quant() or ( not args.paged_kv_cache and args.quant_mode.has_int8_kv_cache()) builder_config = builder.create_builder_config( @@ -478,10 +478,10 @@ def build(rank, args): hidden_act=args.hidden_act, max_position_embeddings=args.n_positions, max_batch_size=args.max_batch_size, + max_beam_width=args.max_beam_width, max_input_len=args.max_input_len, max_output_len=args.max_output_len, max_num_tokens=args.max_num_tokens, - fp8=args.enable_fp8, int8=int8_trt_flag, quant_mode=args.quant_mode, strongly_typed=args.strongly_typed) diff --git a/examples/gptj/summarize.py b/examples/gptj/summarize.py deleted file mode 100644 index b490d3bb73bb..000000000000 --- a/examples/gptj/summarize.py +++ /dev/null @@ -1,416 +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 copy -import json -import os -import random - -import numpy as np -import torch -from datasets import load_dataset, load_metric -from transformers import AutoModelForCausalLM, AutoTokenizer - -import tensorrt_llm -import tensorrt_llm.profiler as profiler -from tensorrt_llm.logger import logger -from tensorrt_llm.quantization import QuantMode - -from build import get_engine_name # isort:skip - - -def TRTGPTJ(args, config): - dtype = config['builder_config']['precision'] - world_size = config['builder_config']['tensor_parallel'] - assert world_size == tensorrt_llm.mpi_world_size(), \ - f'Engine world size ({world_size}) != Runtime world size ({tensorrt_llm.mpi_world_size()})' - - world_size = config['builder_config']['tensor_parallel'] - num_heads = config['builder_config']['num_heads'] // world_size - hidden_size = config['builder_config']['hidden_size'] // world_size - vocab_size = config['builder_config']['vocab_size'] - num_layers = config['builder_config']['num_layers'] - use_gpt_attention_plugin = bool( - config['plugin_config']['gpt_attention_plugin']) - remove_input_padding = config['plugin_config']['remove_input_padding'] - quant_mode = QuantMode(config['builder_config'].get('quant_mode', 0)) - paged_kv_cache = config['plugin_config']['paged_kv_cache'] - tokens_per_block = config['plugin_config']['tokens_per_block'] - - model_config = tensorrt_llm.runtime.ModelConfig( - vocab_size=vocab_size, - num_layers=num_layers, - num_heads=num_heads, - num_kv_heads=num_heads, - hidden_size=hidden_size, - gpt_attention_plugin=use_gpt_attention_plugin, - remove_input_padding=remove_input_padding, - paged_kv_cache=paged_kv_cache, - tokens_per_block=tokens_per_block, - quant_mode=quant_mode, - dtype=dtype) - - runtime_rank = tensorrt_llm.mpi_rank() - runtime_mapping = tensorrt_llm.Mapping(world_size, - runtime_rank, - tp_size=world_size) - torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node) - - engine_name = get_engine_name('gptj', dtype, world_size, runtime_rank) - serialize_path = os.path.join(args.engine_dir, engine_name) - - tensorrt_llm.logger.set_level(args.log_level) - - with open(serialize_path, 'rb') as f: - engine_buffer = f.read() - decoder = tensorrt_llm.runtime.GenerationSession(model_config, - engine_buffer, - runtime_mapping) - - return decoder - - -def main(args): - 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_dir = args.model_dir - - tokenizer = AutoTokenizer.from_pretrained(model_dir, - padding_side='left', - model_max_length=2048, - truncation=True) - tokenizer.pad_token = tokenizer.eos_token - - dataset_cnn = load_dataset("ccdv/cnn_dailymail", - '3.0.0', - cache_dir=args.dataset_path) - - config_path = os.path.join(args.engine_dir, 'config.json') - with open(config_path, 'r') as f: - config = json.load(f) - - max_batch_size = args.batch_size - - # runtime parameters - # repetition_penalty = 1 - top_k = args.top_k - output_len = args.output_len - test_token_num = 923 - # top_p = 0.0 - # random_seed = 5 - temperature = 1 - num_beams = args.num_beams - - pad_id = tokenizer.encode(tokenizer.pad_token, add_special_tokens=False)[0] - end_id = tokenizer.encode(tokenizer.eos_token, add_special_tokens=False)[0] - - if test_trt_llm: - tensorrt_llm_gpt = TRTGPTJ(args, config) - - if test_hf: - model = AutoModelForCausalLM.from_pretrained(model_dir) - model.cuda() - if args.data_type == 'fp16': - model.half() - - def summarize_tensorrt_llm(datapoint): - batch_size = len(datapoint['article']) - - line = copy.copy(datapoint['article']) - line_encoded = [] - input_lengths = [] - for i in range(batch_size): - line[i] = line[i] + ' TL;DR: ' - - line[i] = line[i].strip() - line[i] = line[i].replace(" n't", "n't") - - input_id = tokenizer.encode(line[i], - return_tensors='pt').type(torch.int32) - input_id = input_id[:, -test_token_num:] - - line_encoded.append(input_id) - input_lengths.append(input_id.shape[-1]) - - # do padding, should move outside the profiling to prevent the overhead - max_length = max(input_lengths) - if tensorrt_llm_gpt.remove_input_padding: - line_encoded = [ - torch.tensor(t, dtype=torch.int32).cuda() for t in line_encoded - ] - else: - # do padding, should move outside the profiling to prevent the overhead - for i in range(batch_size): - pad_size = max_length - input_lengths[i] - - pad = torch.ones([1, pad_size]).type(torch.int32) * pad_id - line_encoded[i] = torch.cat( - [torch.tensor(line_encoded[i], dtype=torch.int32), pad], - axis=-1) - - line_encoded = torch.cat(line_encoded, axis=0).cuda() - input_lengths = torch.tensor(input_lengths, - dtype=torch.int32).cuda() - - sampling_config = tensorrt_llm.runtime.SamplingConfig( - end_id=end_id, pad_id=pad_id, top_k=top_k, num_beams=num_beams) - - with torch.no_grad(): - tensorrt_llm_gpt.setup(batch_size, - max_context_length=max_length, - max_new_tokens=output_len, - beam_width=num_beams, - max_kv_cache_length=args.max_kv_cache_len) - - if tensorrt_llm_gpt.remove_input_padding: - output_ids = tensorrt_llm_gpt.decode_batch( - line_encoded, sampling_config) - else: - output_ids = tensorrt_llm_gpt.decode( - line_encoded, - input_lengths, - sampling_config, - ) - - torch.cuda.synchronize() - - # Extract a list of tensors of shape beam_width x output_ids. - output_beams_list, output_ids_list = [], [] - if tensorrt_llm_gpt.mapping.is_first_pp_rank(): - output_beams_list = [ - tokenizer.batch_decode(output_ids[batch_idx, :, - input_lengths[batch_idx]:], - skip_special_tokens=True) - for batch_idx in range(batch_size) - ] - output_ids_list = [ - output_ids[batch_idx, :, input_lengths[batch_idx]:] - for batch_idx in range(batch_size) - ] - return output_beams_list, output_ids_list - - def summarize_hf(datapoint): - batch_size = len(datapoint['article']) - 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}" - ) - - line = copy.copy(datapoint['article']) - for i in range(batch_size): - line[i] = line[i] + ' TL;DR: ' - - line[i] = line[i].strip() - line[i] = line[i].replace(" n't", "n't") - - line_encoded = tokenizer(line, - return_tensors='pt', - padding=True, - truncation=True)["input_ids"].type(torch.int64) - - line_encoded = line_encoded[:, -test_token_num:] - line_encoded = line_encoded.cuda() - - with torch.no_grad(): - output = model.generate(line_encoded, - max_length=len(line_encoded[0]) + - output_len, - top_k=top_k, - temperature=temperature, - eos_token_id=tokenizer.eos_token_id, - pad_token_id=tokenizer.pad_token_id, - num_beams=num_beams, - num_return_sequences=num_beams, - early_stopping=True) - - tokens_list = output[:, len(line_encoded[0]):].tolist() - output = output.reshape([batch_size, num_beams, -1]) - output_lines_list = [ - tokenizer.batch_decode(output[:, i, len(line_encoded[0]):], - skip_special_tokens=True) - for i in range(num_beams) - ] - - return output_lines_list, tokens_list - - if test_trt_llm: - datapoint = dataset_cnn['test'][0:1] - summary, _ = summarize_tensorrt_llm(datapoint) - if runtime_rank == 0: - logger.info( - "---------------------------------------------------------") - logger.info("TensorRT-LLM Generated : ") - logger.info(f" Article : {datapoint['article']}") - logger.info(f"\n Highlights : {datapoint['highlights']}") - logger.info(f"\n Summary : {summary}") - logger.info( - "---------------------------------------------------------") - - if test_hf: - datapoint = dataset_cnn['test'][0:1] - summary, _ = summarize_hf(datapoint) - logger.info("---------------------------------------------------------") - logger.info("HF Generated : ") - logger.info(f" Article : {datapoint['article']}") - logger.info(f"\n Highlights : {datapoint['highlights']}") - logger.info(f"\n Summary : {summary}") - logger.info("---------------------------------------------------------") - - tensorrt_llm_result = [[] for _ in range(num_beams)] - hf_result = [[] for _ in range(num_beams)] - ite_count = 0 - data_point_idx = 0 - - # Support running the set with different order to verify correctness - test_idx = list( - range(min(len(dataset_cnn['test']), max_batch_size * args.max_ite))) - random.seed(args.random_seed) - random.shuffle(test_idx) - while (data_point_idx < len(dataset_cnn['test'])) 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_cnn['test'][test_idx[data_point_idx:( - data_point_idx + max_batch_size)]] - - if test_trt_llm: - profiler.start('tensorrt_llm') - summary_tensorrt_llm, tokens_tensorrt_llm = summarize_tensorrt_llm( - datapoint) - profiler.stop('tensorrt_llm') - - if test_hf: - profiler.start('hf') - summary_hf, tokens_hf = summarize_hf(datapoint) - profiler.stop('hf') - - if runtime_rank == 0: - if test_trt_llm: - for batch_idx in range(len(summary_tensorrt_llm)): - for beam_idx in range(num_beams): - tensorrt_llm_result[beam_idx].append( - tuple([ - datapoint['id'][batch_idx], - summary_tensorrt_llm[batch_idx][beam_idx], - datapoint['highlights'][batch_idx] - ])) - if test_hf: - for beam_idx in range(num_beams): - for batch_idx in range(len(summary_hf[beam_idx])): - hf_result[beam_idx].append( - tuple([ - datapoint['id'][batch_idx], - summary_hf[beam_idx][batch_idx], - datapoint['highlights'][batch_idx] - ])) - - logger.debug('-' * 100) - logger.debug(f"Article : {datapoint['article']}") - if test_trt_llm: - logger.debug(f'TensorRT-LLM Summary: {summary_tensorrt_llm}') - if test_hf: - logger.debug(f'HF Summary: {summary_hf}') - logger.debug(f"highlights : {datapoint['highlights']}") - - data_point_idx += max_batch_size - ite_count += 1 - - if runtime_rank == 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)' - ) - for beam_idx in range(num_beams): - # Because 'rouge' uses sampling to compute the scores, the scores - # would be different when the results are same with different order. - # So, sorting them first to prevent this issue. - metric_tensorrt_llm = load_metric("rouge") - metric_tensorrt_llm.seed = 0 - beams_results = sorted(tensorrt_llm_result[beam_idx]) - - for j in range(len(beams_results)): - metric_tensorrt_llm.add_batch( - predictions=[beams_results[j][1]], - references=[beams_results[j][2]]) - - logger.info(f"TensorRT-LLM beam {beam_idx} result") - computed_metrics_tensorrt_llm = metric_tensorrt_llm.compute() - for key in computed_metrics_tensorrt_llm.keys(): - logger.info( - f' {key} : {computed_metrics_tensorrt_llm[key].mid[2]*100}' - ) - - if args.check_accuracy and beam_idx == 0: - assert computed_metrics_tensorrt_llm['rouge1'].mid[ - 2] * 100 > args.tensorrt_llm_rouge1_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)' - ) - for beam_idx in range(num_beams): - metric_tensorrt_hf = load_metric("rouge") - metric_tensorrt_hf.seed = 0 - beams_results = sorted(hf_result[beam_idx]) - - for j in range(len(beams_results)): - metric_tensorrt_hf.add_batch( - predictions=[beams_results[j][1]], - references=[beams_results[j][2]]) - logger.info(f"HF beam {beam_idx} result") - computed_metrics_hf = metric_tensorrt_hf.compute() - for key in computed_metrics_hf.keys(): - logger.info( - f' {key} : {computed_metrics_hf[key].mid[2]*100}') - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default='EleutherAI/gpt-j-6B') - parser.add_argument('--test_hf', action='store_true') - parser.add_argument('--test_trt_llm', action='store_true') - parser.add_argument('--data_type', - type=str, - choices=['fp32', 'fp16'], - default='fp32') - parser.add_argument('--dataset_path', type=str, default='') - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument('--engine_dir', type=str, default='gptj_engine') - 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_kv_cache_len', - type=int, - default=None, - help='The max kv cache length. \ - If the final sequence length exceeds the kv cache length, we will enable cyclic kv cache. \ - If it is set to None, we will use the max sequence length.') - parser.add_argument('--check_accuracy', action='store_true') - parser.add_argument('--tensorrt_llm_rouge1_threshold', - type=float, - default=15.0) - parser.add_argument('--num_beams', type=int, default=1) - parser.add_argument('--top_k', type=int, default=1) - parser.add_argument('--random_seed', type=int, default=0) - - args = parser.parse_args() - - main(args) diff --git a/examples/gptneox/README.md b/examples/gptneox/README.md index 6dc9c2b59a43..82c62704d8b3 100644 --- a/examples/gptneox/README.md +++ b/examples/gptneox/README.md @@ -4,11 +4,11 @@ This document explains how to build the [GPT-NeoX](https://huggingface.co/Eleuth ## Overview -The TensorRT-LLM GPT-NeoX implementation can be found in [`tensorrt_llm/models/gptneox/model.py`](../../tensorrt_llm/models/gptneox/model.py). The TensorRT-LLM GPT-NeoX example code is located in [`examples/gptneox`](./). There are three main files in that folder: +The TensorRT-LLM GPT-NeoX implementation can be found in [`tensorrt_llm/models/gptneox/model.py`](../../tensorrt_llm/models/gptneox/model.py). The TensorRT-LLM GPT-NeoX example code is located in [`examples/gptneox`](./). There are three main files: * [`build.py`](./build.py) to build the [TensorRT](https://developer.nvidia.com/tensorrt) engine(s) needed to run the GPT-NeoX model, * [`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/cnn_dailymail) dataset using the model. + * and a shared [`../summarize.py`](../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/cnn_dailymail) dataset using the model. ## Support Matrix * FP16 @@ -101,36 +101,36 @@ As previously explained, the first step is to build the TensorRT engine as descr pip install -r requirements.txt ``` -The summarization can be done using the [`summarize.py`](./summarize.py) script as follows: +The summarization can be done using the [`../summarize.py`](../summarize.py) script as follows: ```bash # Run the summarization task using a TensorRT-LLM model and a single GPU. -python3 summarize.py --engine_dir gptneox_engine \ - --model_dir gptneox_model \ - --batch_size 1 \ - --test_trt_llm \ - --tensorrt_llm_rouge1_threshold 14 \ - --data_type fp16 \ - --check_accuracy 2>&1 | tee summary_trt_llm.log +python3 ../summarize.py --engine_dir gptneox_engine \ + --hf_model_dir gptneox_model \ + --batch_size 1 \ + --test_trt_llm \ + --tensorrt_llm_rouge1_threshold 14 \ + --data_type fp16 \ + --check_accuracy 2>&1 | tee summary_trt_llm.log # Run the summarization task using a HF model and a single GPU. -python3 summarize.py --engine_dir gptneox_engine \ - --model_dir gptneox_model \ - --batch_size 1 \ - --test_hf \ - --tensorrt_llm_rouge1_threshold 14 \ - --data_type fp16 \ - --check_accuracy 2>&1 | tee summary_hf.log +python3 ../summarize.py --engine_dir gptneox_engine \ + --hf_model_dir gptneox_model \ + --batch_size 1 \ + --test_hf \ + --tensorrt_llm_rouge1_threshold 14 \ + --data_type fp16 \ + --check_accuracy 2>&1 | tee summary_hf.log # Run the summarization task using a TensorRT-LLM model and 2-way tensor parallelism. mpirun -n 2 --allow-run-as-root \ -python3 summarize.py --engine_dir gptneox_engine_tp2 \ - --model_dir gptneox_model \ - --batch_size 1 \ - --test_trt_llm \ - --tensorrt_llm_rouge1_threshold 14 \ - --data_type fp16 \ - --check_accuracy 2>&1 | tee summary_trt_llm_tp2.log +python3 ../summarize.py --engine_dir gptneox_engine_tp2 \ + --hf_model_dir gptneox_model \ + --batch_size 1 \ + --test_trt_llm \ + --tensorrt_llm_rouge1_threshold 14 \ + --data_type fp16 \ + --check_accuracy 2>&1 | tee summary_trt_llm_tp2.log ``` ## Apply groupwise quantization GPTQ @@ -205,25 +205,25 @@ Install the requirements first. pip install -r requirements.txt ``` -Then use the [`summarize.py`](./summarize.py) script to summarize. +Then use the [`../summarize.py`](../summarize.py) script to summarize. ```bash # Run the summarization task using a TensorRT-LLM model and a single GPU. -python3 summarize.py --engine_dir gptneox_engine_gptq \ - --model_dir gptneox_model \ - --batch_size 1 \ - --test_trt_llm \ - --tensorrt_llm_rouge1_threshold 14 \ - --data_type fp16 \ - --check_accuracy 2>&1 | tee summary_trt_llm_gptq.log +python3 ../summarize.py --engine_dir gptneox_engine_gptq \ + --hf_model_dir gptneox_model \ + --batch_size 1 \ + --test_trt_llm \ + --tensorrt_llm_rouge1_threshold 14 \ + --data_type fp16 \ + --check_accuracy 2>&1 | tee summary_trt_llm_gptq.log # Run the summarization task using a TensorRT-LLM model and 2-way tensor parallelism. mpirun -n 2 --allow-run-as-root \ -python3 summarize.py --engine_dir gptneox_engine_gptq_tp2 \ - --model_dir gptneox_model \ - --batch_size 1 \ - --test_trt_llm \ - --tensorrt_llm_rouge1_threshold 14 \ - --data_type fp16 \ - --check_accuracy 2>&1 | tee summary_trt_llm_gptq_tp2.log +python3 ../summarize.py --engine_dir gptneox_engine_gptq_tp2 \ + --hf_model_dir gptneox_model \ + --batch_size 1 \ + --test_trt_llm \ + --tensorrt_llm_rouge1_threshold 14 \ + --data_type fp16 \ + --check_accuracy 2>&1 | tee summary_trt_llm_gptq_tp2.log ``` diff --git a/examples/gptneox/build.py b/examples/gptneox/build.py index ed5b4ac44f8d..92a7e5d9ec75 100644 --- a/examples/gptneox/build.py +++ b/examples/gptneox/build.py @@ -400,6 +400,7 @@ def build(rank, args): max_position_embeddings=args.n_positions, apply_query_key_layer_scaling=apply_query_key_layer_scaling, max_batch_size=args.max_batch_size, + max_beam_width=args.max_beam_width, max_input_len=args.max_input_len, int8=args.use_weight_only_quant_matmul_plugin or args.use_weight_only_groupwise_quant_matmul_plugin, diff --git a/examples/gptneox/requirements.txt b/examples/gptneox/requirements.txt index f46bff310071..fa85d0a336b5 100644 --- a/examples/gptneox/requirements.txt +++ b/examples/gptneox/requirements.txt @@ -1,2 +1,3 @@ datasets~=2.14.5 rouge_score~=0.1.2 +evaluate~=0.4.1 diff --git a/examples/gptneox/summarize.py b/examples/gptneox/summarize.py deleted file mode 100644 index b9f6d0a45ffc..000000000000 --- a/examples/gptneox/summarize.py +++ /dev/null @@ -1,380 +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 copy -import json -import os - -import numpy as np -import torch -from datasets import load_dataset, load_metric -from transformers import AutoModelForCausalLM, AutoTokenizer - -import tensorrt_llm -import tensorrt_llm.profiler as profiler -from tensorrt_llm.logger import logger - -from build import get_engine_name # isort:skip - - -def TRTGPTNeoX(args, config): - dtype = config['builder_config']['precision'] - world_size = config['builder_config']['tensor_parallel'] - assert world_size == tensorrt_llm.mpi_world_size(), \ - f'Engine world size ({world_size}) != Runtime world size ({tensorrt_llm.mpi_world_size()})' - - world_size = config['builder_config']['tensor_parallel'] - num_heads = config['builder_config']['num_heads'] // world_size - hidden_size = config['builder_config']['hidden_size'] // world_size - vocab_size = config['builder_config']['vocab_size'] - num_layers = config['builder_config']['num_layers'] - use_gpt_attention_plugin = bool( - config['plugin_config']['gpt_attention_plugin']) - remove_input_padding = config['plugin_config']['remove_input_padding'] - - model_config = tensorrt_llm.runtime.ModelConfig( - vocab_size=vocab_size, - num_layers=num_layers, - num_heads=num_heads, - num_kv_heads=num_heads, - hidden_size=hidden_size, - gpt_attention_plugin=use_gpt_attention_plugin, - remove_input_padding=remove_input_padding, - dtype=dtype) - - runtime_rank = tensorrt_llm.mpi_rank() - runtime_mapping = tensorrt_llm.Mapping(world_size, - runtime_rank, - tp_size=world_size) - torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node) - - engine_name = get_engine_name('gptneox', dtype, world_size, runtime_rank) - serialize_path = os.path.join(args.engine_dir, engine_name) - - tensorrt_llm.logger.set_level(args.log_level) - - with open(serialize_path, 'rb') as f: - engine_buffer = f.read() - decoder = tensorrt_llm.runtime.GenerationSession(model_config, - engine_buffer, - runtime_mapping) - - return decoder - - -def main(args): - 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_dir = args.model_dir - - tokenizer = AutoTokenizer.from_pretrained(model_dir, - padding_side='left', - model_max_length=2048, - truncation=True) - tokenizer.pad_token = tokenizer.eos_token - - dataset_cnn = load_dataset("ccdv/cnn_dailymail", - '3.0.0', - cache_dir=args.dataset_path) - - config_path = os.path.join(args.engine_dir, 'config.json') - with open(config_path, 'r') as f: - config = json.load(f) - - max_batch_size = args.batch_size - - # runtime parameters - # repetition_penalty = 1 - top_k = args.top_k - output_len = args.output_len - test_token_num = 923 - # top_p = 0.0 - # random_seed = 5 - temperature = 1 - num_beams = args.num_beams - - pad_id = tokenizer.encode(tokenizer.pad_token, add_special_tokens=False)[0] - end_id = tokenizer.encode(tokenizer.eos_token, add_special_tokens=False)[0] - - if test_trt_llm: - tensorrt_llm_gpt = TRTGPTNeoX(args, config) - - if test_hf: - model = AutoModelForCausalLM.from_pretrained(model_dir) - model.cuda() - if args.data_type == 'fp16': - model.half() - - def summarize_tensorrt_llm(datapoint): - batch_size = len(datapoint['article']) - - line = copy.copy(datapoint['article']) - line_encoded = [] - input_lengths = [] - for i in range(batch_size): - line[i] = line[i] + ' TL;DR: ' - - line[i] = line[i].strip() - line[i] = line[i].replace(" n't", "n't") - - input_id = tokenizer.encode(line[i], - return_tensors='pt').type(torch.int32) - input_id = input_id[:, -test_token_num:] - - line_encoded.append(input_id) - input_lengths.append(input_id.shape[-1]) - - # do padding, should move outside the profiling to prevent the overhead - max_length = max(input_lengths) - if tensorrt_llm_gpt.remove_input_padding: - line_encoded = [ - torch.tensor(t, dtype=torch.int32).cuda() for t in line_encoded - ] - else: - # do padding, should move outside the profiling to prevent the overhead - for i in range(batch_size): - pad_size = max_length - input_lengths[i] - - pad = torch.ones([1, pad_size]).type(torch.int32) * pad_id - line_encoded[i] = torch.cat( - [torch.tensor(line_encoded[i], dtype=torch.int32), pad], - axis=-1) - - line_encoded = torch.cat(line_encoded, axis=0).cuda() - input_lengths = torch.tensor(input_lengths, - dtype=torch.int32).cuda() - - sampling_config = tensorrt_llm.runtime.SamplingConfig( - end_id=end_id, pad_id=pad_id, top_k=top_k, num_beams=num_beams) - - with torch.no_grad(): - tensorrt_llm_gpt.setup(batch_size, - max_context_length=max_length, - max_new_tokens=output_len, - beam_width=num_beams, - max_kv_cache_length=args.max_kv_cache_len) - - if tensorrt_llm_gpt.remove_input_padding: - output_ids = tensorrt_llm_gpt.decode_batch( - line_encoded, sampling_config) - else: - output_ids = tensorrt_llm_gpt.decode( - line_encoded, - input_lengths, - sampling_config, - ) - - torch.cuda.synchronize() - - # Extract a list of tensors of shape beam_width x output_ids. - if tensorrt_llm_gpt.mapping.is_first_pp_rank(): - output_beams_list = [ - tokenizer.batch_decode(output_ids[batch_idx, :, - input_lengths[batch_idx]:], - skip_special_tokens=True) - for batch_idx in range(batch_size) - ] - return output_beams_list, output_ids[:, :, max_length:].tolist() - return [], [] - - def summarize_hf(datapoint): - batch_size = len(datapoint['article']) - 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}" - ) - - line = copy.copy(datapoint['article']) - for i in range(batch_size): - line[i] = line[i] + ' TL;DR: ' - - line[i] = line[i].strip() - line[i] = line[i].replace(" n't", "n't") - - line_encoded = tokenizer(line, - return_tensors='pt', - padding=True, - truncation=True)["input_ids"].type(torch.int64) - - line_encoded = line_encoded[:, -test_token_num:] - line_encoded = line_encoded.cuda() - - with torch.no_grad(): - output = model.generate(line_encoded, - max_length=len(line_encoded[0]) + - output_len, - top_k=top_k, - temperature=temperature, - eos_token_id=tokenizer.eos_token_id, - pad_token_id=tokenizer.pad_token_id, - num_beams=num_beams, - num_return_sequences=num_beams, - early_stopping=True) - - tokens_list = output[:, len(line_encoded[0]):].tolist() - output = output.reshape([batch_size, num_beams, -1]) - output_lines_list = [ - tokenizer.batch_decode(output[:, i, len(line_encoded[0]):], - skip_special_tokens=True) - for i in range(num_beams) - ] - - return output_lines_list, tokens_list - - if test_trt_llm: - datapoint = dataset_cnn['test'][0:1] - summary, _ = summarize_tensorrt_llm(datapoint) - if runtime_rank == 0: - logger.info( - "---------------------------------------------------------") - logger.info("TensorRT-LLM Generated : ") - logger.info(f" Article : {datapoint['article']}") - logger.info(f"\n Highlights : {datapoint['highlights']}") - logger.info(f"\n Summary : {summary}") - logger.info( - "---------------------------------------------------------") - - if test_hf: - datapoint = dataset_cnn['test'][0:1] - summary, _ = summarize_hf(datapoint) - logger.info("---------------------------------------------------------") - logger.info("HF Generated : ") - logger.info(f" Article : {datapoint['article']}") - logger.info(f"\n Highlights : {datapoint['highlights']}") - logger.info(f"\n Summary : {summary}") - logger.info("---------------------------------------------------------") - - metric_tensorrt_llm = [load_metric("rouge") for _ in range(num_beams)] - metric_hf = [load_metric("rouge") for _ in range(num_beams)] - for i in range(num_beams): - metric_tensorrt_llm[i].seed = 0 - metric_hf[i].seed = 0 - - ite_count = 0 - data_point_idx = 0 - while (data_point_idx < len(dataset_cnn['test'])) 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_cnn['test'][data_point_idx:(data_point_idx + - max_batch_size)] - - if test_trt_llm: - profiler.start('tensorrt_llm') - summary_tensorrt_llm, tokens_tensorrt_llm = summarize_tensorrt_llm( - datapoint) - profiler.stop('tensorrt_llm') - - if test_hf: - profiler.start('hf') - summary_hf, tokens_hf = summarize_hf(datapoint) - profiler.stop('hf') - - if runtime_rank == 0: - if test_trt_llm: - for batch_idx in range(len(summary_tensorrt_llm)): - for beam_idx in range(num_beams): - metric_tensorrt_llm[beam_idx].add_batch( - predictions=[ - summary_tensorrt_llm[batch_idx][beam_idx] - ], - references=[datapoint['highlights'][batch_idx]]) - if test_hf: - for beam_idx in range(num_beams): - for i in range(len(summary_hf[beam_idx])): - metric_hf[beam_idx].add_batch( - predictions=[summary_hf[beam_idx][i]], - references=[datapoint['highlights'][i]]) - - logger.debug('-' * 100) - logger.debug(f"Article : {datapoint['article']}") - if test_trt_llm: - logger.debug(f'TensorRT-LLM Summary: {summary_tensorrt_llm}') - if test_hf: - logger.debug(f'HF Summary: {summary_hf}') - logger.debug(f"highlights : {datapoint['highlights']}") - - data_point_idx += max_batch_size - ite_count += 1 - - if runtime_rank == 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)' - ) - for beam_idx in range(num_beams): - logger.info(f"TensorRT-LLM beam {beam_idx} result") - 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].mid[2]*100}' - ) - - if args.check_accuracy and beam_idx == 0: - assert computed_metrics_tensorrt_llm['rouge1'].mid[ - 2] * 100 > args.tensorrt_llm_rouge1_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)' - ) - for beam_idx in range(num_beams): - logger.info(f"HF beam {beam_idx} result") - computed_metrics_hf = metric_hf[beam_idx].compute() - for key in computed_metrics_hf.keys(): - logger.info( - f' {key} : {computed_metrics_hf[key].mid[2]*100}') - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', - type=str, - default='EleutherAI/gpt-neox-20b') - parser.add_argument('--test_hf', action='store_true') - parser.add_argument('--test_trt_llm', action='store_true') - parser.add_argument('--data_type', - type=str, - choices=['fp32', 'fp16'], - default='fp32') - parser.add_argument('--dataset_path', type=str, default='') - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument('--engine_dir', type=str, default='gptneox_engine') - 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_kv_cache_len', - type=int, - default=None, - help='The max kv cache length. \ - If the final sequence length exceeds the kv cache length, we will enable cyclic kv cache. \ - If it is set to None, we will use the max sequence length.') - parser.add_argument('--check_accuracy', action='store_true') - parser.add_argument('--tensorrt_llm_rouge1_threshold', - type=float, - default=15.0) - parser.add_argument('--num_beams', type=int, default=1) - parser.add_argument('--top_k', type=int, default=1) - - args = parser.parse_args() - - main(args) diff --git a/examples/internlm/README.md b/examples/internlm/README.md index 2f7a3cee7b9d..a16628db59d1 100644 --- a/examples/internlm/README.md +++ b/examples/internlm/README.md @@ -4,11 +4,11 @@ This document shows how to build and run InternLM 7B / 20B models in TensorRT-LL ## Overview -The TensorRT-LLM InternLM implementation can be found in [tensorrt_llm/models/internlm/model.py](../../tensorrt_llm/models/internlm/model.py). The TensorRT-LLM InternLM example code is located in [`examples/internlm`](./). There are three main files in that folder:: +The TensorRT-LLM InternLM implementation can be found in [tensorrt_llm/models/internlm/model.py](../../tensorrt_llm/models/internlm/model.py). The TensorRT-LLM InternLM example code is located in [`examples/internlm`](./). There are three main files: * [`build.py`](./build.py) to build the [TensorRT](https://developer.nvidia.com/tensorrt) engine(s) needed to run the InternLM model, * [`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/cnn_dailymail) dataset using the model. + * and a shared [`../summarize.py`](../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/cnn_dailymail) dataset using the model. ## Support Matrix * FP16 / BF16 @@ -139,7 +139,7 @@ python build.py --ft_model_dir=./internlm-chat-20b/smooth_internlm/int8_kv_cache --use_weight_only ``` -Test with `run.py` or `summarize.py`: +Test with `run.py` or `../summarize.py`: ```bash python run.py --max_output_len=120 \ @@ -152,15 +152,15 @@ python run.py --max_output_len=120 \ --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_location ./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-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_location ./internlm-chat-20b \ - --data_type fp16 \ - --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-20b \ + --data_type fp16 \ + --engine_dir ./internlm-chat-20b/trt_engines/int8_kv_cache_weight_only/1-gpu ``` #### SmoothQuant @@ -214,7 +214,7 @@ python build.py --ft_model_dir=./internlm-chat-20b/smooth_internlm/sq0.5/1-gpu/ Note we use `--ft_model_dir` instead of `--model_dir` and `--meta_ckpt_dir` since SmoothQuant model needs INT8 weights and various scales from the binary files. -Test with `run.py` or `summarize.py`: +Test with `run.py` or `../summarize.py`: ```bash python run.py --max_output_len=120 \ @@ -227,15 +227,15 @@ python run.py --max_output_len=120 \ --tokenizer_dir ./internlm-chat-20b/ \ --engine_dir ./internlm-chat-20b/trt_engines/smoothquant/1-gpu -python summarize.py --test_trt_llm --test_hf \ - --hf_model_location ./internlm-chat-7b \ - --data_type fp16 \ - --engine_dir ./internlm-chat-7b/trt_engines/smoothquant/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/smoothquant/1-gpu -python summarize.py --test_trt_llm --test_hf \ - --hf_model_location ./internlm-chat-20b \ - --data_type fp16 \ - --engine_dir ./internlm-chat-20b/trt_engines/smoothquant/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/smoothquant/1-gpu ``` ### Run @@ -280,28 +280,28 @@ mpirun -n 4 --allow-run-as-root \ ```bash # Run summarization using the InternLM 7B model in FP16. -python summarize.py --test_trt_llm --test_hf \ - --hf_model_location ./internlm-chat-7b/ \ - --data_type fp16 \ - --engine_dir ./internlm-chat-7b/trt_engines/fp16/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/fp16/1-gpu/ # Run summarization using the InternLM 7B model quantized to INT8. -python summarize.py --test_trt_llm --test_hf \ - --hf_model_location ./internlm-chat-7b/ \ - --data_type fp16 \ - --engine_dir ./internlm-chat-7b/trt_engines/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/weight_only/1-gpu/ # 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_location ./internlm-chat-7b/ \ - --data_type fp16 \ - --engine_dir ./internlm-chat-7b/trt_engines/fp16/2-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/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_location ./internlm-chat-20b/ \ - --data_type bf16 \ - --engine_dir ./internlm-chat-20b/trt_engines/bf16/4-gpu/ + 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/internlm/build.py b/examples/internlm/build.py index 7d08b7f3614c..40dcadaedad8 100644 --- a/examples/internlm/build.py +++ b/examples/internlm/build.py @@ -494,16 +494,14 @@ def build_rank_engine(builder: Builder, assert args.n_layer % args.pp_size == 0, \ f"num_layers {args.n_layer} must be a multiple of pipeline parallelism size {args.pp_size}" - # Initialize Module - tensorrt_llm_internlm = tensorrt_llm.models.InternLMForCausalLM( + tensorrt_llm_internlm = tensorrt_llm.models.LLaMAForCausalLM( num_layers=args.n_layer, num_heads=args.n_head, num_kv_heads=args.n_kv_head, hidden_size=args.n_embd, vocab_size=args.vocab_size, hidden_act=args.hidden_act, - attn_bias=args.attn_bias, max_position_embeddings=args.n_positions, dtype=dtype, mlp_hidden_size=args.inter_size, @@ -513,6 +511,8 @@ def build_rank_engine(builder: Builder, rotary_scaling=args.rotary_scaling, use_parallel_embedding=args.use_parallel_embedding, embedding_sharding_dim=args.embedding_sharding_dim, + use_fused_mlp=False, + attn_bias=args.attn_bias, quant_mode=args.quant_mode, rms_norm_eps=args.rms_norm_eps) if args.use_smooth_quant: @@ -688,6 +688,7 @@ def build(rank, args): hidden_act=args.hidden_act, max_position_embeddings=args.n_positions, max_batch_size=args.max_batch_size, + max_beam_width=args.max_beam_width, max_input_len=args.max_input_len, max_output_len=args.max_output_len, max_num_tokens=args.max_num_tokens, diff --git a/examples/internlm/requirements.txt b/examples/internlm/requirements.txt index 926de5f08662..ed4f5ac05673 100644 --- a/examples/internlm/requirements.txt +++ b/examples/internlm/requirements.txt @@ -1,3 +1,4 @@ datasets==2.14.5 rouge_score~=0.1.2 sentencepiece~=0.1.99 +evaluate~=0.4.1 diff --git a/examples/internlm/summarize.py b/examples/internlm/summarize.py deleted file mode 100644 index 2199014aafdb..000000000000 --- a/examples/internlm/summarize.py +++ /dev/null @@ -1,414 +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 copy -import json -import os - -import numpy as np -import torch -from datasets import load_dataset, load_metric -from transformers import AutoModelForCausalLM, AutoTokenizer - -import tensorrt_llm -import tensorrt_llm.profiler as profiler -from tensorrt_llm.logger import logger -from tensorrt_llm.quantization import QuantMode - -from build import get_engine_name # isort:skip - - -def TRTInternLM(args, config): - dtype = config['builder_config']['precision'] - tp_size = config['builder_config']['tensor_parallel'] - pp_size = config['builder_config']['pipeline_parallel'] - 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['builder_config']['num_heads'] // tp_size - hidden_size = config['builder_config']['hidden_size'] // tp_size - vocab_size = config['builder_config']['vocab_size'] - num_layers = config['builder_config']['num_layers'] - use_gpt_attention_plugin = bool( - config['plugin_config']['gpt_attention_plugin']) - remove_input_padding = config['plugin_config']['remove_input_padding'] - num_kv_heads = config['builder_config'].get('num_kv_heads', num_heads) - paged_kv_cache = config['plugin_config']['paged_kv_cache'] - tokens_per_block = config['plugin_config']['tokens_per_block'] - use_custom_all_reduce = config['plugin_config'].get('use_custom_all_reduce', - False) - - quant_mode = QuantMode(config['builder_config']['quant_mode']) - if config['builder_config'].get('multi_query_mode', False): - tensorrt_llm.logger.warning( - "`multi_query_mode` config is deprecated. Please rebuild the engine." - ) - num_kv_heads = 1 - num_kv_heads = (num_kv_heads + tp_size - 1) // tp_size - - model_config = tensorrt_llm.runtime.ModelConfig( - vocab_size=vocab_size, - num_layers=num_layers, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - hidden_size=hidden_size, - paged_kv_cache=paged_kv_cache, - tokens_per_block=tokens_per_block, - gpt_attention_plugin=use_gpt_attention_plugin, - remove_input_padding=remove_input_padding, - use_custom_all_reduce=use_custom_all_reduce, - dtype=dtype, - quant_mode=quant_mode) - - runtime_rank = tensorrt_llm.mpi_rank() - runtime_mapping = tensorrt_llm.Mapping(world_size, - runtime_rank, - tp_size=tp_size, - pp_size=pp_size) - torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node) - - engine_name = get_engine_name('internlm', dtype, tp_size, pp_size, - runtime_rank) - serialize_path = os.path.join(args.engine_dir, engine_name) - - tensorrt_llm.logger.set_level(args.log_level) - - profiler.start('load tensorrt_llm engine') - with open(serialize_path, 'rb') as f: - engine_buffer = f.read() - decoder = tensorrt_llm.runtime.GenerationSession(model_config, - engine_buffer, - runtime_mapping) - profiler.stop('load tensorrt_llm engine') - tensorrt_llm.logger.info( - f'Load engine takes: {profiler.elapsed_time_in_sec("load tensorrt_llm engine")} sec' - ) - return decoder - - -def main(args): - 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 - hf_model_location = args.hf_model_location - profiler.start('load tokenizer') - tokenizer = AutoTokenizer.from_pretrained(hf_model_location, - legacy=False, - padding_side='left', - trust_remote_code=True) - profiler.stop('load tokenizer') - tensorrt_llm.logger.info( - f'Load tokenizer takes: {profiler.elapsed_time_in_sec("load tokenizer")} sec' - ) - tokenizer.pad_token = tokenizer.eos_token - - dataset_cnn = load_dataset("ccdv/cnn_dailymail", - '3.0.0', - cache_dir=args.dataset_path) - - max_batch_size = args.batch_size - - # runtime parameters - # repetition_penalty = 1 - top_k = args.top_k - output_len = 100 - test_token_num = 923 - # top_p = 0.0 - # random_seed = 5 - temperature = 1 - num_beams = args.num_beams - - pad_id = tokenizer.encode(tokenizer.pad_token, add_special_tokens=False)[0] - end_id = tokenizer.encode(tokenizer.eos_token, add_special_tokens=False)[0] - - if test_trt_llm: - config_path = os.path.join(args.engine_dir, 'config.json') - with open(config_path, 'r') as f: - config = json.load(f) - - tensorrt_llm_internlm = TRTInternLM(args, config) - - if test_hf: - profiler.start('load HF model') - model = AutoModelForCausalLM.from_pretrained(hf_model_location, - trust_remote_code=True) - profiler.stop('load HF model') - tensorrt_llm.logger.info( - f'Load HF model takes: {profiler.elapsed_time_in_sec("load HF model")} sec' - ) - if args.data_type == 'fp16': - model.half() - elif args.data_type == 'fp32': - model = model.float() - elif args.data_type == 'bf16': - model = model.to(dtype=torch.bfloat16) - # else use dtype in hf config, which is by default - model.cuda() - - def summarize_tensorrt_llm(datapoint): - batch_size = len(datapoint['article']) - - line = copy.copy(datapoint['article']) - line_encoded = [] - input_lengths = [] - for i in range(batch_size): - line[i] = line[i] + ' TL;DR: ' - - line[i] = line[i].strip() - line[i] = line[i].replace(" n't", "n't") - - input_id = tokenizer.encode(line[i], - return_tensors='pt').type(torch.int32) - input_id = input_id[:, -test_token_num:] - - line_encoded.append(input_id) - input_lengths.append(input_id.shape[-1]) - - # do padding, should move outside the profiling to prevent the overhead - max_length = max(input_lengths) - if tensorrt_llm_internlm.remove_input_padding: - line_encoded = [ - torch.as_tensor(t, dtype=torch.int32, device='cuda') - for t in line_encoded - ] - else: - # do padding, should move outside the profiling to prevent the overhead - for i in range(batch_size): - pad_size = max_length - input_lengths[i] - - pad = torch.ones([1, pad_size]).type(torch.int32) * pad_id - line_encoded[i] = torch.cat( - [torch.tensor(line_encoded[i], dtype=torch.int32), pad], - axis=-1) - - line_encoded = torch.cat(line_encoded, axis=0).cuda() - input_lengths = torch.tensor(input_lengths, - dtype=torch.int32).cuda() - - sampling_config = tensorrt_llm.runtime.SamplingConfig( - end_id=end_id, pad_id=pad_id, top_k=top_k, num_beams=num_beams) - - with torch.no_grad(): - tensorrt_llm_internlm.setup(batch_size, - max_context_length=max_length, - max_new_tokens=output_len, - beam_width=num_beams) - - if tensorrt_llm_internlm.remove_input_padding: - output_ids = tensorrt_llm_internlm.decode_batch( - line_encoded, sampling_config) - else: - output_ids = tensorrt_llm_internlm.decode( - line_encoded, - input_lengths, - sampling_config, - ) - - torch.cuda.synchronize() - - # Extract a list of tensors of shape beam_width x output_ids. - if tensorrt_llm_internlm.mapping.is_first_pp_rank(): - output_beams_list = [ - tokenizer.batch_decode(output_ids[batch_idx, :, - input_lengths[batch_idx]:], - skip_special_tokens=True) - for batch_idx in range(batch_size) - ] - return output_beams_list, output_ids[:, :, max_length:].tolist() - return [], [] - - def summarize_hf(datapoint): - batch_size = len(datapoint['article']) - 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}" - ) - - line = copy.copy(datapoint['article']) - for i in range(batch_size): - line[i] = line[i] + ' TL;DR: ' - - line[i] = line[i].strip() - line[i] = line[i].replace(" n't", "n't") - - line_encoded = tokenizer(line, - return_tensors='pt', - padding=True, - truncation=True)["input_ids"].type(torch.int64) - - line_encoded = line_encoded[:, -test_token_num:] - line_encoded = line_encoded.cuda() - - with torch.no_grad(): - output = model.generate(line_encoded, - max_length=len(line_encoded[0]) + - output_len, - top_k=top_k, - temperature=temperature, - eos_token_id=tokenizer.eos_token_id, - pad_token_id=tokenizer.pad_token_id, - num_beams=num_beams, - num_return_sequences=num_beams, - early_stopping=True) - - tokens_list = output[:, len(line_encoded[0]):].tolist() - output = output.reshape([batch_size, num_beams, -1]) - output_lines_list = [ - tokenizer.batch_decode(output[:, i, len(line_encoded[0]):], - skip_special_tokens=True) - for i in range(num_beams) - ] - - return output_lines_list, tokens_list - - if test_trt_llm: - datapoint = dataset_cnn['test'][0:1] - summary, _ = summarize_tensorrt_llm(datapoint) - if runtime_rank == 0: - logger.info( - "---------------------------------------------------------") - logger.info("TensorRT-LLM Generated : ") - logger.info(f" Article : {datapoint['article']}") - logger.info(f"\n Highlights : {datapoint['highlights']}") - logger.info(f"\n Summary : {summary}") - logger.info( - "---------------------------------------------------------") - - if test_hf: - datapoint = dataset_cnn['test'][0:1] - summary, _ = summarize_hf(datapoint) - logger.info("---------------------------------------------------------") - logger.info("HF Generated : ") - logger.info(f" Article : {datapoint['article']}") - logger.info(f"\n Highlights : {datapoint['highlights']}") - logger.info(f"\n Summary : {summary}") - logger.info("---------------------------------------------------------") - - metric_tensorrt_llm = [load_metric("rouge") for _ in range(num_beams)] - metric_hf = [load_metric("rouge") for _ in range(num_beams)] - for i in range(num_beams): - metric_tensorrt_llm[i].seed = 0 - metric_hf[i].seed = 0 - - ite_count = 0 - data_point_idx = 0 - while (data_point_idx < len(dataset_cnn['test'])) 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_cnn['test'][data_point_idx:(data_point_idx + - max_batch_size)] - - if test_trt_llm: - profiler.start('tensorrt_llm') - summary_tensorrt_llm, tokens_tensorrt_llm = summarize_tensorrt_llm( - datapoint) - profiler.stop('tensorrt_llm') - - if test_hf: - profiler.start('hf') - summary_hf, tokens_hf = summarize_hf(datapoint) - profiler.stop('hf') - - if runtime_rank == 0: - if test_trt_llm: - for batch_idx in range(len(summary_tensorrt_llm)): - for beam_idx in range(num_beams): - metric_tensorrt_llm[beam_idx].add_batch( - predictions=[ - summary_tensorrt_llm[batch_idx][beam_idx] - ], - references=[datapoint['highlights'][batch_idx]]) - if test_hf: - for beam_idx in range(num_beams): - for batch_idx in range(len(summary_hf[beam_idx])): - metric_hf[beam_idx].add_batch( - predictions=[summary_hf[beam_idx][batch_idx]], - references=[datapoint['highlights'][batch_idx]]) - - logger.debug('-' * 100) - logger.debug(f"Article : {datapoint['article']}") - if test_trt_llm: - logger.debug(f'TensorRT-LLM Summary: {summary_tensorrt_llm}') - if test_hf: - logger.debug(f'HF Summary: {summary_hf}') - logger.debug(f"highlights : {datapoint['highlights']}") - - data_point_idx += max_batch_size - ite_count += 1 - - if runtime_rank == 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)' - ) - for beam_idx in range(num_beams): - logger.info(f"TensorRT-LLM beam {beam_idx} result") - 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].mid[2]*100}' - ) - - if args.check_accuracy and beam_idx == 0: - assert computed_metrics_tensorrt_llm['rouge1'].mid[ - 2] * 100 > args.tensorrt_llm_rouge1_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)' - ) - for beam_idx in range(num_beams): - logger.info(f"HF beam {beam_idx} result") - computed_metrics_hf = metric_hf[beam_idx].compute() - for key in computed_metrics_hf.keys(): - logger.info( - f' {key} : {computed_metrics_hf[key].mid[2]*100}') - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--hf_model_location', - type=str, - default='internlm-7b-hf') - parser.add_argument('--test_hf', action='store_true') - parser.add_argument('--test_trt_llm', action='store_true') - parser.add_argument('--data_type', - type=str, - choices=['fp32', 'fp16', 'bf16'], - default='auto') - parser.add_argument('--dataset_path', type=str, default='') - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument('--engine_dir', type=str, default='internlm_outputs') - parser.add_argument('--batch_size', type=int, default=1) - parser.add_argument('--max_ite', type=int, default=20) - parser.add_argument('--check_accuracy', action='store_true') - parser.add_argument('--tensorrt_llm_rouge1_threshold', - type=float, - default=15.0) - parser.add_argument('--num_beams', type=int, default=1) - parser.add_argument('--top_k', type=int, default=1) - - args = parser.parse_args() - - main(args) diff --git a/examples/internlm/weight.py b/examples/internlm/weight.py index 7293962c7404..e5052b10cce0 100644 --- a/examples/internlm/weight.py +++ b/examples/internlm/weight.py @@ -27,7 +27,7 @@ import tensorrt_llm.logger as logger from tensorrt_llm._utils import str_dtype_to_torch, torch_to_numpy from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import InternLMForCausalLM +from tensorrt_llm.models import LLaMAForCausalLM from tensorrt_llm.models.quantized.quant import get_dummy_quant_scales from tensorrt_llm.quantization import QuantMode @@ -181,7 +181,7 @@ def parse_ft_config(ini_file): def load_from_hf_internlm( - tensorrt_llm_internlm: tensorrt_llm.models.InternLMForCausalLM, + tensorrt_llm_internlm: tensorrt_llm.models.LLaMAForCausalLM, hf_internlm, mapping=Mapping(), dtype='float32'): @@ -374,7 +374,7 @@ def load_from_hf_internlm( def load_from_meta_internlm( - tensorrt_llm_internlm: tensorrt_llm.models.InternLMForCausalLM, + tensorrt_llm_internlm: tensorrt_llm.models.LLaMAForCausalLM, meta_ckpt_dir, mapping=Mapping(), dtype="float32"): @@ -565,7 +565,7 @@ def gather_embedding(cur_embed, name: str, num_ckpts): return -def load_from_binary(tensorrt_llm_internlm: InternLMForCausalLM, +def load_from_binary(tensorrt_llm_internlm: LLaMAForCausalLM, dir_path, mapping=Mapping(), fp16=False, @@ -1078,7 +1078,7 @@ def preprocess_groupwise_weight_params(weight_name, return -def load_from_awq_internlm(tensorrt_llm_internlm: InternLMForCausalLM, +def load_from_awq_internlm(tensorrt_llm_internlm: LLaMAForCausalLM, quant_ckpt_path, mapping=Mapping(), dtype="float16"): diff --git a/examples/llama/.gitignore b/examples/llama/.gitignore index 8e4cabd04493..b43358a79c1c 100644 --- a/examples/llama/.gitignore +++ b/examples/llama/.gitignore @@ -1,2 +1,3 @@ llama* tokenizer.model +*output* diff --git a/examples/llama/README.md b/examples/llama/README.md index ba80b9a39d89..d17add03a53b 100644 --- a/examples/llama/README.md +++ b/examples/llama/README.md @@ -4,11 +4,11 @@ This document shows how to build and run a LLaMA model in TensorRT-LLM on both s ## Overview -The TensorRT-LLM LLaMA implementation can be found in [tensorrt_llm/models/llama/model.py](../../tensorrt_llm/models/llama/model.py). The TensorRT-LLM LLaMA example code is located in [`examples/llama`](./). There are three main files in that folder:: +The TensorRT-LLM LLaMA implementation can be found in [tensorrt_llm/models/llama/model.py](../../tensorrt_llm/models/llama/model.py). The TensorRT-LLM LLaMA example code is located in [`examples/llama`](./). There are three main files: * [`build.py`](./build.py) to build the [TensorRT](https://developer.nvidia.com/tensorrt) engine(s) needed to run the LLaMA model, * [`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/cnn_dailymail) dataset using the model. + * and a shared [`../summarize.py`](../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/cnn_dailymail) dataset using the model. ## Support Matrix * FP16 @@ -163,6 +163,63 @@ RoPE scaling is supported through GPT Attention Plugin. You can add `--rotary_sc The implementation is identical to Huggingface's. Please refer to https://huggingface.co/docs/transformers/model_doc/llama2#transformers.LlamaConfig.rope_scaling for more details. +### Long context length +To use the model with Long context lengths, it is necessary to add `--multi_block_mode` in the build command to enable faster decoding in multihead attention. + + +A few LLaMA models are fine-tuned for long context length that TRT-LLM can support today. For example https://huggingface.co/Yukang/LongAlpaca-70B employs rotary scaling plus fine-tuning to support up to 32K context length. The following show the steps for running LongAlpaca-70B in TRT-LLM: + + +```bash +# Build 8-GPU engine with long context LLaMA model +python build.py --model_dir ./tmp/LongAlpaca-70B/ \ + --dtype float16 \ + --remove_input_padding \ + --use_gpt_attention_plugin float16 \ + --enable_context_fmha \ + --use_gemm_plugin float16 \ + --output_dir ./tmp/llama/70B/trt_engines/fp16/8-gpu/ \ + --world_size 8 \ + --tp_size 8 \ + --pp_size 1 \ + --multi_block_mode \ + --max_input_len 32768 \ + --max_output_len 16384 \ + --vocab_size=32001 \ + --rotary_scaling linear 8.0 + +# Get the long text data from Gutenberg Project +wget https://www.gutenberg.org/cache/epub/64317/pg64317.txt + +# Run with 8 GPUs +# Notice, `--input_tokens_limit ` is a convenience option to limit the input length for the data. +# It should be set to the maximum context length the model supports. Here the limit is set to 32K. +mpirun -n 8 --allow-run-as-root \ + python run.py \ + --max_output_len 128 \ + --input_tokens_limit 32768 \ + --input_tokens pg64317.txt \ + --engine_dir ./tmp/llama/70B/trt_engines/fp16/8-gpu/ \ + --tokenizer_dir ./tmp/LongAlpaca-70B/ +``` + +Note that if engine is built with contiguous KV cache (i.e., without the flag `--paged_kv_cache`), you may need to reduce the max batch size (`--max_batch_size`) to fit the whole model and the KV cache in the GPU memory. The ballpark estimate for runtime memory consumption is given by + +``` +Total memory = (Model size + KV cache size + Activation memory) / Parallelism +``` + +where +- The model size is `the number of parameters * the size of data type`. +- The KV cache size is `the total number of tokens * the size of KV cache data type * the number of layers * the KV hidden dimension` +- The activation memory is determined by TRT engine, which can be a few GBs regardless of the degree of parallelism used + +For LLaMA v2 70B FP16 weights + FP8 KV cache, the model size is 70B parameters * 2 bytes = 140GB. The KV cache size is 32K tokens * 1 bytes * 80 layers * 2048 KV hidden dimension = 5GB per 32K tokens. We have 145GB spread across 8 GPUs. The end result is ~18GB per GPU plus some GBs of flat scratch/activation memory allocated by TRT engine and the TRT-LLM runtime. + +Note that the KV hidden dimension is derived by the number of KV heads times hidden dimension of each head. LLaMA v2 70B has hidden dimension of 8192, and uses grouped-query attention where 8 key heads and 8 value heads are associated with 64 query heads. Each head has hidden dimension of 8192/64 = 128. So the hidden dimension for KV in total is 128 * 8 * 2 = 2048. + +The total number of tokens is determined by beam width, batch size, and maximum sequence length. + #### INT8 KV cache INT8 KV cache could be enabled to reduce memory footprint. It will bring more performance gains when batch size gets larger. @@ -198,14 +255,14 @@ python build.py --ft_model_dir=/llama/smooth_llama_7B/int8_kv_cache/1-gpu/ \ --use_weight_only ``` -Test with `summarize.py`: +Test with `../summarize.py`: ```bash -python summarize.py --test_trt_llm \ - --hf_model_location /llama-models/llama-7b-hf \ - --data_type fp16 \ - --engine_dir ./tmp/llama/7B/trt_engines/int8_kv_cache_weight_only/1-gpu \ - --test_hf +python ../summarize.py --test_trt_llm \ + --hf_model_dir /llama-models/llama-7b-hf \ + --data_type fp16 \ + --engine_dir ./tmp/llama/7B/trt_engines/int8_kv_cache_weight_only/1-gpu \ + --test_hf ``` **INT8 KV cache + AWQ** @@ -230,14 +287,14 @@ python build.py --model_dir ./tmp/llama/7B/ \ --ft_model_dir /llama/smooth_llama_7B/int8_kv_cache/1-gpu/ # Directory to look for INT8 scale of KV cache ``` -Test with `summarize.py`: +Test with `../summarize.py`: ```bash -python summarize.py --test_trt_llm \ - --hf_model_location /llama-models/llama-7b-hf \ - --data_type fp16 \ - --engine_dir ./tmp/llama/7B/trt_engines/int8_kv_cache_int4_AWQ/1-gpu \ - --test_hf +python ../summarize.py --test_trt_llm \ + --hf_model_dir /llama-models/llama-7b-hf \ + --data_type fp16 \ + --engine_dir ./tmp/llama/7B/trt_engines/int8_kv_cache_int4_AWQ/1-gpu \ + --test_hf ``` #### SmoothQuant @@ -294,13 +351,15 @@ python build.py --model_dir ./tmp/llama/70B \ --quantized_fp8_model_path ./quantized_fp8/llama_tp1_rank0.npz \ --dtype float16 \ --use_gpt_attention_plugin float16 \ - --use_gemm_plugin float16 \ --output_dir ./tmp/llama/70B/trt_engines/fp8/2-gpu/ \ --remove_input_padding \ + --enable_context_fmha \ --enable_fp8 \ --fp8_kv_cache \ + --strongly_typed \ --world_size 2 \ - --tp_size 2 + --tp_size 2 \ + --parallel_build ``` #### Groupwise quantization (AWQ/GPTQ) @@ -405,34 +464,34 @@ python3 run.py --max_output_len=50 \ ```bash # Run summarization using the LLaMA 7B model in FP16. -python summarize.py --test_trt_llm \ - --hf_model_location ./tmp/llama/7B/ \ - --data_type fp16 \ - --engine_dir ./tmp/llama/7B/trt_engines/fp16/1-gpu/ +python ../summarize.py --test_trt_llm \ + --hf_model_dir ./tmp/llama/7B/ \ + --data_type fp16 \ + --engine_dir ./tmp/llama/7B/trt_engines/fp16/1-gpu/ # Run summarization using the LLaMA 7B model quantized to INT8. -python summarize.py --test_trt_llm \ - --hf_model_location ./tmp/llama/7B/ \ - --data_type fp16 \ - --engine_dir ./tmp/llama/7B/trt_engines/weight_only/1-gpu/ +python ../summarize.py --test_trt_llm \ + --hf_model_dir ./tmp/llama/7B/ \ + --data_type fp16 \ + --engine_dir ./tmp/llama/7B/trt_engines/weight_only/1-gpu/ # Run summarization using the LLaMA 7B model in FP16 using two GPUs. mpirun -n 2 --allow-run-as-root \ - python summarize.py --test_trt_llm \ - --hf_model_location ./tmp/llama/7B/ \ - --data_type fp16 \ - --engine_dir ./tmp/llama/7B/trt_engines/fp16/2-gpu/ + python ../summarize.py --test_trt_llm \ + --hf_model_dir ./tmp/llama/7B/ \ + --data_type fp16 \ + --engine_dir ./tmp/llama/7B/trt_engines/fp16/2-gpu/ # Run summarization using the LLaMA 30B model in FP16 using two GPUs. mpirun -n 2 --allow-run-as-root \ - python summarize.py --test_trt_llm \ - --hf_model_location ./tmp/llama/30B/ \ - --data_type fp16 \ - --engine_dir ./tmp/llama/30B/trt_engines/fp16/2-gpu/ + python ../summarize.py --test_trt_llm \ + --hf_model_dir ./tmp/llama/30B/ \ + --data_type fp16 \ + --engine_dir ./tmp/llama/30B/trt_engines/fp16/2-gpu/ ``` -#### Mistral v1.0 -Mistral v1.0 is compatible with LLaMA interface and can be built and run using the same instructions. +#### Mistral v0.1 +Mistral v0.1 is compatible with LLaMA interface and can be built and run using the same instructions. Setting `--max_input_len`, corresponding to the `max_position_embeddings` in the original Mistral config explicitly regulates context size. The `--max_kv_cache_len` parameter is set to the `sliding_window` value in the config and regulates both sliding window attention in the context phase and rolling buffer cache in the generation phase. diff --git a/examples/llama/build.py b/examples/llama/build.py index c460627f3691..10f61e3df0a8 100644 --- a/examples/llama/build.py +++ b/examples/llama/build.py @@ -23,10 +23,11 @@ import torch.multiprocessing as mp from transformers import LlamaConfig, LlamaForCausalLM from weight import (get_scaling_factors, load_from_awq_llama, load_from_binary, - load_from_gptq_llama, load_from_hf_llama, - load_from_meta_llama) + load_from_gptq_llama, load_from_hf_checkpoint, + load_from_hf_llama, load_from_meta_llama) import tensorrt_llm +from tensorrt_llm import profiler from tensorrt_llm._utils import str_dtype_to_trt from tensorrt_llm.builder import Builder from tensorrt_llm.layers.attention import PositionEmbeddingType @@ -35,7 +36,6 @@ from tensorrt_llm.models import quantize_model from tensorrt_llm.network import net_guard from tensorrt_llm.plugin.plugin import ContextFMHAType -from tensorrt_llm.profiler import check_gpt_mem_usage from tensorrt_llm.quantization import QuantMode from weight import parse_ft_config # isort:skip @@ -194,6 +194,9 @@ def parse_arguments(): It is beneifical when batchxnum_heads cannot fully utilize GPU.' ) parser.add_argument('--visualize', default=False, action='store_true') + parser.add_argument('--load_by_shard', + action='store_true', + help='Load a pretrained model shard-by-shard.') parser.add_argument('--enable_debug_output', default=False, action='store_true') @@ -352,6 +355,9 @@ def parse_arguments(): type=int, default=0, help='Setting to a value > 0 enables support for prompt tuning.') + parser.add_argument('--gather_all_token_logits', + action='store_true', + default=False) args = parser.parse_args() tensorrt_llm.logger.set_level(args.log_level) @@ -494,6 +500,7 @@ def build_rank_engine(builder: Builder, assert args.n_layer % args.pp_size == 0, \ f"num_layers {args.n_layer} must be a multiple of pipeline parallelism size {args.pp_size}" + profiler.print_memory_usage(f'Rank {rank} Engine build starts') # Initialize Module tensorrt_llm_llama = tensorrt_llm.models.LLaMAForCausalLM( num_layers=args.n_layer, @@ -553,27 +560,35 @@ def build_rank_engine(builder: Builder, elif args.model_dir is not None: logger.info(f'Loading HF LLaMA ... from {args.model_dir}') tik = time.time() - hf_llama = LlamaForCausalLM.from_pretrained( - args.model_dir, - device_map={ - "model": "cpu", - "lm_head": "cpu" - }, # Load to CPU memory - torch_dtype="auto") + if not args.load_by_shard: + hf_llama = LlamaForCausalLM.from_pretrained( + args.model_dir, + device_map={ + "model": "cpu", + "lm_head": "cpu" + }, # Load to CPU memory + torch_dtype='auto', + ) + load_from_hf_llama(tensorrt_llm_llama, + hf_llama, + mapping=mapping, + dtype=args.dtype) + del hf_llama + else: + load_from_hf_checkpoint(tensorrt_llm_llama, + args.model_dir, + mapping, + dtype=args.dtype) tok = time.time() t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) logger.info(f'HF LLaMA loaded. Total time: {t}') - load_from_hf_llama(tensorrt_llm_llama, - hf_llama, - mapping=mapping, - dtype=args.dtype) - del hf_llama elif args.ft_model_dir is not None: load_from_binary(tensorrt_llm_llama, args.ft_model_dir, mapping, fp16=(args.dtype == 'float16'), multi_query_mode=(args.n_kv_head != args.n_head)) + profiler.print_memory_usage(f'Rank {rank} model weight loaded.') # Module -> Network network = builder.create_network() @@ -632,7 +647,7 @@ def build_rank_engine(builder: Builder, args.max_beam_width, args.max_num_tokens, prompt_embedding_table_size=args.max_prompt_embedding_table_size, - ) + gather_all_token_logits=args.gather_all_token_logits) tensorrt_llm_llama(*inputs) if args.enable_debug_output: # mark intermediate nodes' outputs @@ -674,6 +689,8 @@ def build(rank, args): # skip other ranks if parallel_build is enabled if args.parallel_build and cur_rank != rank: continue + tik = time.time() + # NOTE: when only int8 kv cache is used together with paged kv cache no int8 tensors are exposed to TRT int8_trt_flag = args.quant_mode.has_act_or_weight_quant() or ( not args.paged_kv_cache and args.quant_mode.has_int8_kv_cache()) @@ -692,6 +709,7 @@ def build(rank, args): hidden_act=args.hidden_act, max_position_embeddings=args.n_positions, max_batch_size=args.max_batch_size, + max_beam_width=args.max_beam_width, max_input_len=args.max_input_len, max_output_len=args.max_output_len, max_num_tokens=args.max_num_tokens, @@ -701,6 +719,7 @@ def build(rank, args): opt_level=args.builder_opt, max_prompt_embedding_table_size=args. max_prompt_embedding_table_size, + gather_all_token_logits=args.gather_all_token_logits, ) engine_name = get_engine_name(MODEL_NAME, args.dtype, args.tp_size, args.pp_size, cur_rank) @@ -715,7 +734,7 @@ def build(rank, args): kv_dtype = str_dtype_to_trt('int8') elif args.quant_mode.has_fp8_kv_cache(): kv_dtype = str_dtype_to_trt('fp8') - check_gpt_mem_usage( + profiler.check_gpt_mem_usage( engine=engine, kv_dtype=kv_dtype, use_gpt_attention_plugin=args.use_gpt_attention_plugin, @@ -735,6 +754,12 @@ def build(rank, args): serialize_engine(engine, os.path.join(args.output_dir, engine_name)) del engine + profiler.print_memory_usage(f'Rank {cur_rank} Engine serialized') + + tok = time.time() + t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) + logger.info( + f'Rank {cur_rank} Engine build time: {t} - {tok - tik} (sec)') if rank == 0: ok = builder.save_timing_cache( diff --git a/examples/llama/hf_llama_convert.py b/examples/llama/hf_llama_convert.py index f16627c30d04..c98bf1901123 100644 --- a/examples/llama/hf_llama_convert.py +++ b/examples/llama/hf_llama_convert.py @@ -23,12 +23,15 @@ import torch import torch.multiprocessing as multiprocessing from convert import split_and_save_weight, str_to_np_dtype +from datasets import load_dataset from smoothquant import (capture_activation_range, smooth_gemm, smooth_gemm_fc1_gate) from tqdm import tqdm from transformers import LlamaForCausalLM, LlamaTokenizer from transformers.models.llama.modeling_llama import LlamaDecoderLayer +from tensorrt_llm.logger import logger + def merge_qkv_scales(q_name, hf_model, scales, llama_qkv_para): layer_name_q = q_name.replace(".weight", "") @@ -170,7 +173,14 @@ def hf_gpt_converter(args): saved_dir = Path(args.out_dir) / f"{infer_tp}-gpu" saved_dir.mkdir(parents=True, exist_ok=True) - model = LlamaForCausalLM.from_pretrained(args.in_file, device_map="auto") + model = LlamaForCausalLM.from_pretrained(args.in_file, + torch_dtype="auto", + device_map="auto", + trust_remote_code=True) + if args.load_model_on_cpu: + model = model.float() + model = model.cpu() + torch.cuda.empty_cache() act_range = {} llama_qkv_para = {} @@ -180,9 +190,17 @@ def hf_gpt_converter(args): if args.smoothquant is not None or args.calibrate_kv_cache: os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get( "TOKENIZERS_PARALLELISM", "false") + if args.load_model_on_cpu: + logger.warning( + "Note that running capture_activation_range on cpu would be very small." + ) + dataset = load_dataset("ccdv/cnn_dailymail", + '3.0.0', + cache_dir=args.dataset_cache_dir) act_range = capture_activation_range( model, - LlamaTokenizer.from_pretrained(args.in_file, padding_side='left')) + LlamaTokenizer.from_pretrained(args.in_file, padding_side='left'), + dataset) if args.smoothquant is not None: smooth_llama_model(model, act_range, args.smoothquant, llama_qkv_para, llama_smoother) @@ -217,6 +235,8 @@ def hf_gpt_converter(args): continue ft_name = gpt_to_ft_name(name) + if args.convert_model_on_cpu: + param = param.cpu() if name.replace(".weight", "") in llama_smoother.keys(): smoother = llama_smoother[name.replace(".weight", "")] smoother = smoother.detach().cpu().numpy() @@ -319,6 +339,12 @@ def hf_gpt_converter(args): type=str, default="fp32", choices=["fp32", "fp16"]) + parser.add_argument("--dataset-cache-dir", + type=str, + default=None, + help="cache dir to load the hugging face dataset") + parser.add_argument("--load-model-on-cpu", action="store_true") + parser.add_argument("--convert-model-on-cpu", action="store_true") args = parser.parse_args() print("\n=============== Argument ===============") @@ -328,5 +354,6 @@ def hf_gpt_converter(args): assert (args.calibrate_kv_cache or args.smoothquant), \ "Either INT8 kv cache or SmoothQuant must be enabled for this script. Otherwise you can directly build engines from HuggingFace checkpoints, no need to do this FT-format conversion. " + logger.set_level("info") hf_gpt_converter(args) diff --git a/examples/llama/requirements.txt b/examples/llama/requirements.txt index a6789a325ba8..55f57b4de633 100644 --- a/examples/llama/requirements.txt +++ b/examples/llama/requirements.txt @@ -1,3 +1,4 @@ datasets==2.14.6 +evaluate~=0.4.1 rouge_score~=0.1.2 sentencepiece~=0.1.99 diff --git a/examples/llama/run.py b/examples/llama/run.py index 75512abea0b6..ac02340beaff 100644 --- a/examples/llama/run.py +++ b/examples/llama/run.py @@ -16,6 +16,7 @@ import csv import json from pathlib import Path +from typing import Union import numpy as np import torch @@ -46,6 +47,8 @@ def read_config(config_path: Path): use_gpt_attention_plugin = config['plugin_config']['gpt_attention_plugin'] remove_input_padding = config['plugin_config']['remove_input_padding'] dtype = config['builder_config']['precision'] + gather_all_token_logits = config['builder_config'][ + 'gather_all_token_logits'] tp_size = config['builder_config']['tensor_parallel'] pp_size = config['builder_config']['pipeline_parallel'] world_size = tp_size * pp_size @@ -85,6 +88,7 @@ def read_config(config_path: Path): remove_input_padding=remove_input_padding, dtype=dtype, quant_mode=quant_mode, + gather_all_token_logits=gather_all_token_logits, use_custom_all_reduce=use_custom_all_reduce, max_prompt_embedding_table_size=max_prompt_embedding_table_size) @@ -92,7 +96,8 @@ def read_config(config_path: Path): def parse_input(input_text: str, input_file: str, tokenizer, end_id: int, - remove_input_padding: bool): + remove_input_padding: bool, input_tokens_limit: Union[int, + None]): input_tokens = [] if input_file is None: input_tokens.append( @@ -108,10 +113,23 @@ def parse_input(input_text: str, input_file: str, tokenizer, end_id: int, for row in inputs: row = row[row != end_id] input_tokens.append(row) + elif input_file.endswith('.txt'): + with open(input_file, 'r', encoding='utf-8', + errors='replace') as txt_file: + input_text = txt_file.read() + input_tokens.append( + tokenizer.encode(input_text, add_special_tokens=False)) else: print('Input file format not supported.') raise SystemExit + # Cap max input tokens + if input_tokens_limit is not None: + print( + f"Maximum input number of tokens found as {max([len(x) for x in input_tokens])};" + f" will be capped to {input_tokens_limit}") + input_tokens = [x[-input_tokens_limit:] for x in input_tokens] + input_ids = None input_lengths = torch.tensor([len(x) for x in input_tokens], dtype=torch.int32, @@ -176,7 +194,6 @@ def print_output(output_ids, input_lengths, max_output_len, tokenizer, print(f'Output: \"{output_text}\"') output_ids = output_ids.reshape((-1, output_ids.size(2))) - print(output_ids) if output_csv is not None: output_file = Path(output_csv) @@ -218,6 +235,11 @@ def parse_arguments(): help= 'CSV or Numpy file containing tokenized input. Alternative to text input.', default=None) + parser.add_argument( + '--input_tokens_limit', + type=int, + help='Truncate input tokens if number exceeds the set limit value', + default=None) parser.add_argument('--output_csv', type=str, help='CSV file where the tokenized output is stored.', @@ -260,6 +282,7 @@ def generate( streaming_interval: int = 5, prompt_table: Path = None, tasks: str = None, + input_tokens_limit: Union[None, int] = None, ): tensorrt_llm.logger.set_level(log_level) @@ -294,10 +317,13 @@ def generate( if runtime_rank == 0: print(f"Running the {dtype} engine ...") - input_ids, input_lengths = parse_input(input_text, input_file, tokenizer, - EOS_TOKEN, - model_config.remove_input_padding) - print(input_ids) + input_ids, input_lengths = parse_input( + input_text, + input_file, + tokenizer, + EOS_TOKEN, + model_config.remove_input_padding, + input_tokens_limit=input_tokens_limit) max_input_length = torch.max(input_lengths).item() decoder.setup(input_lengths.size(0), @@ -333,6 +359,16 @@ def generate( print_output(output_ids, input_lengths, max_output_len, tokenizer, output_csv, output_npy, sequence_lengths) + if model_config.gather_all_token_logits: + if runtime_mapping.is_last_pp_rank(): + print( + f"context_logits.shape: {outputs['context_logits'].shape}") + print( + f"generation_logits.shape: {len(outputs['generation_logits']), outputs['generation_logits'][0].shape}" + ) + print(outputs['context_logits']) + print(outputs['generation_logits']) + if __name__ == '__main__': args = parse_arguments() diff --git a/examples/llama/smoothquant.py b/examples/llama/smoothquant.py index 4e4145cb4ebb..f1a7acd462a4 100644 --- a/examples/llama/smoothquant.py +++ b/examples/llama/smoothquant.py @@ -145,12 +145,15 @@ def smooth_ln_fcs(ln, fcs, act_scales, alpha=0.5): @torch.no_grad() -def capture_activation_range(model, tokenizer, num_samples=512, seq_len=512): +def capture_activation_range(model, + tokenizer, + dataset, + num_samples=512, + seq_len=512): model.eval() - next(model.parameters()).device + device = next(model.parameters()).device act_scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) - test_token_num = 923 tokenizer.pad_token = tokenizer.eos_token def stat_tensor(name, tensor, act_scales, key): @@ -181,22 +184,18 @@ def stat_input_hook(m, x, y, name): m.register_forward_hook( functools.partial(stat_input_hook, name=name))) - from datasets import load_dataset - dataset_cnn = load_dataset("ccdv/cnn_dailymail", '3.0.0') - for i in tqdm(range(num_samples), desc="calibrating model"): - datapoint = dataset_cnn['train'][i:i + 1] + datapoint = dataset['train'][i:i + 1] line = copy.copy(datapoint['article']) line[0] = line[0] + ' TL;DR: ' line[0] = line[0].strip() line[0] = line[0].replace(" n't", "n't") - line_encoded = tokenizer(line, - return_tensors="pt", - padding=True, - truncation=True)["input_ids"].type(torch.int64) - line_encoded = line_encoded[:, -test_token_num:] - line_encoded = line_encoded.cuda() - model(line_encoded) + input_ids = tokenizer(line, + return_tensors="pt", + max_length=seq_len, + padding=True, + truncation=True).input_ids.to(device) + model(input_ids) for h in hooks: h.remove() diff --git a/examples/llama/summarize_long.py b/examples/llama/summarize_long.py index 0ca1e41fa5cb..f2ce0178a2f5 100644 --- a/examples/llama/summarize_long.py +++ b/examples/llama/summarize_long.py @@ -18,12 +18,14 @@ import torch from datasets import load_dataset, load_metric -from summarize import TRTLLaMA from transformers import AutoModelForCausalLM, LlamaTokenizer import tensorrt_llm import tensorrt_llm.profiler as profiler from tensorrt_llm.logger import logger +from tensorrt_llm.quantization import QuantMode + +from build import get_engine_name # isort:skip def parse_args(): @@ -71,6 +73,76 @@ def parse_args(): return args +def TRTLLaMA(args, config): + dtype = config['builder_config']['precision'] + tp_size = config['builder_config']['tensor_parallel'] + pp_size = config['builder_config']['pipeline_parallel'] + 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['builder_config']['num_heads'] // tp_size + hidden_size = config['builder_config']['hidden_size'] // tp_size + vocab_size = config['builder_config']['vocab_size'] + num_layers = config['builder_config']['num_layers'] + use_gpt_attention_plugin = bool( + config['plugin_config']['gpt_attention_plugin']) + remove_input_padding = config['plugin_config']['remove_input_padding'] + num_kv_heads = config['builder_config'].get('num_kv_heads', num_heads) + paged_kv_cache = config['plugin_config']['paged_kv_cache'] + tokens_per_block = config['plugin_config']['tokens_per_block'] + use_custom_all_reduce = config['plugin_config'].get('use_custom_all_reduce', + False) + + quant_mode = QuantMode(config['builder_config']['quant_mode']) + if config['builder_config'].get('multi_query_mode', False): + tensorrt_llm.logger.warning( + "`multi_query_mode` config is deprecated. Please rebuild the engine." + ) + num_kv_heads = 1 + num_kv_heads = (num_kv_heads + tp_size - 1) // tp_size + + model_config = tensorrt_llm.runtime.ModelConfig( + vocab_size=vocab_size, + num_layers=num_layers, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + hidden_size=hidden_size, + paged_kv_cache=paged_kv_cache, + tokens_per_block=tokens_per_block, + gpt_attention_plugin=use_gpt_attention_plugin, + remove_input_padding=remove_input_padding, + use_custom_all_reduce=use_custom_all_reduce, + dtype=dtype, + quant_mode=quant_mode) + + runtime_rank = tensorrt_llm.mpi_rank() + runtime_mapping = tensorrt_llm.Mapping(world_size, + runtime_rank, + tp_size=tp_size, + pp_size=pp_size) + torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node) + + engine_name = get_engine_name('llama', dtype, tp_size, pp_size, + runtime_rank) + serialize_path = os.path.join(args.engine_dir, engine_name) + + tensorrt_llm.logger.set_level(args.log_level) + + profiler.start('load tensorrt_llm engine') + with open(serialize_path, 'rb') as f: + engine_buffer = f.read() + decoder = tensorrt_llm.runtime.GenerationSession(model_config, + engine_buffer, + runtime_mapping) + profiler.stop('load tensorrt_llm engine') + tensorrt_llm.logger.info( + f'Load engine takes: {profiler.elapsed_time_in_sec("load tensorrt_llm engine")} sec' + ) + return decoder + + def get_long_texts(dataset_openweb): for datapoint in dataset_openweb["train"]: text = datapoint["text"] diff --git a/examples/llama/weight.py b/examples/llama/weight.py index 0211c5f892ce..2973703049bc 100644 --- a/examples/llama/weight.py +++ b/examples/llama/weight.py @@ -135,13 +135,21 @@ def extract_layer_idx(name): return None -def split(v, tp_size, idx, dim=0): +def split(v: Union[np.ndarray, torch.Tensor], + tp_size: int, + tp_rank: int, + dim=0): if tp_size == 1: return v - if len(v.shape) == 1: - return np.ascontiguousarray(np.split(v, tp_size)[idx].copy()) + assert len(v.shape) > 1 or dim == 0 + if isinstance(v, np.ndarray): + return np.ascontiguousarray( + np.split(v, tp_size, axis=dim)[tp_rank].copy()) else: - return np.ascontiguousarray(np.split(v, tp_size, axis=dim)[idx].copy()) + assert v.shape[dim] % tp_size == 0, \ + 'Unable to split: shape={v.shape} (dim={dim}) tp_size={tp_size}.' + split_size = v.shape[dim] // tp_size + return v.split(split_size, dim=dim)[tp_rank].clone().detach() def dup_kv_weight(v, num_head, tp_size): @@ -151,7 +159,7 @@ def dup_kv_weight(v, num_head, tp_size): v = v.reshape(num_head, head_size, -1)[:, None, :, :].expand(num_head, reps, head_size, v.shape[1]) - return v.reshape(num_head * reps * head_size, -1).clone() + return v.reshape(num_head * reps * head_size, -1).clone().detach() def parse_ft_config(ini_file): @@ -237,6 +245,14 @@ def load_from_hf_llama(tensorrt_llm_llama: tensorrt_llm.models.LLaMAForCausalLM, tensorrt_llm_llama.ln_f.weight.value = v elif 'lm_head.weight' in k: if mapping.is_last_pp_rank(): + vocab_size = tensorrt_llm_llama.vocab_embedding.num_embeddings + if vocab_size % mapping.tp_size != 0: + # padding + vocab_size_padded = tensorrt_llm_llama.lm_head.out_features * mapping.tp_size + pad_width = vocab_size_padded - vocab_size + v = np.pad(v, ((0, pad_width), (0, 0)), + 'constant', + constant_values=0) tensorrt_llm_llama.lm_head.weight.value = np.ascontiguousarray( split(v, mapping.tp_size, mapping.tp_rank)) else: @@ -337,7 +353,183 @@ def load_from_hf_llama(tensorrt_llm_llama: tensorrt_llm.models.LLaMAForCausalLM, tok = time.time() t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) tensorrt_llm.logger.info(f'Weights loaded. Total time: {t}') - return + + +def load_from_hf_checkpoint( + tensorrt_llm_llama: tensorrt_llm.models.LLaMAForCausalLM, + model_dir: Union[str, Path], + mapping=Mapping(), + dtype: Union[str, torch.dtype] = torch.float32, +): + tensorrt_llm.logger.info('Loading weights from HF LLaMA...') + tik = time.time() + if isinstance(dtype, str): + dtype = tensorrt_llm._utils.str_dtype_to_torch(dtype) + + model_dir = Path(model_dir) + + from transformers import AutoConfig + hf_config = AutoConfig.from_pretrained(model_dir) + + quant_mode = getattr(tensorrt_llm_llama, 'quant_mode', QuantMode(0)) + if quant_mode.is_int8_weight_only(): + plugin_weight_only_quant_type = torch.int8 + elif quant_mode.is_int4_weight_only(): + plugin_weight_only_quant_type = torch.quint4x2 + use_weight_only = quant_mode.is_weight_only() + num_kv_heads = tensorrt_llm_llama.num_kv_heads + mha_mode = num_kv_heads == tensorrt_llm_llama.num_heads + + # Load examples/common/utils.py + import sys + sys.path.append(str(Path(__file__).parent.parent)) + from common import utils + + layers_range = tensorrt_llm_llama.get_transformer_layers( + mapping, tensorrt_llm_llama.num_layers) + + def _is_qkv_weight(name): + for k in ['q_proj', 'k_proj', 'v_proj']: + if 'self_attn' in name and k in name: + return True + return False + + # Function to make a fused qkv matrix. + def _fuse_qkv(name, params): + # if param[name] is None: + # return None + i = utils.retrieved_layer_index_from_name(name) + prefix = f'model.layers.{i}.self_attn.' + q_weight = params[prefix + 'q_proj.weight'] + k_weight = params[prefix + 'k_proj.weight'] + v_weight = params[prefix + 'v_proj.weight'] + if not mha_mode: + head_size = tensorrt_llm_llama.hidden_size // tensorrt_llm_llama.num_heads + if num_kv_heads < mapping.tp_size: + # duplicate the KV heads up to tensor_parallel + k_weight = dup_kv_weight(k_weight, num_kv_heads, + mapping.tp_size) + v_weight = dup_kv_weight(v_weight, num_kv_heads, + mapping.tp_size) + assert (k_weight.shape[0] % (mapping.tp_size * head_size)) == 0 + assert (v_weight.shape[0] % (mapping.tp_size * head_size)) == 0 + qkv_weight = [q_weight, k_weight, v_weight] + else: + qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0) + # To skip other weights (q / k / v) + for k in ['q_proj.weight', 'k_proj.weight', 'v_proj.weight']: + params[prefix + k] = None + return qkv_weight + + for model_file in utils.iterate_shard_files(model_dir, + rank=mapping.tp_rank, + progress_bar=False): + logger.debug(f'Loading file {str(model_file)}...') + model_params = utils.load_state_dict(model_file, dtype=dtype) + for name, param in model_params.items(): + logger.debug(f'Converting weight {name}...') + i = utils.retrieved_layer_index_from_name(name) + if i is None: + layer = None + else: + if i not in layers_range: + continue + layer = tensorrt_llm_llama.layers[i - layers_range[0]] + + if 'model.embed_tokens.weight' in name: + if hf_config.tie_word_embeddings: + # lm_head.weight has the same weights as embedding + if mapping.is_last_pp_rank(): + tensorrt_llm_llama.lm_head.weight.value = split( + param, mapping.tp_size, mapping.tp_rank) + if tensorrt_llm_llama.use_parallel_embedding: + param = split(param, mapping.tp_size, mapping.tp_rank, + tensorrt_llm_llama.embedding_sharding_dim) + if mapping.is_first_pp_rank(): + tensorrt_llm_llama.vocab_embedding.weight.value = param + elif 'model.norm.weight' in name: + if mapping.is_last_pp_rank(): + tensorrt_llm_llama.ln_f.weight.value = param + elif 'lm_head.weight' in name: + if mapping.is_last_pp_rank(): + tensorrt_llm_llama.lm_head.weight.value = split( + param, mapping.tp_size, mapping.tp_rank) + elif 'input_layernorm.weight' in name: + layer.input_layernorm.weight.value = param + elif 'post_attention_layernorm.weight' in name: + layer.post_layernorm.weight.value = param + elif _is_qkv_weight(name) and model_params[name] is not None: + param = _fuse_qkv(name, model_params) + if not mha_mode: + assert isinstance(param, list) and len(param) == 3 + wq = split(param[0], mapping.tp_size, mapping.tp_rank) + wk = split(param[1], mapping.tp_size, mapping.tp_rank) + wv = split(param[2], mapping.tp_size, mapping.tp_rank) + split_v = torch.cat((wq, wk, wv)) + else: + q_emb = param.shape[0] // 3 + model_emb = param.shape[1] + param = param.reshape(3, q_emb, model_emb) + split_v = split(param, + mapping.tp_size, + mapping.tp_rank, + dim=1) + split_v = split_v.reshape(3 * (q_emb // mapping.tp_size), + model_emb) + + if use_weight_only: + param = split_v.transpose() + processed_torch_weights, torch_weight_scales = \ + torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( + param, plugin_weight_only_quant_type) + layer.attention.qkv.weight.value = processed_torch_weights + layer.attention.qkv.per_channel_scale.value = torch_weight_scales + else: + layer.attention.qkv.weight.value = split_v + elif 'self_attn.o_proj.weight' in name: + split_v = split(param, mapping.tp_size, mapping.tp_rank, dim=1) + if use_weight_only: + processed_torch_weights, torch_weight_scales = \ + torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( + split_v.transpose(), plugin_weight_only_quant_type) + layer.attention.dense.weight.value = processed_torch_weights + layer.attention.dense.per_channel_scale.value = torch_weight_scales + else: + layer.attention.dense.weight.value = split_v + elif 'mlp.up_proj.weight' in name: + split_v = split(param, mapping.tp_size, mapping.tp_rank, dim=0) + if use_weight_only: + processed_torch_weights, torch_weight_scales = \ + torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( + split_v.transpose(), plugin_weight_only_quant_type) + layer.mlp.gate.weight.value = processed_torch_weights + layer.mlp.gate.per_channel_scale.value = torch_weight_scales + else: + layer.mlp.gate.weight.value = split_v + elif 'mlp.down_proj.weight' in name: + split_v = split(param, mapping.tp_size, mapping.tp_rank, dim=1) + if use_weight_only: + processed_torch_weights, torch_weight_scales = \ + torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( + split_v.transpose(), plugin_weight_only_quant_type) + layer.mlp.proj.weight.value = processed_torch_weights + layer.mlp.proj.per_channel_scale.value = torch_weight_scales + else: + layer.mlp.proj.weight.value = split_v + elif 'mlp.gate_proj.weight' in name: + split_v = split(param, mapping.tp_size, mapping.tp_rank, dim=0) + if use_weight_only: + processed_torch_weights, torch_weight_scales = \ + torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( + split_v.transpose(), plugin_weight_only_quant_type) + layer.mlp.fc.weight.value = processed_torch_weights + layer.mlp.fc.per_channel_scale.value = torch_weight_scales + else: + layer.mlp.fc.weight.value = split_v + del model_params + tok = time.time() + t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) + tensorrt_llm.logger.info(f'Weights loaded. Total time: {t}') def load_from_meta_llama( @@ -648,9 +840,10 @@ def set_smoother(module, dir_path, base_name, shape, rank): tensorrt_llm_llama.lm_head.weight.value = np.ascontiguousarray( split(lm_head_weight, mapping.tp_size, mapping.tp_rank)) + layers_per_pipeline_stage = tensorrt_llm_llama.num_layers // mapping.pp_size layers_range = list( - range(mapping.pp_rank * tensorrt_llm_llama.num_layers, - (mapping.pp_rank + 1) * tensorrt_llm_llama.num_layers, 1)) + range(mapping.pp_rank * layers_per_pipeline_stage, + (mapping.pp_rank + 1) * layers_per_pipeline_stage, 1)) for i in layers_range: n_groups = n_head // n_kv_head @@ -658,7 +851,7 @@ def set_smoother(module, dir_path, base_name, shape, rank): 3 * n_embd // mapping.tp_size) if not multi_query_mode else ( n_embd // mapping.tp_size + (n_embd // n_head * n_groups) // mapping.tp_size * 2) - idx = i - mapping.pp_rank * tensorrt_llm_llama.num_layers + idx = i - mapping.pp_rank * layers_per_pipeline_stage tensorrt_llm_llama.layers[idx].input_layernorm.weight.value = (fromfile( dir_path, 'model.layers.' + str(i) + '.input_layernorm.weight.bin')) t = fromfile( @@ -684,7 +877,7 @@ def set_smoother(module, dir_path, base_name, shape, rank): torch.tensor(t), plugin_weight_only_quant_type) dst.value = processed_torch_weights.numpy() scales = tensorrt_llm_llama.layers[ - i].attention.qkv.per_channel_scale + idx].attention.qkv.per_channel_scale scales.value = torch_weight_scales.numpy() else: dst.value = np.ascontiguousarray(np.transpose(t, [1, 0])) @@ -711,7 +904,7 @@ def set_smoother(module, dir_path, base_name, shape, rank): torch.tensor(t), plugin_weight_only_quant_type) dst.value = processed_torch_weights.numpy() scales = tensorrt_llm_llama.layers[ - i].attention.dense.per_channel_scale + idx].attention.dense.per_channel_scale scales.value = torch_weight_scales.numpy() else: dst.value = np.ascontiguousarray(np.transpose(t, [1, 0])) @@ -738,11 +931,11 @@ def set_smoother(module, dir_path, base_name, shape, rank): quant_per_channel, rank=mapping.tp_rank) elif use_weight_only: - dst = tensorrt_llm_llama.layers[i].mlp.fc.weight + dst = tensorrt_llm_llama.layers[idx].mlp.fc.weight processed_torch_weights, torch_weight_scales = torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( torch.tensor(t), plugin_weight_only_quant_type) dst.value = processed_torch_weights.numpy() - scales = tensorrt_llm_llama.layers[i].mlp.fc.per_channel_scale + scales = tensorrt_llm_llama.layers[idx].mlp.fc.per_channel_scale scales.value = torch_weight_scales.numpy() else: tensorrt_llm_llama.layers[ @@ -766,11 +959,11 @@ def set_smoother(module, dir_path, base_name, shape, rank): quant_per_channel, rank=mapping.tp_rank) elif use_weight_only: - dst = tensorrt_llm_llama.layers[i].mlp.gate.weight + dst = tensorrt_llm_llama.layers[idx].mlp.gate.weight processed_torch_weights, torch_weight_scales = torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( torch.tensor(t), plugin_weight_only_quant_type) dst.value = processed_torch_weights.numpy() - scales = tensorrt_llm_llama.layers[i].mlp.gate.per_channel_scale + scales = tensorrt_llm_llama.layers[idx].mlp.gate.per_channel_scale scales.value = torch_weight_scales.numpy() else: tensorrt_llm_llama.layers[ @@ -794,11 +987,11 @@ def set_smoother(module, dir_path, base_name, shape, rank): 'model.layers.' + str(i) + '.mlp.proj', [1, inter_size // mapping.tp_size], mapping.tp_rank) elif use_weight_only: - dst = tensorrt_llm_llama.layers[i].mlp.proj.weight + dst = tensorrt_llm_llama.layers[idx].mlp.proj.weight processed_torch_weights, torch_weight_scales = torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( torch.tensor(t), plugin_weight_only_quant_type) dst.value = processed_torch_weights.numpy() - scales = tensorrt_llm_llama.layers[i].mlp.proj.per_channel_scale + scales = tensorrt_llm_llama.layers[idx].mlp.proj.per_channel_scale scales.value = torch_weight_scales.numpy() else: tensorrt_llm_llama.layers[idx].mlp.proj.weight.value = ( diff --git a/examples/mpt/README.md b/examples/mpt/README.md index 894091b8654c..aee6ae6ea63b 100644 --- a/examples/mpt/README.md +++ b/examples/mpt/README.md @@ -1,6 +1,6 @@ # MPT -This document explains how to build the [MPT](https://huggingface.co/mosaicml/mpt-7b) model using TensorRT-LLM and run on a single GPU and a single node with multiple GPUs +This document explains how to build the [MPT](https://huggingface.co/mosaicml/mpt-7b) model using TensorRT-LLM and run on a single GPU and a single node with multiple GPUs. ## Overview Currently we use `tensorrt_llm.models.GPTLMHeadModel` to build TRT engine for MPT models. @@ -12,13 +12,14 @@ Support for float16, float32 and bfloat16 conversion. Just change `data_type` fl * INT8 & INT4 Weight-Only * FP8 KV CACHE * Tensor Parallel + * MHA, MQA & GQA * STRONGLY TYPED #### MPT 7B ### 1. Convert weights from HF Transformers to FT format -The [`hf_gpt_convert.py`](./convert_hf_mpt_to_ft.py) script allows you to convert weights from HF Transformers format to FT format. +The [`convert_hf_mpt_to_ft.py`](./convert_hf_mpt_to_ft.py) script allows you to convert weights from HF Transformers format to FT format. ```bash python convert_hf_mpt_to_ft.py -i mosaicml/mpt-7b -o ./ft_ckpts/mpt-7b/fp16/ -t float16 @@ -85,7 +86,6 @@ Examples of build invocations: ```bash # Build 4-GPU MPT-30B float16 engines -# ALiBi is not supported with GPT attention plugin so we can't use that plugin to increase runtime performance python3 build.py --world_size=4 \ --parallel_build \ --max_batch_size 64 \ @@ -103,3 +103,98 @@ python3 build.py --world_size=4 \ # Run 4-GPU MPT7B TRT engine on a sample input prompt mpirun -n 4 --allow-run-as-root python run.py --engine_dir ./trt_engines/mpt-30b/fp16/4-gpu/ --max_output_len 10 ``` + +#### Replit Code V-1.5 3B +Same commands can be changed to convert [Replit Code V-1.5 3B](https://huggingface.co/replit/replit-code-v1_5-3b) to TRT LLM format. Below is an example to build Replit Code V-1.5 3B fp16 2-way tensor parallelized TRT engine. + +### 1. Convert weights from HF Transformers to FT format + +The [`convert_hf_mpt_to_ft.py`](./convert_hf_mpt_to_ft.py) script allows you to convert weights from HF Transformers format to FT format. + + +```bash +python convert_hf_mpt_to_ft.py -i ./replit-code-v1_5-3b -o ./ft_ckpts/replit-code-v1_5-3b/bf16/ --tensor_parallelism 2 -t bfloat16 +``` + +`--infer_gpu_num 2` is used to convert to FT format with 2-way tensor parallelism + + +### 2. Build TensorRT engine(s) + +Examples of build invocations: + +```bash +# Build 2-GPU Replit Code V-1.5 3B bfloat16 engines +python3 build.py --world_size=2 \ + --parallel_build \ + --max_batch_size 16 \ + --max_input_len 512 \ + --max_output_len 64 \ + --use_gpt_attention_plugin \ + --use_gemm_plugin \ + --model_dir ./ft_ckpts/replit-code-v1_5-3b/bf16/2-gpu \ + --output_dir=./trt_engines/replit-code-v1_5-3b/bf16/2-gpu +``` +Here is the partial output of above command. + +```bash +[11/15/2023-02:47:50] [TRT] [I] Total Activation Memory: 738233344 +[11/15/2023-02:47:51] [TRT] [I] Total Weights Memory: 3523622456 +[11/15/2023-02:47:51] [TRT] [I] [MemUsageChange] Init cuBLAS/cuBLASLt: CPU +0, GPU +64, now: CPU 8316, GPU 5721 (MiB) +[11/15/2023-02:47:51] [TRT] [I] [MemUsageChange] Init cuDNN: CPU +0, GPU +64, now: CPU 8316, GPU 5785 (MiB) +[11/15/2023-02:47:51] [TRT] [I] [MemUsageStats] Peak memory usage of TRT CPU/GPU memory allocators: CPU 192 MiB, GPU 3361 MiB +[11/15/2023-02:47:51] [TRT] [I] [MemUsageChange] TensorRT-managed allocation in building engine: CPU +0, GPU +3361, now: CPU 0, GPU 3361 (MiB) +[11/15/2023-02:47:51] [TRT] [I] [MemUsageStats] Peak memory usage during Engine building and serialization: CPU: 12851 MiB +[11/15/2023-02:47:51] [TRT-LLM] [I] Total time of building gpt_bfloat16_tp2_rank1.engine: 00:00:04 +[11/15/2023-02:47:51] [TRT-LLM] [I] Serializing engine to trt_engines/replit-code-v1_5-3b/bf16/2-gpu/gpt_bfloat16_tp2_rank1.engine... +[11/15/2023-02:48:02] [TRT-LLM] [I] Engine serialized. Total time: 00:00:10 +[11/15/2023-02:48:02] [TRT-LLM] [I] Timing cache serialized to model.cache +[11/15/2023-02:48:02] [TRT-LLM] [I] Total time of building all 2 engines: 00:01:21 +``` + +### 3. Run TRT engine to check if the build was correct + +```bash +# Run 2-GPU Replit Code V-1.5 3B TRT engine on a sample input prompt +mpirun -n 2 --allow-run-as-root python run.py --engine_dir ./trt_engines/replit-code-v1_5-3b/bf16/2-gpu/ --max_output_len 64 --input_text "def fibonacci" --tokenizer ./replit-code-v1_5-3b/ +``` + +Here is the output of above command. +```bash +Input: "def fibonacci" +Output: "(n): + if n == 0: + return 0 + elif n == 1: + return 1 + else: + return fibonacci(n-1) + fibonacci(n-2) + +print(fibonacci(10))" +``` +#### FP8 Post-Training Quantization + +The example below uses the NVIDIA AMMO (AlgorithMic Model Optimization) toolkit for the model quantization process. + +First make sure AMMO toolkit is installed (see [examples/quantization/README.md](/examples/quantization/README.md#preparation)) + +After successfully running the script, the output should be in .npz format, e.g. `quantized_fp8/llama_tp_1_rank0.npz`, +where FP8 scaling factors are stored. + +```bash +# Quantize MPT 7B into FP8 and export a single-rank checkpoint +python examples/quantization/quantize.py --model_dir .mosaicml/mpt-7b \ + --dtype float16 \ + --qformat fp8 \ + --export_path ./quantized_fp8 + +# Build MPT 7B TP using binary checkpoint + PTQ scaling factors from the single-rank checkpoint +python build.py --model_dir ft_ckpts/mpt-7b/fp16 \ + --quantized_fp8_model_path ./quantized_fp8/mpt_tp1_rank0.npz \ + --use_gpt_attention_plugin \ + --use_gemm_plugin \ + --output_dir trt_engines/mpt-7b/fp8/1-gpu/ \ + --remove_input_padding \ + --enable_fp8 \ + --fp8_kv_cache +``` diff --git a/examples/mpt/build.py b/examples/mpt/build.py index f214f7aee723..6f5fe99bae6f 100644 --- a/examples/mpt/build.py +++ b/examples/mpt/build.py @@ -31,7 +31,7 @@ from tensorrt_llm.plugin.plugin import ContextFMHAType from tensorrt_llm.quantization import QuantMode -from weight import load_from_ft, parse_ft_config, check_embedding_share # isort:skip +from weight import get_scaling_factors, load_from_ft, parse_ft_config, check_embedding_share # isort:skip MODEL_NAME = "gpt" @@ -87,6 +87,7 @@ def parse_arguments(args): parser.add_argument('--n_positions', type=int, default=1024) parser.add_argument('--n_embd', type=int, default=1024) parser.add_argument('--n_head', type=int, default=16) + parser.add_argument('--n_kv_head', type=int, default=None) parser.add_argument('--hidden_act', type=str, default='gelu') parser.add_argument( '--rotary_pct', @@ -153,14 +154,6 @@ def parse_arguments(args): help= 'The path to save the serialized engine files, timing cache file and model configs' ) - parser.add_argument( - "--multi_query_mode", - "-mq", - default=False, - action='store_true', - help= - "Whether this model uses multi-query attention mechanism (default: False)" - ) parser.add_argument('--remove_input_padding', default=False, action='store_true') @@ -300,6 +293,11 @@ def parse_arguments(args): choices=PositionEmbeddingType.choices(), help='Set the position embedding type.', ) + parser.add_argument( + '--quantized_fp8_model_path', + type=str, + default=None, + help='Path of a quantized model checkpoint in .npz format') args = parser.parse_args(args) logger.set_level(args.log_level) @@ -309,7 +307,8 @@ def parse_arguments(args): if args.model_dir is not None: logger.info(f"Setting model configuration from {args.model_dir}.") - n_embd, n_head, n_layer, n_positions, vocab_size, _, hidden_act, rotary_pct, bias, inter_size, multi_query_mode, dtype, prompt_num_tasks, prompt_max_vocab_size, position_embedding_type = parse_ft_config( + + n_embd, n_head, n_layer, n_positions, vocab_size, _, hidden_act, rotary_pct, bias, inter_size, n_kv_head, dtype, prompt_num_tasks, prompt_max_vocab_size, position_embedding_type = parse_ft_config( Path(args.model_dir) / "config.ini") args.n_embd = n_embd args.n_head = n_head @@ -321,7 +320,7 @@ def parse_arguments(args): args.bias = bias args.dtype = dtype args.inter_size = inter_size - args.multi_query_mode = multi_query_mode + args.n_kv_head = n_kv_head args.position_embedding_type = position_embedding_type plugins_args = [ 'use_gpt_attention_plugin', 'use_gemm_plugin', 'use_layernorm_plugin', @@ -432,38 +431,27 @@ def build_rank_engine(builder: Builder, apply_query_key_layer_scaling, quant_mode=args.quant_mode, bias=args.bias, - multi_query_mode=args.multi_query_mode, + num_kv_heads=args.n_kv_head, use_prompt_tuning=args.max_prompt_embedding_table_size > 0, use_parallel_embedding=args.use_parallel_embedding, embedding_sharding_dim=args.embedding_sharding_dim, share_embedding_table=share_embedding_table) - if args.use_smooth_quant or args.use_weight_only: - tensorrt_llm_gpt = quantize_model(tensorrt_llm_gpt, args.quant_mode) + quantize_kwargs = {} + if args.enable_fp8 or args.fp8_kv_cache: + logger.info(f'Loading scaling factors from ' + f'{args.quantized_fp8_model_path}') + quant_scales = get_scaling_factors(args.quantized_fp8_model_path, + num_layers=args.n_layer, + quant_mode=args.quant_mode) + quantize_kwargs = {"quant_scales": quant_scales} + tensorrt_llm_gpt = quantize_model(tensorrt_llm_gpt, args.quant_mode, + **quantize_kwargs) if args.model_dir is not None: - gpt_dummy_fp8_scaling_factors = { - 'fc_act': [0.5 for _ in range(args.n_layer)], - 'fc_weights': [0.5 for _ in range(args.n_layer)], - 'proj_act': [0.5 for _ in range(args.n_layer)], - 'proj_weights': [0.5 for _ in range(args.n_layer)], - 'qkv_act': [0.5 for _ in range(args.n_layer)], - 'qkv_weights': [0.5 for _ in range(args.n_layer)], - 'qkv_output': [0.5 for _ in range(args.n_layer)], - 'dense_act': [0.5 for _ in range(args.n_layer)], - 'dense_weights': [0.5 for _ in range(args.n_layer)], - } - - load_from_ft(tensorrt_llm_gpt, - args.model_dir, - rank, - args.world_size, - args.dtype, - args.use_parallel_embedding, - args.embedding_sharding_dim, - share_embedding_table, - scaling_factors=gpt_dummy_fp8_scaling_factors - if args.enable_fp8 else None) + load_from_ft(tensorrt_llm_gpt, args.model_dir, rank, args.world_size, + args.dtype, args.use_parallel_embedding, + args.embedding_sharding_dim, share_embedding_table) # Module -> Network network = builder.create_network() @@ -566,18 +554,19 @@ def build(rank, args): parallel_build=args.parallel_build, num_layers=args.n_layer, num_heads=args.n_head, + num_kv_heads=args.n_kv_head if args.n_kv_head else args.n_head, hidden_size=args.n_embd, vocab_size=args.vocab_size, hidden_act=args.hidden_act, max_position_embeddings=args.n_positions, apply_query_key_layer_scaling=apply_query_key_layer_scaling, max_batch_size=args.max_batch_size, + max_beam_width=args.max_beam_width, max_input_len=args.max_input_len, max_output_len=args.max_output_len, max_num_tokens=args.max_num_tokens, int8=int8_trt_flag, opt_level=args.builder_opt, - multi_query_mode=args.multi_query_mode, strongly_typed=args.strongly_typed, use_prompt_tuning=args.max_prompt_embedding_table_size > 0, quant_mode=args.quant_mode, diff --git a/examples/mpt/convert_hf_mpt_to_ft.py b/examples/mpt/convert_hf_mpt_to_ft.py index a8d1cd177e19..2f77113d0142 100644 --- a/examples/mpt/convert_hf_mpt_to_ft.py +++ b/examples/mpt/convert_hf_mpt_to_ft.py @@ -90,10 +90,6 @@ def convert_weight_to_ft_each(out_dir: str, tensor_parallelism: int, for j in range(tensor_parallelism): save_path = os.path.join(out_dir, f'model.{tensor_name}.{j}.bin') split_vals[j].tofile(save_path) - if config['no_bias']: - fake_weight_path = os.path.join(out_dir, f'model.{tensor_name}.bin') - write_zero_bias(tensor_name, fake_weight_path, data.shape[-1], - data_type) elif tensor_name.find('mlp.dense_4h_to_h.weight') != -1: assert data.shape == ( @@ -105,10 +101,6 @@ def convert_weight_to_ft_each(out_dir: str, tensor_parallelism: int, for j in range(tensor_parallelism): save_path = os.path.join(out_dir, f'model.{tensor_name}.{j}.bin') split_vals[j].tofile(save_path) - if config['no_bias']: - fake_weight_path = os.path.join(out_dir, f'model.{tensor_name}.bin') - write_zero_bias(tensor_name, fake_weight_path, data.shape[-1], - data_type) elif tensor_name.find('mlp.dense_h_to_4h.weight') != -1: assert data.shape == ( @@ -121,9 +113,6 @@ def convert_weight_to_ft_each(out_dir: str, tensor_parallelism: int, for j in range(tensor_parallelism): save_path = os.path.join(out_dir, f'model.{tensor_name}.{j}.bin') split_vals[j].tofile(save_path) - if config['no_bias']: - write_zero_bias(tensor_name, save_path, split_vals[j].shape[-1], - data_type) elif tensor_name.find('mlp.dense_h_to_4h.bias') != -1: assert data.shape == ( @@ -147,21 +136,38 @@ def convert_weight_to_ft_each(out_dir: str, tensor_parallelism: int, split_vals[j].tofile(save_path) elif tensor_name.find('attention.query_key_value.weight') != -1: - assert data.shape == ( - 3 * config['d_model'], - config['d_model']), f'unexpected dim for {tensor_name}' - # nn.Linear weights are transposed - data = data.T - - data = data.reshape(config['d_model'], 3, config['d_model']) - split_vals = np.split(data, tensor_parallelism, axis=-1) + if 'kv_n_heads' in config['attn_config']: + # Multi-query or grouped query attention + head_dim = config['d_model'] // config['n_heads'] + assert data.shape == ( + config['d_model'] + + 2 * config['attn_config']['kv_n_heads'] * head_dim, + config['d_model']), f'unexpected dim for {tensor_name}' + # nn.Linear weights are transposed + data = data.T + w_q, w_k, w_v = np.split(data, [ + config['d_model'], config['d_model'] + + (config['attn_config']['kv_n_heads'] * head_dim) + ], + axis=-1) + w_q_split = np.split(w_q, tensor_parallelism, axis=-1) + w_k_split = np.split(w_k, tensor_parallelism, axis=-1) + w_v_split = np.split(w_v, tensor_parallelism, axis=-1) + split_vals = [ + np.concatenate((w_q_split[i], w_k_split[i], w_v_split[i]), + axis=-1) for i in range(tensor_parallelism) + ] + else: + # Multi-head attention + assert data.shape == (3 * config['d_model'], config['d_model']) + # nn.Linear weights are transposed + data = data.T + data = data.reshape(config['d_model'], 3, config['d_model']) + split_vals = np.split(data, tensor_parallelism, axis=-1) for j in range(tensor_parallelism): save_path = os.path.join(out_dir, f'model.{tensor_name}.{j}.bin') split_vals[j].tofile(save_path) - if config['no_bias']: - write_zero_bias(tensor_name, save_path, - (3, split_vals[j].shape[-1]), data_type) else: raise RuntimeError(f'Tensor with name {tensor_name} is not handled') @@ -240,6 +246,12 @@ def convert_mpt_to_ft(model_name_or_path: str, raise RuntimeError( 'qk_ln is enabled for this MPT model. This may not work as expected in FT. Use --force to force a conversion.' ) + if 'kv_n_heads' in hf_config['attn_config']: + config['gpt']['n_kv_head'] = str( + hf_config['attn_config']['kv_n_heads']) + + if 'no_bias' in hf_config and hf_config['no_bias']: + config['gpt']['bias'] = str(False) with open(os.path.join(out_dir, 'config.ini'), 'w') as configfile: config.write(configfile) diff --git a/examples/mpt/run.py b/examples/mpt/run.py index ff7b88a61032..5111573cde9b 100644 --- a/examples/mpt/run.py +++ b/examples/mpt/run.py @@ -22,6 +22,7 @@ from transformers import AutoTokenizer, T5Tokenizer import tensorrt_llm +from tensorrt_llm.quantization import QuantMode from tensorrt_llm.runtime import ModelConfig, SamplingConfig from build import get_engine_name # isort:skip @@ -39,11 +40,16 @@ def read_config(config_path: Path): hidden_size = config['builder_config']['hidden_size'] // world_size vocab_size = config['builder_config']['vocab_size'] num_layers = config['builder_config']['num_layers'] - multi_query_mode = config['builder_config']['multi_query_mode'] + num_kv_heads = config['builder_config'].get('num_kv_heads', num_heads) paged_kv_cache = config['plugin_config']['paged_kv_cache'] tokens_per_block = config['plugin_config']['tokens_per_block'] - num_kv_heads = 1 if multi_query_mode else num_heads dtype = config['builder_config']['precision'] + quant_mode = QuantMode(config['builder_config']['quant_mode']) + + num_kv_heads = (num_kv_heads + world_size - 1) // world_size + assert (num_heads % world_size) == 0 + num_heads = num_heads // world_size + hidden_size = hidden_size // world_size model_config = ModelConfig(num_heads=num_heads, num_kv_heads=num_kv_heads, @@ -54,6 +60,7 @@ def read_config(config_path: Path): remove_input_padding=remove_input_padding, paged_kv_cache=paged_kv_cache, tokens_per_block=tokens_per_block, + quant_mode=quant_mode, dtype=dtype) dtype = config['builder_config']['precision'] diff --git a/examples/mpt/weight.py b/examples/mpt/weight.py index 104e1684bef9..a0c2b0155a7a 100644 --- a/examples/mpt/weight.py +++ b/examples/mpt/weight.py @@ -1,18 +1,100 @@ import configparser import time from pathlib import Path +from typing import Dict, List, Optional, Union import numpy as np import torch import tensorrt_llm -from tensorrt_llm._utils import (pad_vocab_size, str_dtype_to_np, - str_dtype_to_torch) +from tensorrt_llm._utils import pad_vocab_size, str_dtype_to_np from tensorrt_llm.functional import is_gated_activation from tensorrt_llm.models import GPTLMHeadModel +from tensorrt_llm.models.quantized.quant import get_dummy_quant_scales from tensorrt_llm.quantization import QuantMode +def get_scaling_factors( + model_path: Union[str, Path], + num_layers: int, + quant_mode: Optional[QuantMode] = None, +) -> Optional[Dict[str, List[int]]]: + """ Get the scaling factors for MPT model + + Returns a dictionary of scaling factors for the selected layers of the + MPT model. + + Args: + model_path (str): Path to the quantized MPT model + layers (list): List of layers to get the scaling factors for. If None, + all layers are selected. + + Returns: + dict: Dictionary of scaling factors for the selected layers of the + LLaMA model. + + example: + + { + 'qkv_act': qkv_act_scale, + 'qkv_weights': qkv_weights_scale, + 'qkv_output' : qkv_outputs_scale, + 'dense_act': dense_act_scale, + 'dense_weights': dense_weights_scale, + 'fc_act': fc_act_scale, + 'fc_weights': fc_weights_scale, + 'proj_act': proj_act_scale, + 'proj_weights': proj_weights_scale, + } + """ + + if model_path is None: + logger.warning(f"--quantized_fp8_model_path not specified. " + f"Initialize quantization scales automatically.") + return get_dummy_quant_scales(num_layers) + weight_dict = np.load(model_path) + + # yapf: disable + scaling_factor = { + 'qkv_act': [], + 'qkv_weights': [], + 'qkv_output': [], + 'dense_act': [], + 'dense_weights': [], + 'fc_act': [], + 'fc_weights': [], + 'proj_act': [], + 'proj_weights': [], + } + + for layer in range(num_layers): + scaling_factor['qkv_act'].append(max( + weight_dict[f'_np:layers:{layer}:attention:qkv:q:activation_scaling_factor'].item(), + weight_dict[f'_np:layers:{layer}:attention:qkv:k:activation_scaling_factor'].item(), + weight_dict[f'_np:layers:{layer}:attention:qkv:v:activation_scaling_factor'].item() + )) + scaling_factor['qkv_weights'].append(max( + weight_dict[f'_np:layers:{layer}:attention:qkv:q:weights_scaling_factor'].item(), + weight_dict[f'_np:layers:{layer}:attention:qkv:k:weights_scaling_factor'].item(), + weight_dict[f'_np:layers:{layer}:attention:qkv:v:weights_scaling_factor'].item() + )) + if quant_mode is not None and quant_mode.has_fp8_kv_cache(): + # Not calibrarting KV cache. + scaling_factor['qkv_output'].append(1.0) + scaling_factor['dense_act'].append(weight_dict[f'_np:layers:{layer}:attention:dense:activation_scaling_factor'].item()) + scaling_factor['dense_weights'].append(weight_dict[f'_np:layers:{layer}:attention:dense:weights_scaling_factor'].item()) + scaling_factor['fc_act'].append(weight_dict[f'_np:layers:{layer}:mlp:fc:activation_scaling_factor'].item()) + scaling_factor['fc_weights'].append(weight_dict[f'_np:layers:{layer}:mlp:fc:weights_scaling_factor'].item()) + scaling_factor['proj_act'].append(weight_dict[f'_np:layers:{layer}:mlp:proj:activation_scaling_factor'].item()) + scaling_factor['proj_weights'].append(weight_dict[f'_np:layers:{layer}:mlp:proj:weights_scaling_factor'].item()) + # yapf: enable + for k, v in scaling_factor.items(): + assert len(v) == num_layers, \ + f'Expect scaling factor {k} of length {num_layers}, got {len(v)}' + + return scaling_factor + + def gen_suffix(rank, use_smooth_quant, quant_per_channel): suffix = f"{rank}.bin" if use_smooth_quant: @@ -61,10 +143,15 @@ def parse_ft_config(ini_file): if inter_size is None: inter_size = 4 * n_embd + n_kv_head = gpt_config.getint('gpt', 'n_kv_head', fallback=None) multi_query_mode = gpt_config.getboolean('gpt', 'multi_query_mode', fallback=False) + assert not (multi_query_mode and n_kv_head and n_kv_head != 1), \ + "if multi_query_mode is enabled, n_kv_head must be 1 or unset" + if multi_query_mode: + n_kv_head = 1 prompt_num_tasks = gpt_config.getint('gpt', 'prompt_num_tasks', fallback=0) prompt_max_vocab_size = gpt_config.getint('gpt', 'prompt_max_vocab_size', @@ -72,7 +159,7 @@ def parse_ft_config(ini_file): pos_embedding_type = gpt_config.get('gpt', 'position_embedding_type', fallback='alibi') - return n_embd, n_head, n_layer, n_positions, vocab_size, do_layer_norm_before, hidden_act, rotary_pct, bias, inter_size, multi_query_mode, dtype, prompt_num_tasks, prompt_max_vocab_size, pos_embedding_type + return n_embd, n_head, n_layer, n_positions, vocab_size, do_layer_norm_before, hidden_act, rotary_pct, bias, inter_size, n_kv_head, dtype, prompt_num_tasks, prompt_max_vocab_size, pos_embedding_type def check_embedding_share(dir_path): @@ -90,8 +177,7 @@ def load_from_ft(tensorrt_llm_gpt: GPTLMHeadModel, dtype='float32', use_parallel_embedding=False, sharding_dim=0, - share_embedding_table=False, - scaling_factors=None): + share_embedding_table=False): tensorrt_llm.logger.info('Loading weights from FT...') tik = time.time() @@ -100,7 +186,7 @@ def load_from_ft(tensorrt_llm_gpt: GPTLMHeadModel, plugin_weight_only_quant_type = torch.int8 elif quant_mode.is_int4_weight_only(): plugin_weight_only_quant_type = torch.quint4x2 - n_embd, n_head, n_layer, n_positions, vocab_size, do_layer_norm_before, hidden_act, rotary_pct, bias, inter_size, multi_query_mode, *_ = parse_ft_config( + n_embd, n_head, n_layer, n_positions, vocab_size, do_layer_norm_before, hidden_act, rotary_pct, bias, inter_size, n_kv_head, *_ = parse_ft_config( Path(dir_path) / 'config.ini') np_dtype = str_dtype_to_np(dtype) @@ -162,9 +248,6 @@ def set_smoothquant_scale_factors(module, # Int8 KV cache use_int8_kv_cache = quant_mode.has_int8_kv_cache() - #Enable FP8 Gemm - enable_fp8_qdq = quant_mode.has_fp8_qdq() - # Debug suffix = gen_suffix(rank, use_smooth_quant, quant_per_channel) # The type of weights. @@ -218,12 +301,18 @@ def set_smoothquant_scale_factors(module, constant_values=0) tensorrt_llm_gpt.lm_head.weight.value = np.ascontiguousarray( split(lm_head_weight, tensor_parallel, rank)) - fake_fp8_sf_dt = np.float32 for i in range(n_layer): - c_attn_out_dim = (3 * n_embd // - tensor_parallel) if not multi_query_mode else ( - n_embd // tensor_parallel + - (n_embd // n_head) * 2) + head_dim = n_embd // n_head + if n_kv_head == 1: + # multi-query attention. + c_attn_out_dim = (n_embd // tensor_parallel) + (head_dim * 2) + elif n_kv_head: + # grouped-query attention. + c_attn_out_dim = (n_embd // tensor_parallel + + (head_dim * n_kv_head * 2) // tensor_parallel) + else: + # multi-head attention. + c_attn_out_dim = 3 * n_embd // tensor_parallel tensorrt_llm_gpt.layers[i].input_layernorm.weight.value = (fromfile( dir_path, 'model.layers.' + str(i) + '.input_layernorm.weight.bin')) tensorrt_llm_gpt.layers[i].input_layernorm.bias.value = (fromfile( @@ -262,19 +351,6 @@ def set_smoothquant_scale_factors(module, if t is not None: dst = tensorrt_llm_gpt.layers[i].attention.qkv.bias dst.value = np.ascontiguousarray(t) - if enable_fp8_qdq: - tensorrt_llm_gpt.layers[ - i].attention.qkv.activation_scaling_factor.value = np.array( - [scaling_factors['qkv_act'][i]], dtype=fake_fp8_sf_dt) - tensorrt_llm_gpt.layers[ - i].attention.qkv.weights_scaling_factor.value = np.array( - [scaling_factors['qkv_weights'][i]], dtype=fake_fp8_sf_dt) - tensorrt_llm_gpt.layers[ - i].attention.kv_orig_quant_scale.value = np.array( - [scaling_factors['qkv_output'][i]], dtype=np.float32) - tensorrt_llm_gpt.layers[ - i].attention.kv_quant_orig_scale.value = np.array( - [1.0 / scaling_factors['qkv_output'][i]], dtype=np.float32) dst = tensorrt_llm_gpt.layers[i].attention.dense.weight t = fromfile( @@ -307,13 +383,6 @@ def set_smoothquant_scale_factors(module, dst.value = fromfile( dir_path, 'model.layers.' + str(i) + '.attention.dense.bias.bin') - if enable_fp8_qdq: - tensorrt_llm_gpt.layers[ - i].attention.dense.activation_scaling_factor.value = np.array( - [scaling_factors['dense_act'][i]], dtype=fake_fp8_sf_dt) - tensorrt_llm_gpt.layers[ - i].attention.dense.weights_scaling_factor.value = np.array( - [scaling_factors['dense_weights'][i]], dtype=fake_fp8_sf_dt) dst = tensorrt_llm_gpt.layers[i].post_layernorm.weight dst.value = fromfile( @@ -364,13 +433,6 @@ def set_smoothquant_scale_factors(module, tensorrt_llm_gpt.layers[ i].mlp.gate.weight.value = np.ascontiguousarray( np.transpose(t, [1, 0])) - if enable_fp8_qdq: - tensorrt_llm_gpt.layers[ - i].mlp.fc.activation_scaling_factor.value = np.array( - [scaling_factors['fc_act'][i]], dtype=fake_fp8_sf_dt) - tensorrt_llm_gpt.layers[ - i].mlp.fc.weights_scaling_factor.value = np.array( - [scaling_factors['fc_weights'][i]], dtype=fake_fp8_sf_dt) t = fromfile( dir_path, @@ -413,147 +475,6 @@ def set_smoothquant_scale_factors(module, i].attention.kv_orig_quant_scale.value = 1.0 / t tensorrt_llm_gpt.layers[i].attention.kv_quant_orig_scale.value = t - if enable_fp8_qdq: - tensorrt_llm_gpt.layers[ - i].mlp.proj.activation_scaling_factor.value = np.array( - [scaling_factors['proj_act'][i]], dtype=fake_fp8_sf_dt) - tensorrt_llm_gpt.layers[ - i].mlp.proj.weights_scaling_factor.value = np.array( - [scaling_factors['proj_weights'][i]], dtype=fake_fp8_sf_dt) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - tensorrt_llm.logger.info(f'Weights loaded. Total time: {t}') - - -def load_from_hf_gpt(tensorrt_llm_gpt: GPTLMHeadModel, - hf_gpt, - rank=0, - tensor_parallel=1, - dtype='float32', - multi_query_mode=False): - tensorrt_llm.logger.info('Loading weights from HF GPT...') - tik = time.time() - - valid_lm_head_weight = False - hidden_size = tensorrt_llm_gpt._hidden_size - head_size = tensorrt_llm_gpt._num_heads // hidden_size - for k, v in hf_gpt.state_dict().items(): - torch_dtype = str_dtype_to_torch(dtype) - v = v.to(torch_dtype).cpu().numpy() - if 'wte.weight' in k: - tensorrt_llm_gpt.embedding.vocab_embedding.weight.value = v - elif 'wpe.weight' in k: - tensorrt_llm_gpt.embedding.position_embedding.weight.value = v - elif 'ln_f.weight' in k: - tensorrt_llm_gpt.ln_f.weight.value = v - elif 'ln_f.bias' in k: - tensorrt_llm_gpt.ln_f.bias.value = v - elif 'lm_head.weight' in k: - tensorrt_llm_gpt.lm_head.weight.value = np.ascontiguousarray( - split(v, tensor_parallel, rank)) - valid_lm_head_weight = True - else: - layer_idx = extract_layer_idx(k) - if layer_idx is None: - continue - idx = int(layer_idx) - if 'ln_1.weight' in k: - tensorrt_llm_gpt.layers[idx].input_layernorm.weight.value = v - elif 'ln_1.bias' in k: - tensorrt_llm_gpt.layers[idx].input_layernorm.bias.value = v - elif 'attn.c_attn.weight' in k: - if multi_query_mode: - # HF-StarCoder uses torch.nn.Linear - w_qkv = v.reshape(hidden_size + 2 * head_size, 3, - hidden_size) - w_q, w_kv = np.split(w_qkv, [hidden_size, 2 * head_size]) - w_q = split(w_q, tensor_parallel, rank) - dst = tensorrt_llm_gpt.layers[idx].attention.qkv.weight - dst.value = np.ascontiguousarray(np.concatenate(w_q, w_kv)) - else: - # HF-GPT uses Conv1D instead of Linear - v = v.transpose() - dst = tensorrt_llm_gpt.layers[idx].attention.qkv.weight - dst.value = np.ascontiguousarray( - split(v, tensor_parallel, rank)) - elif 'attn.c_attn.bias' in k: - if multi_query_mode: - v.reshape(hidden_size + 2 * head_size, 3) - bias_q, bias_kv = np.split(w_qkv, - [hidden_size, 2 * head_size]) - bias_q = split(bias_q, tensor_parallel, rank) - dst = tensorrt_llm_gpt.layers[idx].attention.qkv.bias - dst.value = np.ascontiguousarray( - np.concatenate(bias_q, bias_kv)) - else: - dst = tensorrt_llm_gpt.layers[idx].attention.qkv.bias - dst.value = np.ascontiguousarray( - split(v, tensor_parallel, rank)) - elif 'attn.q_attn.weight' in k: - # Get the corresponding kv_atten.weight: - # ex: transformer.h.23.attn.kv_attn.weight - u = hf_gpt.state_dict()[k.replace('q_attn', 'kv_attn')] - u = u.to(torch_dtype).cpu().numpy(force=True) - # HF-SantaCoder uses transformer.Conv1D so we transpose to match shape - # In addition, kv_head must be broadcasted to all ranks so split is not applied - v = split(v.transpose(), tensor_parallel, rank) # W_q - u = u.transpose() # W_kv - dst = tensorrt_llm_gpt.layers[idx].attention.qkv.weight - dst.value = np.ascontiguousarray(np.concatenate((v, u))) - elif 'attn.q_attn.bias' in k: - # Get the corresponding kv_atten.bias: - # ex: transformer.h.23.attn.kv_attn.bias - u = hf_gpt.state_dict()[k.replace('q_attn', 'kv_attn')] - u = u.to(torch_dtype).cpu().numpy(force=True) - v = split(v, tensor_parallel, rank) - dst = tensorrt_llm_gpt.layers[idx].attention.qkv.bias - dst.value = np.ascontiguousarray(np.concatenate((v, u))) - elif 'attn.c_proj.weight' in k: - v = v.transpose() - dst = tensorrt_llm_gpt.layers[idx].attention.dense.weight - dst.value = np.ascontiguousarray( - split(v, tensor_parallel, rank, dim=1)) - elif 'attn.c_proj.bias' in k: - dst = tensorrt_llm_gpt.layers[idx].attention.dense.bias - dst.value = v - elif 'ln_2.weight' in k: - dst = tensorrt_llm_gpt.layers[idx].post_layernorm.weight - dst.value = v - elif 'ln_2.bias' in k: - dst = tensorrt_llm_gpt.layers[idx].post_layernorm.bias - dst.value = v - elif 'mlp.c_fc.weight' in k: - v = v.transpose() - tensorrt_llm_gpt.layers[ - idx].mlp.fc.weight.value = np.ascontiguousarray( - split(v, tensor_parallel, rank)) - elif 'mlp.c_fc.bias' in k: - tensorrt_llm_gpt.layers[ - idx].mlp.fc.bias.value = np.ascontiguousarray( - split(v, tensor_parallel, rank)) - elif 'mlp.c_proj.weight' in k: - v = v.transpose() - tensorrt_llm_gpt.layers[ - idx].mlp.proj.weight.value = np.ascontiguousarray( - split(v, tensor_parallel, rank, dim=1)) - elif 'mlp.c_proj.bias' in k: - tensorrt_llm_gpt.layers[idx].mlp.proj.bias.value = v - - if not valid_lm_head_weight: - # Use wte as lm_head weight to match the load_from_ft implementation. - lm_head_weight = tensorrt_llm_gpt.embedding.vocab_embedding.weight._value - vocab_size = hf_gpt.config.vocab_size - if vocab_size % tensor_parallel != 0: - # padding - vocab_size_padded = tensorrt_llm_gpt.lm_head.out_features * tensor_parallel - pad_width = vocab_size_padded - vocab_size - lm_head_weight = np.pad(lm_head_weight, ((0, pad_width), (0, 0)), - 'constant', - constant_values=0) - tensorrt_llm_gpt.lm_head.weight.value = np.ascontiguousarray( - split(lm_head_weight, tensor_parallel, rank)) - tok = time.time() t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) tensorrt_llm.logger.info(f'Weights loaded. Total time: {t}') diff --git a/examples/opt/README.md b/examples/opt/README.md index c5a39cc768c2..c2aab59776eb 100644 --- a/examples/opt/README.md +++ b/examples/opt/README.md @@ -5,13 +5,12 @@ multiple GPUs or multiple nodes with multiple GPUs. ## Overview -The TensorRT-LLM OPT implementation can be found in [`tensorrt_llm/models/opt/model.py`](../../tensorrt_llm/models/opt/model.py). The TensorRT-LLM OPT example -code is located in [`examples/opt`](./). There are four main files in that folder: +The TensorRT-LLM OPT implementation can be found in [`tensorrt_llm/models/opt/model.py`](../../tensorrt_llm/models/opt/model.py). The TensorRT-LLM OPT example code is located in [`examples/opt`](./). There are four main files: * [`hf_opt_convert.py`](./hf_opt_convert.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the [FasterTransformer (FT)](https://github.com/NVIDIA/FasterTransformer) format, * [`build.py`](./build.py) to build the [TensorRT](https://developer.nvidia.com/tensorrt) engine(s) needed to run the OPT model, - * [`summarize.py`](./summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/cnn_dailymail) dataset using the model. + * and a shared [`../summarize.py`](../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/cnn_dailymail) dataset using the model. ## Support Matrix * FP16 @@ -151,44 +150,48 @@ The script can also perform the same summarization using the HF OPT model. ```bash # OPT-125M -python3 summarize.py --engine_dir trt_engine/opt-125m/fp16/1-gpu \ - --test_hf \ - --batch_size 1 \ - --test_trt_llm \ - --hf_model_location opt-125m \ - --data_type fp16 \ - --check_accuracy \ - --tensorrt_llm_rouge1_threshold=14 +python3 ../summarize.py --engine_dir trt_engine/opt-125m/fp16/1-gpu \ + --test_hf \ + --batch_size 1 \ + --test_trt_llm \ + --hf_model_dir opt-125m \ + --data_type fp16 \ + --check_accuracy \ + --tensorrt_llm_rouge1_threshold=14 \ + --no_add_special_tokens # OPT-350M -python3 summarize.py --engine_dir trt_engine/opt-350m/fp16/1-gpu \ - --test_hf \ - --batch_size 1 \ - --test_trt_llm \ - --hf_model_location opt-350m \ - --data_type fp16 \ - --check_accuracy \ - --tensorrt_llm_rouge1_threshold=20 +python3 ../summarize.py --engine_dir trt_engine/opt-350m/fp16/1-gpu \ + --test_hf \ + --batch_size 1 \ + --test_trt_llm \ + --hf_model_dir opt-350m \ + --data_type fp16 \ + --check_accuracy \ + --tensorrt_llm_rouge1_threshold=20 \ + --no_add_special_tokens # OPT-2.7B -python3 summarize.py --engine_dir trt_engine/opt-2.7b/fp16/1-gpu \ - --test_hf \ - --batch_size 1 \ - --test_trt_llm \ - --hf_model_location opt-2.7b \ - --data_type fp16 \ - --check_accuracy \ - --tensorrt_llm_rouge1_threshold=21 +python3 ../summarize.py --engine_dir trt_engine/opt-2.7b/fp16/1-gpu \ + --test_hf \ + --batch_size 1 \ + --test_trt_llm \ + --hf_model_dir opt-2.7b \ + --data_type fp16 \ + --check_accuracy \ + --tensorrt_llm_rouge1_threshold=21 \ + --no_add_special_tokens # OPT-66B mpirun -n 4 --allow-run-as-root \ - python3 summarize.py --engine_dir trt_engines/opt-66b/fp16/4-gpu \ - --batch_size 1 \ - --test_trt_llm \ - --hf_model_location opt-66b \ - --data_type fp16 \ - --check_accuracy \ - --tensorrt_llm_rouge1_threshold=21 + python3 ../summarize.py --engine_dir trt_engines/opt-66b/fp16/4-gpu \ + --batch_size 1 \ + --test_trt_llm \ + --hf_model_dir opt-66b \ + --data_type fp16 \ + --check_accuracy \ + --tensorrt_llm_rouge1_threshold=21 \ + --no_add_special_tokens ``` #### Fused MultiHead Attention (FMHA) diff --git a/examples/opt/build.py b/examples/opt/build.py index 8093617986de..b5ec19ca5472 100644 --- a/examples/opt/build.py +++ b/examples/opt/build.py @@ -361,9 +361,11 @@ def build(rank, args): hidden_act=args.hidden_act, max_position_embeddings=args.n_positions, max_batch_size=args.max_batch_size, + max_beam_width=args.max_beam_width, max_input_len=args.max_input_len, max_output_len=args.max_output_len, - use_prompt_tuning=args.max_prompt_embedding_table_size > 0, + max_prompt_embedding_table_size=args. + max_prompt_embedding_table_size, int8=(args.quant_mode.has_act_or_weight_quant() or args.quant_mode.has_int8_kv_cache()), strongly_typed=args.strongly_typed) diff --git a/examples/opt/requirements.txt b/examples/opt/requirements.txt index f46bff310071..fb205c8dd536 100644 --- a/examples/opt/requirements.txt +++ b/examples/opt/requirements.txt @@ -1,2 +1,3 @@ datasets~=2.14.5 +evaluate~=0.4.1 rouge_score~=0.1.2 diff --git a/examples/opt/summarize.py b/examples/opt/summarize.py deleted file mode 100644 index 6f81a269d75f..000000000000 --- a/examples/opt/summarize.py +++ /dev/null @@ -1,377 +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 copy -import json -import os - -import numpy as np -import torch -from datasets import load_dataset, load_metric -from transformers import AutoModelForCausalLM, AutoTokenizer - -import tensorrt_llm -import tensorrt_llm.profiler as profiler -from tensorrt_llm.logger import logger - -from build import get_engine_name # isort:skip - - -def TRTOPT(args, config): - dtype = config['builder_config']['precision'] - world_size = config['builder_config']['tensor_parallel'] - assert world_size == tensorrt_llm.mpi_world_size(), \ - f'Engine world size ({world_size}) != Runtime world size ({tensorrt_llm.mpi_world_size()})' - - use_gpt_attention_plugin = bool( - config['plugin_config']['gpt_attention_plugin']) - world_size = config['builder_config']['tensor_parallel'] - num_heads = config['builder_config']['num_heads'] // world_size - hidden_size = config['builder_config']['hidden_size'] // world_size - vocab_size = config['builder_config']['vocab_size'] - num_layers = config['builder_config']['num_layers'] - remove_input_padding = config['plugin_config']['remove_input_padding'] - - model_config = tensorrt_llm.runtime.ModelConfig( - vocab_size=vocab_size, - num_layers=num_layers, - num_heads=num_heads, - num_kv_heads=num_heads, - hidden_size=hidden_size, - gpt_attention_plugin=use_gpt_attention_plugin, - remove_input_padding=remove_input_padding, - dtype=dtype) - - runtime_rank = tensorrt_llm.mpi_rank() - runtime_mapping = tensorrt_llm.Mapping(world_size, - runtime_rank, - tp_size=world_size) - torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node) - - engine_name = get_engine_name('opt', dtype, world_size, runtime_rank) - serialize_path = os.path.join(args.engine_dir, engine_name) - - tensorrt_llm.logger.set_level(args.log_level) - - with open(serialize_path, 'rb') as f: - engine_buffer = f.read() - decoder = tensorrt_llm.runtime.GenerationSession(model_config, - engine_buffer, - runtime_mapping) - - return decoder - - -def main(args): - 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 - hf_model_location = args.hf_model_location - tokenizer = AutoTokenizer.from_pretrained(hf_model_location, - padding_side='left') - tokenizer.pad_token = tokenizer.eos_token - - dataset_cnn = load_dataset("ccdv/cnn_dailymail", - '3.0.0', - cache_dir=args.dataset_path) - - config_path = os.path.join(args.engine_dir, 'config.json') - with open(config_path, 'r') as f: - config = json.load(f) - - max_batch_size = args.batch_size - - # runtime parameters - # repetition_penalty = 1 - top_k = args.top_k - output_len = 100 - test_token_num = 923 - # top_p = 0.0 - # random_seed = 5 - temperature = 1 - num_beams = args.num_beams - - # model hyper-parameters - pad_id = tokenizer.encode(tokenizer.pad_token, add_special_tokens=False)[0] - end_id = tokenizer.encode(tokenizer.eos_token, add_special_tokens=False)[0] - - if test_trt_llm: - tensorrt_llm_gpt = TRTOPT(args, config) - - if test_hf: - model = AutoModelForCausalLM.from_pretrained(hf_model_location) - model.cuda() - if args.data_type == 'fp16': - model.half() - - def summarize_tensorrt_llm(datapoint): - batch_size = len(datapoint['article']) - - line = copy.copy(datapoint['article']) - line_encoded = [] - input_lengths = [] - for i in range(batch_size): - line[i] = line[i] + ' TL;DR: ' - - line[i] = line[i].strip() - line[i] = line[i].replace(" n't", "n't") - - input_id = tokenizer.encode(line[i], - return_tensors='pt', - add_special_tokens=False).type( - torch.int32) - input_id = input_id[:, -test_token_num:] - - line_encoded.append(input_id) - input_lengths.append(input_id.shape[-1]) - - # do padding, should move outside the profiling to prevent the overhead - max_length = max(input_lengths) - if tensorrt_llm_gpt.remove_input_padding: - line_encoded = [ - torch.tensor(t, dtype=torch.int32).cuda() for t in line_encoded - ] - else: - # do padding, should move outside the profiling to prevent the overhead - for i in range(batch_size): - pad_size = max_length - input_lengths[i] - - pad = torch.ones([1, pad_size]).type(torch.int32) * pad_id - line_encoded[i] = torch.cat( - [torch.tensor(line_encoded[i], dtype=torch.int32), pad], - axis=-1) - - line_encoded = torch.cat(line_encoded, axis=0).cuda() - input_lengths = torch.tensor(input_lengths, - dtype=torch.int32).cuda() - - sampling_config = tensorrt_llm.runtime.SamplingConfig( - end_id=end_id, pad_id=pad_id, top_k=top_k, num_beams=num_beams) - - with torch.no_grad(): - tensorrt_llm_gpt.setup(batch_size, - max_context_length=max_length, - max_new_tokens=output_len, - beam_width=num_beams, - max_kv_cache_length=args.max_kv_cache_len) - - if tensorrt_llm_gpt.remove_input_padding: - output_ids = tensorrt_llm_gpt.decode_batch( - line_encoded, sampling_config) - else: - output_ids = tensorrt_llm_gpt.decode( - line_encoded, - input_lengths, - sampling_config, - ) - - torch.cuda.synchronize() - - # Extract a list of tensors of shape beam_width x output_ids. - if tensorrt_llm_gpt.mapping.is_first_pp_rank(): - output_beams_list = [ - tokenizer.batch_decode(output_ids[batch_idx, :, - input_lengths[batch_idx]:], - skip_special_tokens=True) - for batch_idx in range(batch_size) - ] - return output_beams_list, output_ids[:, :, max_length:].tolist() - return [], [] - - def summarize_hf(datapoint): - batch_size = len(datapoint['article']) - 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}" - ) - - line = copy.copy(datapoint['article']) - for i in range(batch_size): - line[i] = line[i] + ' TL;DR: ' - - line[i] = line[i].strip() - line[i] = line[i].replace(" n't", "n't") - - line_encoded = tokenizer(line, - return_tensors='pt', - padding=True, - truncation=True)["input_ids"].type(torch.int64) - - line_encoded = line_encoded[:, -test_token_num:] - line_encoded = line_encoded.cuda() - - with torch.no_grad(): - output = model.generate(line_encoded, - max_length=len(line_encoded[0]) + - output_len, - top_k=top_k, - temperature=temperature, - eos_token_id=tokenizer.eos_token_id, - pad_token_id=tokenizer.pad_token_id, - num_beams=num_beams, - num_return_sequences=num_beams, - early_stopping=True) - - tokens_list = output[:, len(line_encoded[0]):].tolist() - output = output.reshape([batch_size, num_beams, -1]) - output_lines_list = [ - tokenizer.batch_decode(output[:, i, len(line_encoded[0]):], - skip_special_tokens=True) - for i in range(num_beams) - ] - - return output_lines_list, tokens_list - - if test_trt_llm: - datapoint = dataset_cnn['test'][0:1] - summary, _ = summarize_tensorrt_llm(datapoint) - if runtime_rank == 0: - logger.info( - "---------------------------------------------------------") - logger.info("TensorRT-LLM Generated : ") - logger.info(f" Article : {datapoint['article']}") - logger.info(f"\n Highlights : {datapoint['highlights']}") - logger.info(f"\n Summary : {summary}") - logger.info( - "---------------------------------------------------------") - - if test_hf: - datapoint = dataset_cnn['test'][0:1] - summary, _ = summarize_hf(datapoint) - logger.info("---------------------------------------------------------") - logger.info("HF Generated : ") - logger.info(f" Article : {datapoint['article']}") - logger.info(f"\n Highlights : {datapoint['highlights']}") - logger.info(f"\n Summary : {summary}") - logger.info("---------------------------------------------------------") - - metric_tensorrt_llm = [load_metric("rouge") for _ in range(num_beams)] - metric_hf = [load_metric("rouge") for _ in range(num_beams)] - for i in range(num_beams): - metric_tensorrt_llm[i].seed = 0 - metric_hf[i].seed = 0 - - ite_count = 0 - data_point_idx = 0 - while (data_point_idx < len(dataset_cnn['test'])) 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_cnn['test'][data_point_idx:(data_point_idx + - max_batch_size)] - - if test_trt_llm: - profiler.start('tensorrt_llm') - summary_tensorrt_llm, tokens_tensorrt_llm = summarize_tensorrt_llm( - datapoint) - profiler.stop('tensorrt_llm') - - if test_hf: - profiler.start('hf') - summary_hf, tokens_hf = summarize_hf(datapoint) - profiler.stop('hf') - - if runtime_rank == 0: - if test_trt_llm: - for batch_idx in range(len(summary_tensorrt_llm)): - for beam_idx in range(num_beams): - metric_tensorrt_llm[beam_idx].add_batch( - predictions=[ - summary_tensorrt_llm[batch_idx][beam_idx] - ], - references=[datapoint['highlights'][batch_idx]]) - if test_hf: - for beam_idx in range(num_beams): - for batch_idx in range(len(summary_hf[beam_idx])): - metric_hf[beam_idx].add_batch( - predictions=[summary_hf[beam_idx][batch_idx]], - references=[datapoint['highlights'][batch_idx]]) - - logger.debug('-' * 100) - logger.debug(f"Article : {datapoint['article']}") - if test_trt_llm: - logger.debug(f'TensorRT-LLM Summary: {summary_tensorrt_llm}') - if test_hf: - logger.debug(f'HF Summary: {summary_hf}') - logger.debug(f"highlights : {datapoint['highlights']}") - - data_point_idx += max_batch_size - ite_count += 1 - - if runtime_rank == 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)' - ) - for beam_idx in range(num_beams): - logger.info(f"TensorRT-LLM beam {beam_idx} result") - 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].mid[2]*100}' - ) - - if args.check_accuracy and beam_idx == 0: - assert computed_metrics_tensorrt_llm['rouge1'].mid[ - 2] * 100 > args.tensorrt_llm_rouge1_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)' - ) - for beam_idx in range(num_beams): - logger.info(f"HF beam {beam_idx} result") - computed_metrics_hf = metric_hf[beam_idx].compute() - for key in computed_metrics_hf.keys(): - logger.info( - f' {key} : {computed_metrics_hf[key].mid[2]*100}') - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--hf_model_location', type=str, default='opt-350m') - parser.add_argument('--test_hf', action='store_true') - parser.add_argument('--test_trt_llm', action='store_true') - parser.add_argument('--data_type', - type=str, - choices=['fp32', 'fp16'], - default='fp32') - parser.add_argument('--dataset_path', type=str, default="") - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument('--engine_dir', type=str, default='gpt_outputs') - parser.add_argument('--batch_size', type=int, default=1) - parser.add_argument('--max_ite', type=int, default=20) - parser.add_argument('--max_kv_cache_len', - type=int, - default=None, - help='The max kv cache length. \ - If the final sequence length exceeds the kv cache length, we will enable cyclic kv cache. \ - If it is set to None, we will use the max sequence length.') - parser.add_argument('--check_accuracy', action='store_true') - parser.add_argument('--tensorrt_llm_rouge1_threshold', - type=float, - default=15.0) - parser.add_argument('--num_beams', type=int, default=1) - parser.add_argument('--top_k', type=int, default=1) - - args = parser.parse_args() - - main(args) diff --git a/examples/quantization/quantize.py b/examples/quantization/quantize.py new file mode 100644 index 000000000000..2affbb502e95 --- /dev/null +++ b/examples/quantization/quantize.py @@ -0,0 +1,143 @@ +# 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. +""" +Adapted from examples/quantization/hf_ptq.py +""" + +import argparse +import random + +import numpy as np +import torch +from datasets import load_dataset +from torch.utils.data import DataLoader +from transformers import AutoModelForCausalLM, AutoTokenizer + +from tensorrt_llm._utils import str_dtype_to_torch +from tensorrt_llm.logger import logger +from tensorrt_llm.models.quantized.ammo import quantize_and_export + + +def get_calib_dataloader(data="cnn_dailymail", + tokenizer=None, + batch_size=1, + calib_size=16, + block_size=512, + cache_dir=None): + print("Loading calibration dataset") + if data == "pileval": + dataset = load_dataset( + "json", + data_files="https://the-eye.eu/public/AI/pile/val.jsonl.zst", + split="train", + cache_dir=cache_dir) + dataset = dataset["text"][:calib_size] + elif data == "cnn_dailymail": + dataset = load_dataset("cnn_dailymail", + name="3.0.0", + split="train", + cache_dir=cache_dir) + dataset = dataset["article"][:calib_size] + else: + raise NotImplementedError + + batch_encoded = tokenizer.batch_encode_plus(dataset, + return_tensors="pt", + padding=True, + max_length=block_size) + batch_encoded = batch_encoded["input_ids"] + batch_encoded = batch_encoded.cuda() + + calib_dataloader = DataLoader(batch_encoded, + batch_size=batch_size, + shuffle=False) + + return calib_dataloader + + +def get_tokenizer(ckpt_path, **kwargs): + logger.info(f"Loading tokenizer from {ckpt_path}") + tokenizer = AutoTokenizer.from_pretrained(ckpt_path, + padding_side="left", + **kwargs) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + return tokenizer + + +def get_model(ckpt_path, dtype="float16", cache_dir=None): + logger.info(f"Loading model from {ckpt_path}") + torch_dtype = str_dtype_to_torch(dtype) + model = AutoModelForCausalLM.from_pretrained( + ckpt_path, + device_map="auto", + trust_remote_code=True, + torch_dtype=torch_dtype, + ) + model.eval() + model = model.to(memory_format=torch.channels_last) + return model + + +def get_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model_dir", + type=str, + required=True, + help="Directory of a HF model checkpoint") + parser.add_argument("--dtype", help="Model data type.", default="float16") + parser.add_argument("--qformat", + type=str, + choices=['fp8', 'int8_sq', 'int4_awq'], + default='fp8', + help='Quantization format.') + parser.add_argument("--calib_size", + type=int, + default=128, + help="Number of samples for calibration.") + parser.add_argument("--export_path", default="exported_model") + parser.add_argument("--cache_dir", + type=str, + default=None, + help="Directory of dataset cache.") + parser.add_argument('--seed', type=int, default=None, help='Random seed') + args = parser.parse_args() + return args + + +def main(): + if not torch.cuda.is_available(): + raise EnvironmentError("GPU is required for inference.") + + args = get_args() + + if args.seed is not None: + random.seed(args.seed) + np.random.seed(args.seed) + + tokenizer = get_tokenizer(args.model_dir, cache_dir=args.cache_dir) + model = get_model(args.model_dir, args.dtype, cache_dir=args.cache_dir) + + calib_dataloader = get_calib_dataloader(tokenizer=tokenizer, + calib_size=args.calib_size, + cache_dir=args.cache_dir) + model = quantize_and_export(model, + qformat=args.qformat, + calib_dataloader=calib_dataloader, + export_path=args.export_path) + + +if __name__ == "__main__": + main() diff --git a/examples/qwen/README.md b/examples/qwen/README.md new file mode 100644 index 000000000000..3f9f3baf69c0 --- /dev/null +++ b/examples/qwen/README.md @@ -0,0 +1,366 @@ +# Qwen + +This document shows how to build and run a Qwen model in TensorRT-LLM on both single GPU, single node multi-GPU and multi-node multi-GPU. + +## Overview + +The TensorRT-LLM Qwen implementation can be found in [model.py](model.py). The TensorRT-LLM Qwen example code is located in [`examples/qwen`](./). There are three main files in that folder:: + + * [`build.py`](./build.py) to build the [TensorRT](https://developer.nvidia.com/tensorrt) engine(s) needed to run the Qwen model, + * [`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/cnn_dailymail) dataset using the model. + +## Support Matrix + * FP16 + * INT8 & INT4 Weight-Only + * SmoothQuant + * INT8 KV CACHE + * Tensor Parallel + * STRONGLY TYPED + +## Usage + +The TensorRT-LLM Qwen example code locates at [examples/qwen](./). 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) + +Need to prepare the HF Qwen checkpoint first by following the guides here [Qwen-7B-Chat](https://huggingface.co/Qwen/Qwen-7B-Chat) or [Qwen-14B-Chat](https://huggingface.co/Qwen/Qwen-14B-Chat) + +Create a `tmp/Qwen` directory to store the weights downloaded from huaggingface. +```bash +mkdir -p ./tmp/Qwen +``` + +Store Qwen-7B-Chat or Qwen-14B-Chat separately. +- for Qwen-7B-Chat +```bash +mv Qwen-7B-Chat ./tmp/Qwen/7B +``` +- for Qwen-14B-Chat +``` +mv Qwen-14B-Chat ./tmp/Qwen/14B +``` + +TensorRT-LLM Qwen builds TensorRT engine(s) from HF checkpoint. If no checkpoint directory is specified, TensorRT-LLM will build engine(s) with dummy weights. + +Normally `build.py` only requires single GPU, but if you've already got all the GPUs needed while inferencing, you could enable parallelly building to make the engine building process faster by adding `--parallel_build` argument. Please note that currently `parallel_build` feature only supports single node. + +Here're some examples: + +```bash +# Build a single-GPU float16 engine from HF weights. +# use_gpt_attention_plugin is necessary in Qwen. +# Try use_gemm_plugin to prevent accuracy issue. +# It is recommend to use --remove_input_padding along with --use_gpt_attention_plugin for better performance + +# Build the Qwen 7B model using a single GPU and FP16. +python build.py --hf_model_dir ./tmp/Qwen/7B/ \ + --dtype float16 \ + --remove_input_padding \ + --use_gpt_attention_plugin float16 \ + --enable_context_fmha \ + --use_gemm_plugin float16 \ + --output_dir ./tmp/Qwen/7B/trt_engines/fp16/1-gpu/ + +# Build the Qwen 7B model using a single GPU and BF16. +python build.py --hf_model_dir ./tmp/Qwen/7B/ \ + --dtype bfloat16 \ + --remove_input_padding \ + --use_gpt_attention_plugin bfloat16 \ + --enable_context_fmha \ + --use_gemm_plugin bfloat16 \ + --output_dir ./tmp/Qwen/7B/trt_engines/bf16/1-gpu/ + +# Build the Qwen 7B model using a single GPU and apply INT8 weight-only quantization. +python build.py --hf_model_dir ./tmp/Qwen/7B/ \ + --dtype float16 \ + --remove_input_padding \ + --use_gpt_attention_plugin float16 \ + --use_gemm_plugin float16 \ + --use_weight_only \ + --weight_only_precision int8 \ + --output_dir ./tmp/Qwen/7B/trt_engines/int8_weight_only/1-gpu/ + +# Build the Qwen 7B model using a single GPU and apply INT4 weight-only quantization. +python build.py --hf_model_dir ./tmp/Qwen/7B/ \ + --dtype float16 \ + --remove_input_padding \ + --use_gpt_attention_plugin float16 \ + --use_gemm_plugin float16 \ + --use_weight_only \ + --weight_only_precision int4 \ + --output_dir ./tmp/Qwen/7B/trt_engines/int4_weight_only/1-gpu/ + +# Build Qwen 7B using 2-way tensor parallelism. +python build.py --hf_model_dir ./tmp/Qwen/7B/ \ + --dtype float16 \ + --remove_input_padding \ + --use_gpt_attention_plugin float16 \ + --enable_context_fmha \ + --use_gemm_plugin float16 \ + --output_dir ./tmp/Qwen/7B/trt_engines/fp16/2-gpu/ \ + --world_size 2 \ + --tp_size 2 + +# Build Qwen 7B using 2-way tensor parallelism and 2-way pipeline parallelism. +python build.py --hf_model_dir ./tmp/Qwen/7B/ \ + --dtype float16 \ + --remove_input_padding \ + --use_gpt_attention_plugin float16 \ + --enable_context_fmha \ + --use_gemm_plugin float16 \ + --output_dir ./tmp/Qwen/7B/trt_engines/fp16/2-gpu/ \ + --world_size 4 \ + --tp_size 2 \ + --pp_size 2 + +# Build Qwen 14B using 2-way tensor parallelism. +python build.py --hf_model_dir ./tmp/Qwen/14B \ + --dtype float16 \ + --remove_input_padding \ + --use_gpt_attention_plugin float16 \ + --enable_context_fmha \ + --use_gemm_plugin float16 \ + --output_dir ./tmp/Qwen/14B/trt_engines/fp16/2-gpu/ \ + --world_size 2 \ + --tp_size 2 +``` +**Demo output of engine building:** +```python +python3 build.py --hf_model_dir /llm-models/Qwen-7B-Chat/ --output_dir /engine_qwen +``` +``` +[11/09/2023-00:57:06] [TRT-LLM] [I] Serially build TensorRT engines. +[11/09/2023-00:57:06] [TRT] [I] [MemUsageChange] Init CUDA: CPU +14, GPU +0, now: CPU 118, GPU 427 (MiB) +[11/09/2023-00:57:08] [TRT] [I] [MemUsageChange] Init builder kernel library: CPU +1974, GPU +350, now: CPU 2227, GPU 777 (MiB) +[11/09/2023-00:57:08] [TRT-LLM] [W] Invalid timing cache, using freshly created one +[11/09/2023-00:57:14] [TRT-LLM] [I] Loading HF QWen ... from /llm-models/Qwen-7B-Chat/ +...... +[11/09/2023-01:01:34] [TRT] [I] [MemUsageStats] Peak memory usage during Engine building and serialization: CPU: 47322 MiB +[11/09/2023-01:01:34] [TRT-LLM] [I] Total time of building qwen_float16_tp1_rank0.engine: 00:03:44 +[11/09/2023-01:01:34] [TRT-LLM] [I] Config saved to /engine_qwen/config.json. +[11/09/2023-01:01:34] [TRT-LLM] [I] Serializing engine to /engine_qwen/qwen_float16_tp1_rank0.engine... +[11/09/2023-01:01:49] [TRT-LLM] [I] Engine serialized. Total time: 00:00:14 +[11/09/2023-01:01:49] [TRT-LLM] [I] Timing cache serialized to /engine_qwen/model.cache +[11/09/2023-01:01:50] [TRT-LLM] [I] Total time of building all 1 engines: 00:04:43 +``` + + +#### INT8 weight only + INT8 KV cache +For INT8 KV cache, [`hf_qwen_convert.py`](./hf_qwen_convert.py) features a +`--calibrate-kv-cache, -kv` option. Setting `-kv` will calibrate the model, +and then export the scaling factors needed for INT8 KV cache inference. + + +Example: + +```bash +python3 hf_qwen_convert.py \ + -i ./tmp/Qwen/7B/ \ + -o ./tmp/Qwen/7B/int8_kv_cache/ \ + --calibrate-kv-cache -t float16 +``` + +[`build.py`](./build.py) add new options for the support of INT8 KV cache. + +`--int8_kv_cache` is the command-line option to enable INT8 KV cache. + +In addition, it could be combined with INT8 weight-only quantization, as follows: + +Examples of INT8 weight-only quantization + INT8 KV cache + +```bash +# Build model with both INT8 weight-only and INT8 KV cache enabled +python build.py --ft_dir_path ./tmp/Qwen/7B/int8_kv_cache/1-gpu/ \ + --dtype float16 \ + --hf_model_dir ./tmp/Qwen/7B \ + --use_gpt_attention_plugin float16 \ + --use_gemm_plugin float16 \ + --output_dir ./tmp/Qwen/7B/trt_engines/int8_kv_cache_weight_only/1-gpu \ + --int8_kv_cache \ + --use_weight_only +``` + +- run +```bash +python3 run.py --max_new_tokens=50 \ + --tokenizer_dir ./tmp/Qwen/7B/ \ + --engine_dir=./tmp/Qwen/7B/trt_engines/int8_kv_cache_weight_only/1-gpu +``` + +Test with `summarize.py`: + + +- validate huggingface +```bash +python3 summarize.py --backend=hf \ + --tokenizer_dir ./tmp/Qwen/7B \ + --hf_model_dir ./tmp/Qwen/7B +``` + +- validate trt-llm +```bash +python3 summarize.py --backend=trt_llm \ + --tokenizer_dir ./tmp/Qwen/7B \ + --engine_dir ./tmp/Qwen/7B/trt_engines/int8_kv_cache_weight_only/1-gpu +``` + +#### SmoothQuant + +The smoothquant supports both Qwen v1 and Qwen v2. 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 +python3 hf_qwen_convert.py -i ./tmp/Qwen/7B -o ./tmp/Qwen/7B/sq0.5/ -sq 0.5 --tensor-parallelism 1 --storage-type float16 +``` + +[`build.py`](./build.py) add new options for the support of INT8 inference of SmoothQuant models. + +`--use_smooth_quant` 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_tensor_ mode. +python3 build.py --ft_dir_path=./tmp/Qwen/7B/sq0.5/1-gpu/ \ + --use_smooth_quant \ + --hf_model_dir ./tmp/Qwen/7B \ + --output_dir ./tmp/Qwen/7B/trt_engines/sq0.5/1-gpu/ + +# Build model for SmoothQuant in the _per_token_ + _per_channel_ mode +python3 build.py --ft_dir_path=./tmp/Qwen/7B/sq0.5/1-gpu/ \ + --use_smooth_quant \ + --per_token \ + --per_channel \ + --hf_model_dir ./tmp/Qwen/7B \ + --output_dir ./tmp/Qwen/7B/trt_engines/sq0.5/1-gpu/ +``` + +- run +```bash +python3 run.py --max_new_tokens=50 \ + --tokenizer_dir ./tmp/Qwen/7B/ \ + --engine_dir=./tmp/Qwen/7B/trt_engines/sq0.5/1-gpu/ +``` + +- summarize +```bash +python summarize.py --backend=trt_llm \ + --tokenizer_dir ./tmp/Qwen/7B/ \ + --data_type fp16 \ + --engine_dir=./tmp/Qwen/7B/trt_engines/sq0.5/1-gpu/ +``` + + +### Run + +To run a TensorRT-LLM Qwen model using the engines generated by build.py + +```bash +# With fp16 inference +python3 run.py --max_new_tokens=50 \ + --tokenizer_dir ./tmp/Qwen/7B/ \ + --engine_dir=./tmp/Qwen/7B/trt_engines/fp16/1-gpu/ + +# With bf16 inference +python3 run.py --max_new_tokens=50 \ + --tokenizer_dir ./tmp/Qwen/7B/ \ + --engine_dir=./tmp/Qwen/7B/trt_engines/bf16/1-gpu + +# With int8 weight only inference +python3 run.py --max_new_tokens=50 \ + --tokenizer_dir ./tmp/Qwen/7B/ \ + --engine_dir=./tmp/Qwen/7B/trt_engines/int8_weight_only/1-gpu/ + +# With int4 weight only inference +python3 run.py --max_new_tokens=50 \ + --tokenizer_dir ./tmp/Qwen/7B/ \ + --engine_dir=./tmp/Qwen/7B/trt_engines/int4_weight_only/1-gpu/ +``` + +**Demo output of run.py:** +```python +python3 run.py --tokenizer_dir /llm-models/Qwen-7B-Chat/ --engine_dir /engine_qwen +``` +``` +Loading engine from /engine_qwen/qwen_float16_tp1_rank0.engine +Input: "<|im_start|>system +You are a helpful assistant.<|im_end|> +<|im_start|>user +你好，请问你叫什么？<|im_end|> +<|im_start|>assistant +" +Output: "我是来自阿里云的大规模语言模型，我叫通义千问。" +``` + +### Summarization using the Qwen model + +```bash +# Run summarization using the Qwen 7B model in FP16. +python summarize.py --backend=trt_llm \ + --tokenizer_dir ./tmp/Qwen/7B/ \ + --data_type fp16 \ + --engine_dir ./tmp/Qwen/7B/trt_engines/fp16/1-gpu/ + +# Run summarization using the Qwen 7B model in BF16. +python summarize.py --backend=trt_llm \ + --tokenizer_dir ./tmp/Qwen/7B/ \ + --data_type fp16 \ + --engine_dir ./tmp/Qwen/7B/trt_engines/bf16/1-gpu/ + +# Run summarization using the Qwen 7B model quantized to INT8. +python summarize.py --backend=trt_llm \ + --tokenizer_dir ./tmp/Qwen/7B/ \ + --data_type fp16 \ + --engine_dir ./tmp/Qwen/7B/trt_engines/int8_weight_only/1-gpu/ + +# Run summarization using the Qwen 7B model quantized to INT4. +python summarize.py --backend=trt_llm \ + --tokenizer_dir ./tmp/Qwen/7B/ \ + --data_type fp16 \ + --engine_dir ./tmp/Qwen/7B/trt_engines/int4_weight_only/1-gpu/ + +# Run summarization using the Qwen 7B model in FP16 using two GPUs. +mpirun -n 2 --allow-run-as-root \ + python summarize.py --backend=trt_llm \ + --tokenizer_dir ./tmp/Qwen/7B/ \ + --data_type fp16 \ + --engine_dir ./tmp/Qwen/7B/trt_engines/fp16/2-gpu/ + +# Run summarization using the Qwen 14B model in FP16 using two GPUs. +mpirun -n 2 --allow-run-as-root \ + python summarize.py --backend=trt_llm \ + --tokenizer_dir ./tmp/Qwen/14B/ \ + --data_type fp16 \ + --engine_dir ./tmp/Qwen/14B/trt_engines/fp16/2-gpu/ +``` +**Demo output of summarize.py:** +```python +python3 summarize.py --backend=trt_llm --tokenizer_dir /llm-models/Qwen-7B-Chat/ --engine_dir /engine_qwen +``` +``` +[11/09/2023-02:21:10] [TRT-LLM] [I] Load tokenizer takes: 0.4043385982513428 sec +Downloading builder script: 100%|███████████████████████████████████████████| 9.27k/9.27k [00:00<00:00, 35.4MB/s] +Downloading and preparing dataset cnn_dailymail/3.0.0 to /root/.cache/huggingface/datasets/ccdv___cnn_dailymail/3 +...... +[11/09/2023-02:23:33] [TRT-LLM] [I] + Highlights : ['James Best, who played the sheriff on "The Dukes of Hazzard," died Monday at 88 .\n"Hazzard" ran from 1979 to 1985 and was among the most popular shows on TV .'] +[11/09/2023-02:23:33] [TRT-LLM] [I] + Summary : [['Actor James Best, known for his portrayal of bumbling sheriff Rosco P. Coltrane on TV\'s "The Dukes of Hazzard," has died at 88 after a brief illness. Best\'s career spanned decades in theater and Hollywood, but it was his role in "The Dukes of Hazzard" that made him a household name. The show ran for seven seasons from 1979 to 1985 and became a hit on TV, spawning TV movies, an animated series and video games. Best\'s portrayal of Rosco was beloved by fans for his childlike enthusiasm and goofy catchphrases. He is survived by friends and colleagues who paid tribute to him on social media.']] +[11/09/2023-02:23:33] [TRT-LLM] [I] --------------------------------------------------------- +load rouge ... +Downloading builder script: 5.60kB [00:00, 18.9MB/s] +load rouge done +[11/09/2023-02:24:06] [TRT-LLM] [I] TensorRT-LLM (total latency: 30.13867211341858 sec) +[11/09/2023-02:24:06] [TRT-LLM] [I] TensorRT-LLM beam 0 result +[11/09/2023-02:24:06] [TRT-LLM] [I] rouge1 : 26.35215119137573 +[11/09/2023-02:24:06] [TRT-LLM] [I] rouge2 : 9.507814774384485 +[11/09/2023-02:24:06] [TRT-LLM] [I] rougeL : 18.171982659482865 +[11/09/2023-02:24:06] [TRT-LLM] [I] rougeLsum : 21.10413175647868 +``` + +## Credits +This Qwen model example exists thanks to Tltin (TltinDeng01@gmail.com) and Zhaohb (zhaohbcloud@126.com). diff --git a/examples/qwen/benchmark.py b/examples/qwen/benchmark.py new file mode 100644 index 000000000000..2353b6de9ad8 --- /dev/null +++ b/examples/qwen/benchmark.py @@ -0,0 +1,378 @@ +# 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. +"""Benchmark offline inference throughput.""" +import argparse +import json +import os +import random +import time +from typing import List, Tuple + +import torch +from run import QWenForCausalLMGenerationSession, get_model +from tqdm import tqdm, trange +from transformers import (AutoModelForCausalLM, AutoTokenizer, + PreTrainedTokenizerBase) +from utils.utils import get_stop_words_ids, make_context + +now_dir = os.path.dirname(os.path.abspath(__file__)) + +MAX_INPUT_LEN = 2048 +MAX_SEQ_LEN = 4096 + +TRT_MAX_BATCH_SIZE = 2 +TEMPERATURE = 1.0 +TOP_P = 0.5 +TOP_K = 1 + + +def sample_requests( + tokenizer: PreTrainedTokenizerBase, + dataset_path: str, + num_requests: int, + chat_format: str = "chatml", +) -> List[Tuple[str, int, int]]: + # Load the dataset. + with open(dataset_path) as f: + dataset = json.load(f) + # Filter out the conversations with less than 2 turns. + dataset = [data for data in dataset if len(data["conversations"]) >= 2] + # Only keep the first two turns of each conversation. + dataset = [(data["conversations"][0]["value"], + data["conversations"][1]["value"]) for data in dataset] + + # Tokenize the prompts and completions. + tokenized_dataset = [] + for i in trange(len(dataset), desc="Tokenizing for sample"): + prompt = dataset[i][0] + output_text = dataset[i][1] + raw_text, prompt_tokens = make_context(tokenizer=tokenizer, + query=prompt, + max_input_length=MAX_INPUT_LEN, + chat_format=chat_format) + new_token_len = len(tokenizer(output_text).input_ids) + tokenized_dataset.append((raw_text, prompt_tokens, new_token_len)) + + # Filter out too long sequences. + filtered_dataset: List[Tuple[str, int, int]] = [] + for prompt, prompt_token_ids, new_token_len in tokenized_dataset: + prompt_len = len(prompt_token_ids) + if prompt_len < 4 or new_token_len < 4: + # Prune too short sequences. + continue + if prompt_len > MAX_INPUT_LEN or (prompt_len + + new_token_len) > MAX_SEQ_LEN: + # Prune too long sequences. + continue + # limit by MAX_SEQ_LEN + filtered_dataset.append((prompt, prompt_len, new_token_len)) + + # Sample the requests. + sampled_requests = random.sample(filtered_dataset, num_requests) + return sampled_requests + + +def run_trt_llm( + requests: List[Tuple[str, int, int]], + engine_dir: str, + tokenizer_dir: str, + n: int, + max_batch_size: int, +) -> float: + global_max_input_len = MAX_INPUT_LEN + global_max_output_len = MAX_SEQ_LEN + if max_batch_size > TRT_MAX_BATCH_SIZE: + raise Exception( + "max batch size {} must be lower than trt_max_batch_size {}".format( + max_batch_size, TRT_MAX_BATCH_SIZE)) + (model_config, sampling_config, runtime_mapping, runtime_rank, + serialize_path, remove_input_padding, tokenizer, eos_token_id, + pad_token_id) = get_model( + tokenizer_dir=tokenizer_dir, + engine_dir=engine_dir, + ) + with open(serialize_path, 'rb') as f: + engine_buffer = f.read() + decoder = QWenForCausalLMGenerationSession(model_config, engine_buffer, + runtime_mapping) + + # Add the requests to the engine. + sampling_config.num_beams = n + sampling_config.temperature = 0.0 if n > 1 else TEMPERATURE + sampling_config.top_p = TOP_P + sampling_config.top_k = TOP_K + start = time.time() + pad_id = tokenizer.im_end_id + + batch: List[str] = [] + max_new_tokens = 0 + total_num_tokens = [] + for i, (prompt, prompt_len, new_token_len) in tqdm(enumerate(requests), + total=len(requests)): + # Add the prompt to the batch. + batch.append(prompt) + max_new_tokens = max(max_new_tokens, new_token_len) + if len(batch) < max_batch_size and i < len(requests) - 1: + continue + input_ids = [] + input_lengths = [] + for input_text in batch: + input_id = tokenizer( + input_text, + return_tensors="pt", + truncation=True, + max_length=global_max_input_len, + ).input_ids.type(torch.int32) + input_ids.append(input_id) + input_lengths.append(input_id.shape[-1]) + # padding + max_length = max(input_lengths) + # do padding, should move outside the profiling to prevent the overhead + for i in range(len(input_ids)): + pad_size = max_length - input_lengths[i] + + pad = torch.ones([1, pad_size]).type(torch.int32) * pad_id + input_ids[i] = torch.cat([torch.IntTensor(input_ids[i]), pad], + axis=-1) + # do inference + input_ids = torch.cat(input_ids, axis=0).cuda() + input_lengths = torch.IntTensor(input_lengths).type(torch.int32).cuda() + output_ids = decoder.generate( + input_ids=input_ids, + input_lengths=input_lengths, + sampling_config=sampling_config, + max_new_tokens=min(max_new_tokens, + global_max_output_len - input_ids.shape[1]), + ) + pure_output_ids = [] + for i in range(len(batch)): + temp_ids = output_ids[i, input_lengths[i]:] + pure_ids = [] + for i in range(len(temp_ids)): + if temp_ids[i] in [tokenizer.im_start_id, tokenizer.im_end_id]: + pure_ids = temp_ids[:i + 1] + break + if len(pure_ids) == 0: + pure_ids = temp_ids + pure_output_ids.append(pure_ids) + # get the output text + output_texts = [ + tokenizer.decode(out_ids, skip_special_tokens=True) + for out_ids in pure_output_ids + ] + # get the total num of tokens + output_lengths = [len(out_ids) for out_ids in pure_output_ids] + assert len(output_lengths) == len(batch) + for input_len, new_token_len in zip(input_lengths, output_lengths): + total_num_tokens.append(input_len + new_token_len) + batch = [] + max_new_tokens = 0 + + end = time.time() + during = end - start + sum_total_num_tokens = sum(total_num_tokens) + return during, sum_total_num_tokens + + +def run_hf( + requests: List[Tuple[str, int, int]], + model: str, + tokenizer: PreTrainedTokenizerBase, + n: int, + max_batch_size: int, + chat_format: str = "chatml", +) -> float: + global_max_input_len = MAX_INPUT_LEN + global_max_output_len = MAX_SEQ_LEN + llm = AutoModelForCausalLM.from_pretrained(model, + torch_dtype=torch.bfloat16, + trust_remote_code=True) + if llm.config.model_type == "llama": + # To enable padding in the HF backend. + tokenizer.pad_token = tokenizer.eos_token + elif llm.config.model_type == "qwen": + tokenizer.pad_token = tokenizer.decode(tokenizer.im_end_id) + llm = llm.cuda() + stop_words_ids = [] + stop_words_ids.extend(get_stop_words_ids(chat_format, tokenizer)) + stop_words_ids2 = [idx for ids in stop_words_ids for idx in ids] + pbar = tqdm(total=len(requests)) + start = time.time() + total_num_tokens = [] + batch: List[str] = [] + input_lengths: List[int] = [] + max_prompt_len = 0 + max_new_tokens = 0 + for i in range(len(requests)): + prompt, prompt_len, new_token_len = requests[i] + # Add the prompt to the batch. + batch.append(prompt) + input_lengths.append(prompt_len) + max_prompt_len = max(max_prompt_len, prompt_len) + max_new_tokens = max(max_new_tokens, new_token_len) + if len(batch) < max_batch_size and i != len(requests) - 1: + # Check if we can add more requests to the batch. + _, next_prompt_len, next_output_len = requests[i + 1] + temp_input_max = max(max_prompt_len, next_prompt_len) + temp_new_token_max = max(max_new_tokens, next_output_len) + if temp_input_max <= global_max_input_len and \ + (temp_input_max + temp_new_token_max) <= global_max_output_len: + continue + # Generate the sequences. + input_ids = tokenizer( + batch, + return_tensors="pt", + padding=True, + truncation=True, + max_length=global_max_input_len, + ).input_ids + + # limit the max_new_tokens + max_new_tokens = min(max_new_tokens, + global_max_output_len - input_ids.shape[1]) + llm_outputs = llm.generate( + input_ids=input_ids.cuda(), + do_sample=True, + stop_words_ids=stop_words_ids, + num_return_sequences=n, + top_k=TOP_K, + top_p=TOP_P, + temperature=TEMPERATURE, + use_cache=True, + max_new_tokens=max_new_tokens, + ) + pure_output_ids = llm_outputs[:, input_ids.shape[-1]:] + # get the output text + output_texts = tokenizer.batch_decode(pure_output_ids, + skip_special_tokens=True) + output_lengths = [] + for out_ids in pure_output_ids: + early_stop = False + for i in range(len(out_ids)): + if out_ids[i] in stop_words_ids2: + output_lengths.append(i + 1) + early_stop = True + break + if not early_stop: + output_lengths.append(len(out_ids)) + assert len(output_lengths) == len(batch) + for input_len, new_token_len in zip(input_lengths, output_lengths): + total_num_tokens.append(input_len + new_token_len) + pbar.update(len(batch)) + + # Clear the batch. + batch = [] + input_lengths = [] + max_prompt_len = 0 + max_new_tokens = 0 + end = time.time() + during = end - start + sum_total_num_tokens = sum(total_num_tokens) + return during, sum_total_num_tokens + + +def main(args: argparse.Namespace): + print(args) + random.seed(args.seed) + + # Sample the requests. + tokenizer = AutoTokenizer.from_pretrained( + args.tokenizer_dir, + padding_side='left', + trust_remote_code=True, + ) + requests = sample_requests(tokenizer=tokenizer, + dataset_path=args.dataset, + num_requests=args.num_prompts, + chat_format=args.chat_format) + + if args.backend == "trt_llm": + elapsed_time, total_num_tokens = run_trt_llm( + requests=requests, + engine_dir=args.engine_dir, + tokenizer_dir=args.tokenizer_dir, + n=args.n, + max_batch_size=args.trt_max_batch_size, + ) + elif args.backend == "hf": + elapsed_time, total_num_tokens = run_hf( + requests=requests, + model=args.hf_model_dir, + tokenizer=tokenizer, + n=args.n, + max_batch_size=args.hf_max_batch_size, + ) + else: + raise ValueError(f"Unknown backend: {args.backend}") + print(f"Throughput: {len(requests) / elapsed_time:.2f} requests/s, " + f"{total_num_tokens / elapsed_time:.2f} tokens/s") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Benchmark the throughput.") + parser.add_argument( + "--backend", + type=str, + choices=["trt_llm", "hf"], + default="trt_llm", + ) + parser.add_argument("--dataset", + type=str, + default=os.path.join( + now_dir, + "ShareGPT_V3_unfiltered_cleaned_split.json"), + help="Path to the dataset.") + parser.add_argument("--hf_model_dir", type=str, default=None) + parser.add_argument("--tokenizer_dir", + type=str, + default=".", + help="Directory containing the tokenizer.model.") + parser.add_argument('--engine_dir', type=str, default='qwen_outputs') + parser.add_argument("--n", + type=int, + default=1, + help="Number of generated sequences per prompt.") + parser.add_argument("--num-prompts", + type=int, + default=100, + help="Number of prompts to process.") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--hf_max_batch_size", + type=int, + default=1, + help="Maximum batch size for HF backend.") + + parser.add_argument("--trt_max_batch_size", + type=int, + default=1, + help="Maximum batch size for TRT-LLM backend.") + parser.add_argument("--chat-format", + type=str, + default="chatml", + choices=["chatml", "raw"], + help="choice the model format, base or chat") + args = parser.parse_args() + + if args.backend == "trt-llm": + if args.trt_max_batch_size is None: + raise ValueError( + "trt max batch size is required for TRT-LLM backend.") + elif args.backend == "hf": + if args.hf_max_batch_size is None: + raise ValueError("hf max batch size is required for HF backend.") + if args.tokenizer_dir is None: + args.tokenizer_dir = args.hf_model + + main(args) diff --git a/examples/qwen/build.py b/examples/qwen/build.py new file mode 100644 index 000000000000..3908b0857d1e --- /dev/null +++ b/examples/qwen/build.py @@ -0,0 +1,623 @@ +# 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 + +import tensorrt as trt +import torch +import torch.multiprocessing as mp +from transformers import AutoConfig, AutoModelForCausalLM +from weight import load_from_ft, load_from_hf_qwen + +import tensorrt_llm +from tensorrt_llm._utils import str_dtype_to_trt +from tensorrt_llm.builder import Builder +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models import quantize_model +from tensorrt_llm.network import net_guard +from tensorrt_llm.plugin.plugin import ContextFMHAType +from tensorrt_llm.quantization import QuantMode + +MODEL_NAME = "qwen" + +import onnx +import tensorrt as trt +from onnx import TensorProto, helper + +now_dir = os.path.dirname(os.path.abspath(__file__)) + + +def trt_dtype_to_onnx(dtype): + if dtype == trt.float16: + return TensorProto.DataType.FLOAT16 + elif dtype == trt.float32: + return TensorProto.DataType.FLOAT + elif dtype == trt.int32: + return TensorProto.DataType.INT32 + else: + raise TypeError("%s is not supported" % dtype) + + +def to_onnx(network, path): + inputs = [] + for i in range(network.num_inputs): + network_input = network.get_input(i) + inputs.append( + helper.make_tensor_value_info( + network_input.name, trt_dtype_to_onnx(network_input.dtype), + list(network_input.shape))) + + outputs = [] + for i in range(network.num_outputs): + network_output = network.get_output(i) + outputs.append( + helper.make_tensor_value_info( + network_output.name, trt_dtype_to_onnx(network_output.dtype), + list(network_output.shape))) + + nodes = [] + for i in range(network.num_layers): + layer = network.get_layer(i) + layer_inputs = [] + for j in range(layer.num_inputs): + ipt = layer.get_input(j) + if ipt is not None: + layer_inputs.append(layer.get_input(j).name) + layer_outputs = [ + layer.get_output(j).name for j in range(layer.num_outputs) + ] + nodes.append( + helper.make_node(str(layer.type), + name=layer.name, + inputs=layer_inputs, + outputs=layer_outputs, + domain="com.nvidia")) + + onnx_model = helper.make_model(helper.make_graph(nodes, + 'attention', + inputs, + outputs, + initializer=None), + producer_name='NVIDIA') + onnx.save(onnx_model, path) + + +def get_engine_name(model, dtype, tp_size, pp_size, rank): + if pp_size == 1: + return '{}_{}_tp{}_rank{}.engine'.format(model, dtype, tp_size, rank) + return '{}_{}_tp{}_pp{}_rank{}.engine'.format(model, dtype, tp_size, + pp_size, rank) + + +def serialize_engine(engine, path): + logger.info(f'Serializing engine to {path}...') + tik = time.time() + with open(path, 'wb') as f: + f.write(bytearray(engine)) + tok = time.time() + t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) + logger.info(f'Engine serialized. Total time: {t}') + + +def parse_arguments(): + parser = argparse.ArgumentParser() + parser.add_argument('--world_size', + type=int, + default=1, + help='world size, only support tensor parallelism now') + parser.add_argument('--tp_size', type=int, default=1) + parser.add_argument('--pp_size', type=int, default=1) + parser.add_argument('--hf_model_dir', type=str, default=None) + parser.add_argument('--ft_dir_path', type=str, default=None) + parser.add_argument('--dtype', + type=str, + default='float16', + choices=['float32', 'bfloat16', 'float16']) + 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('--log_level', + type=str, + default='info', + choices=[ + 'internal_error', + 'error', + 'warning', + 'info', + 'verbose', + ]) + parser.add_argument('--vocab_size', type=int, default=32000) + parser.add_argument('--n_layer', type=int, default=32) + parser.add_argument('--n_positions', type=int, default=2048) + parser.add_argument('--n_embd', type=int, default=4096) + parser.add_argument('--n_head', type=int, default=32) + parser.add_argument('--n_kv_head', type=int, default=None) + parser.add_argument('--inter_size', type=int, default=11008) + parser.add_argument('--hidden_act', type=str, default='silu') + parser.add_argument('--max_batch_size', type=int, default=2) + parser.add_argument('--max_input_len', type=int, default=2048) + parser.add_argument('--max_output_len', type=int, default=2048) + parser.add_argument('--max_beam_width', type=int, default=1) + parser.add_argument('--rotary_base', type=float, default=10000.0) + parser.add_argument('--rotary_scaling', nargs=2, type=str, default=None) + parser.add_argument('--use_gpt_attention_plugin', + nargs='?', + type=str, + default="float16", + choices=['float16', 'bfloat16', 'float32', None]) + parser.add_argument('--use_gemm_plugin', + nargs='?', + type=str, + default="float16", + choices=['float16', 'bfloat16', 'float32', None]) + parser.add_argument('--parallel_build', default=False, action='store_true') + parser.add_argument('--enable_context_fmha', + default=False, + action='store_true') + parser.add_argument('--enable_context_fmha_fp32_acc', + default=False, + action='store_true') + parser.add_argument('--visualize', default=False, action='store_true') + parser.add_argument('--enable_debug_output', + default=False, + action='store_true') + parser.add_argument('--gpus_per_node', type=int, default=8) + parser.add_argument('--builder_opt', type=int, default=None) + parser.add_argument( + '--output_dir', + type=str, + default='qwen_outputs', + help= + 'The path to save the serialized engine files, timing cache file and model configs' + ) + parser.add_argument('--remove_input_padding', + default=False, + action='store_true') + # Arguments related to the quantization of the model. + parser.add_argument( + '--use_smooth_quant', + default=False, + action="store_true", + help= + 'Use the SmoothQuant method to quantize activations and weights for the various GEMMs.' + 'See --per_channel and --per_token for finer-grained quantization options.' + ) + parser.add_argument( + '--per_channel', + default=False, + action="store_true", + help= + 'By default, we use a single static scaling factor for the GEMM\'s result. ' + 'per_channel instead uses a different static scaling factor for each channel. ' + 'The latter is usually more accurate, but a little slower.') + parser.add_argument( + '--per_token', + default=False, + action="store_true", + help= + 'By default, we use a single static scaling factor to scale activations in the int8 range. ' + 'per_token chooses at run time, and for each token, a custom scaling factor. ' + 'The latter is usually more accurate, but a little slower.') + + parser.add_argument( + '--per_group', + default=False, + action="store_true", + help= + 'By default, we use a single static scaling factor to scale weights in the int4 range. ' + 'per_group chooses at run time, and for each group, a custom scaling factor. ' + 'The flag is built for GPTQ/AWQ quantization.') + + parser.add_argument( + '--use_weight_only', + default=False, + action="store_true", + help='Quantize weights for the various GEMMs to INT4/INT8.' + 'See --weight_only_precision to set the precision') + + parser.add_argument( + '--weight_only_precision', + const='int8', + type=str, + nargs='?', + default='int8', + choices=['int8', 'int4'], + help= + 'Define the precision for the weights when using weight-only quantization.' + 'You must also use --use_weight_only for that argument to have an impact.' + ) + parser.add_argument( + '--use_inflight_batching', + action="store_true", + default=False, + help="Activates inflight batching mode of gptAttentionPlugin.") + parser.add_argument( + '--paged_kv_cache', + action="store_true", + default=False, + help= + 'By default we use contiguous KV cache. By setting this flag you enable paged KV cache' + ) + parser.add_argument('--tokens_per_block', + type=int, + default=64, + help='Number of tokens per block in paged KV cache') + + parser.add_argument( + '--max_num_tokens', + type=int, + default=None, + help='Define the max number of tokens supported by the engine') + + parser.add_argument( + '--int8_kv_cache', + default=False, + action="store_true", + help= + 'By default, we use dtype for KV cache. int8_kv_cache chooses int8 quantization for KV' + ) + 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=1, # Meta does TP on hidden dim + 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( + '--strongly_typed', + default=False, + action="store_true", + help= + 'This option is introduced with trt 9.1.0.1+ and will reduce the building time significantly for fp8.' + ) + parser.add_argument( + '--use_custom_all_reduce', + action='store_true', + help= + 'Activates latency-optimized algorithm for all-reduce instead of NCCL.') + + args = parser.parse_args() + assert not ( + args.use_smooth_quant and args.use_weight_only + ), "You cannot enable both SmoothQuant and INT8 weight-only together." + + if not args.remove_input_padding: + if args.use_gpt_attention_plugin: + logger.warning( + f"It is recommended to specify --remove_input_padding when using GPT attention plugin" + ) + + if args.use_inflight_batching: + if not args.use_gpt_attention_plugin: + args.use_gpt_attention_plugin = 'float16' + logger.info( + f"Using GPT attention plugin for inflight batching mode. Setting to default '{args.use_gpt_attention_plugin}'" + ) + if not args.remove_input_padding: + args.remove_input_padding = True + logger.info( + "Using remove input padding for inflight batching mode.") + if not args.paged_kv_cache: + args.paged_kv_cache = True + logger.info("Using paged KV cache for inflight batching mode.") + + if args.use_smooth_quant: + args.quant_mode = QuantMode.use_smooth_quant(args.per_token, + args.per_channel) + elif args.use_weight_only: + if args.per_group: + args.quant_mode = QuantMode.from_description( + quantize_weights=True, + quantize_activations=False, + per_token=False, + per_channel=False, + per_group=True, + use_int4_weights=True) + else: + args.quant_mode = QuantMode.use_weight_only( + args.weight_only_precision == 'int4') + else: + args.quant_mode = QuantMode(0) + + if args.int8_kv_cache: + args.quant_mode = args.quant_mode.set_int8_kv_cache() + + if args.hf_model_dir is not None: + hf_config = AutoConfig.from_pretrained( + args.hf_model_dir, + trust_remote_code=True, + ) + args.inter_size = hf_config.intermediate_size # override the inter_size for QWen + args.n_embd = hf_config.hidden_size + args.n_head = hf_config.num_attention_heads + if hasattr(hf_config, "num_key_value_heads"): + args.n_kv_head = hf_config.num_key_value_heads + args.n_layer = hf_config.num_hidden_layers + args.n_positions = hf_config.max_position_embeddings + args.vocab_size = hf_config.vocab_size + args.hidden_act = "silu" + args.kv_channels = hf_config.kv_channels + args.rotary_emb_base = hf_config.rotary_emb_base + assert args.use_gpt_attention_plugin is not None, "QWen must use gpt attention plugin" + if args.n_kv_head is not None and args.n_kv_head != args.n_head: + assert (args.n_head % args.n_kv_head) == 0, \ + "MQA/GQA requires the number of heads to be divisible by the number of K/V heads." + assert args.n_kv_head == args.tp_size, \ + "The current implementation of GQA requires the number of K/V heads to match the number of GPUs." \ + "This limitation will be removed in a future version." + + assert args.pp_size * args.tp_size == args.world_size + + if args.max_num_tokens is not None: + assert args.enable_context_fmha + + return args + + +def build_rank_engine(builder: Builder, + builder_config: tensorrt_llm.builder.BuilderConfig, + engine_name, rank, multi_query_mode, args): + ''' + @brief: Build the engine on the given rank. + @param rank: The rank to build the engine. + @param args: The cmd line arguments. + @return: The built engine. + ''' + kv_dtype = str_dtype_to_trt(args.dtype) + mapping = Mapping(world_size=args.world_size, + rank=rank, + tp_size=args.tp_size, + pp_size=args.pp_size) + + # Initialize Module + tensorrt_llm_qwen = tensorrt_llm.models.QWenForCausalLM( + num_layers=args.n_layer, + num_heads=args.n_head, + num_kv_heads=args.n_kv_head, + hidden_size=args.n_embd, + seq_length=args.max_input_len, + vocab_size=args.vocab_size, + hidden_act=args.hidden_act, + max_position_embeddings=args.n_positions, + dtype=kv_dtype, + mlp_hidden_size=args.inter_size, + neox_rotary_style=True, + mapping=mapping, + rotary_base=args.rotary_base, + rotary_scaling=args.rotary_scaling, + use_parallel_embedding=args.use_parallel_embedding, + embedding_sharding_dim=args.embedding_sharding_dim, + quant_mode=args.quant_mode, + ) + quantize_kwargs = {} + if args.use_smooth_quant or args.use_weight_only: + if args.weight_only_precision == 'int4_awq': + quantize_kwargs = { + "group_size": args.group_size, + "zero": False, + "pre_quant_scale": True, + "exclude_modules": [], + } + elif args.weight_only_precision == 'int4_gptq': + quantize_kwargs = { + "group_size": args.group_size, + "zero": True, + "pre_quant_scale": False, + } + tensorrt_llm_qwen = quantize_model(tensorrt_llm_qwen, args.quant_mode, + **quantize_kwargs) + ft_dir_path = args.ft_dir_path + if args.hf_model_dir is not None and \ + (ft_dir_path is None or not os.path.exists(ft_dir_path)): + logger.info(f'Loading HF QWen ... from {args.hf_model_dir}') + tik = time.time() + hf_qwen = AutoModelForCausalLM.from_pretrained( + args.hf_model_dir, + device_map={ + "transformer": "cpu", + "lm_head": "cpu" + }, # Load to CPU memory + torch_dtype="auto", + trust_remote_code=True, + ) + tok = time.time() + t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) + logger.info(f'HF QWen loaded. Total time: {t}') + load_from_hf_qwen(tensorrt_llm_qwen, + hf_qwen, + mapping, + max_position_embeddings=args.n_positions, + kv_channels=args.kv_channels, + rotary_emb_base=args.rotary_emb_base, + dtype=args.dtype, + multi_query_mode=multi_query_mode) + del hf_qwen + elif ft_dir_path is not None: + dir_path = ft_dir_path + logger.info(f'Loading FT QWen ... from {ft_dir_path}') + load_from_ft(tensorrt_llm_qwen, + dir_path, + mapping, + dtype=args.dtype, + multi_query_mode=multi_query_mode) + else: + raise ValueError( + "You must specify either --hf_model_dir or --ft_dir_path") + + # Module -> Network + network = builder.create_network() + network.trt_network.name = engine_name + if args.use_gpt_attention_plugin: + network.plugin_config.set_gpt_attention_plugin( + dtype=args.use_gpt_attention_plugin) + if args.use_gemm_plugin: + network.plugin_config.set_gemm_plugin(dtype=args.use_gemm_plugin) + # Quantization plugins. + if args.use_smooth_quant: + network.plugin_config.set_smooth_quant_gemm_plugin(dtype=args.dtype) + network.plugin_config.set_rmsnorm_quantization_plugin(dtype=args.dtype) + network.plugin_config.set_quantize_tensor_plugin() + network.plugin_config.set_quantize_per_token_plugin() + assert not (args.enable_context_fmha and args.enable_context_fmha_fp32_acc) + if args.enable_context_fmha: + network.plugin_config.set_context_fmha(ContextFMHAType.enabled) + if args.enable_context_fmha_fp32_acc: + network.plugin_config.set_context_fmha( + ContextFMHAType.enabled_with_fp32_acc) + if args.use_weight_only: + if args.per_group: + network.plugin_config.set_weight_only_groupwise_quant_matmul_plugin( + dtype='float16') + else: + network.plugin_config.set_weight_only_quant_matmul_plugin( + dtype='float16') + if args.world_size > 1: + network.plugin_config.set_nccl_plugin(args.dtype, + args.use_custom_all_reduce) + if args.remove_input_padding: + network.plugin_config.enable_remove_input_padding() + + if args.paged_kv_cache: + network.plugin_config.enable_paged_kv_cache(args.tokens_per_block) + + with net_guard(network): + # Prepare + network.set_named_parameters(tensorrt_llm_qwen.named_parameters()) + + # Forward + inputs = tensorrt_llm_qwen.prepare_inputs( + max_batch_size=args.max_batch_size, + max_input_len=args.max_input_len, + max_new_tokens=args.max_output_len, + use_cache=True, + max_beam_width=args.max_beam_width, + max_num_tokens=args.max_num_tokens, + ) + tensorrt_llm_qwen(*inputs) + if args.enable_debug_output: + # mark intermediate nodes' outputs + for k, v in tensorrt_llm_qwen.named_network_outputs(): + v = v.trt_tensor + v.name = k + network.trt_network.mark_output(v) + v.dtype = kv_dtype + if args.visualize: + model_path = os.path.join(args.output_dir, 'test.onnx') + to_onnx(network.trt_network, model_path) + + engine = None + + # Network -> Engine + engine = builder.build_engine(network, builder_config) + if rank == 0: + config_path = os.path.join(args.output_dir, 'config.json') + builder.save_config(builder_config, config_path) + return engine + + +def build(rank, args): + torch.cuda.set_device(rank % args.gpus_per_node) + tensorrt_llm.logger.set_level(args.log_level) + if not os.path.exists(args.output_dir): + os.makedirs(args.output_dir) + multi_query_mode = (args.n_kv_head + is not None) and (args.n_kv_head != args.n_head) + + # when doing serializing build, all ranks share one engine + builder = Builder() + + cache = None + for cur_rank in range(args.world_size): + # skip other ranks if parallel_build is enabled + if args.parallel_build and cur_rank != rank: + continue + int8_trt_flag = args.quant_mode.has_act_or_weight_quant() or ( + not args.paged_kv_cache and args.quant_mode.has_int8_kv_cache()) + builder_config = builder.create_builder_config( + name=MODEL_NAME, + precision=args.dtype, + timing_cache=args.timing_cache if cache is None else cache, + tensor_parallel=args.tp_size, + pipeline_parallel=args.pp_size, + parallel_build=args.parallel_build, + num_layers=args.n_layer, + num_heads=args.n_head, + hidden_size=args.n_embd, + vocab_size=args.vocab_size, + hidden_act=args.hidden_act, + max_position_embeddings=args.n_positions, + max_batch_size=args.max_batch_size, + max_beam_width=args.max_beam_width, + max_input_len=args.max_input_len, + max_output_len=args.max_output_len, + max_num_tokens=args.max_num_tokens, + int8=int8_trt_flag, + fp8=args.quant_mode.has_fp8_qdq(), + quant_mode=args.quant_mode, + strongly_typed=args.strongly_typed, + opt_level=args.builder_opt) + engine_name = get_engine_name(MODEL_NAME, args.dtype, args.tp_size, + args.pp_size, cur_rank) + engine = build_rank_engine(builder, builder_config, engine_name, + cur_rank, multi_query_mode, args) + assert engine is not None, f'Failed to build engine for rank {cur_rank}' + + if cur_rank == 0: + # Use in-memory timing cache for multiple builder passes. + if not args.parallel_build: + cache = builder_config.trt_builder_config.get_timing_cache() + + serialize_engine(engine, os.path.join(args.output_dir, engine_name)) + + if rank == 0: + ok = builder.save_timing_cache( + builder_config, os.path.join(args.output_dir, "model.cache")) + assert ok, "Failed to save timing cache." + + +if __name__ == '__main__': + args = parse_arguments() + logger.set_level(args.log_level) + tik = time.time() + if args.parallel_build and args.world_size > 1 and \ + torch.cuda.device_count() >= args.world_size: + logger.warning( + f'Parallelly build TensorRT engines. Please make sure that all of the {args.world_size} GPUs are totally free.' + ) + mp.spawn(build, nprocs=args.world_size, args=(args, )) + else: + args.parallel_build = False + logger.info('Serially build TensorRT engines.') + build(0, args) + + tok = time.time() + t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) + logger.info(f'Total time of building all {args.world_size} engines: {t}') diff --git a/examples/qwen/hf_qwen_convert.py b/examples/qwen/hf_qwen_convert.py new file mode 100644 index 000000000000..a642163cd96d --- /dev/null +++ b/examples/qwen/hf_qwen_convert.py @@ -0,0 +1,361 @@ +# 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. +''' +Convert huggingface QWen-7B-Chat model to numpy file. +Use https://huggingface.co/Qwen/Qwen-7B-Chat as demo. +''' +import argparse +import configparser +import dataclasses +import json +import os +from pathlib import Path + +import torch +import torch.multiprocessing as multiprocessing +from smoothquant import capture_activation_range, smooth_gemm, smooth_gemm_mlp +from tqdm import tqdm +from transformers import AutoModelForCausalLM # transformers-4.10.0-py3 +from transformers import AutoTokenizer, GenerationConfig +# for debug +from utils.convert import split_and_save_weight + +from tensorrt_llm._utils import str_dtype_to_torch, torch_to_numpy + +now_dir = os.path.dirname(os.path.abspath(__file__)) + + +@dataclasses.dataclass(frozen=True) +class ProgArgs: + out_dir: str + in_file: str + max_input_len: int = 2048 + tensor_parallelism: int = 1 + processes: int = 1 + calibrate_kv_cache: bool = False + smoothquant: float = None + model: str = "qwen" + storage_type: str = "fp32" + dataset_cache_dir: str = None + + @staticmethod + def parse(args=None) -> 'ProgArgs': + parser = argparse.ArgumentParser( + formatter_class=argparse.RawTextHelpFormatter) + parser.add_argument('--out-dir', + '-o', + type=str, + help='file name of output directory', + required=True) + parser.add_argument('--in-file', + '-i', + type=str, + help='file name of input checkpoint file', + required=True) + parser.add_argument( + '--max_input_len', + type=int, + help= + "This should be consistent with the max_input_len you used when building engine.", + default=2048) + parser.add_argument('--tensor-parallelism', + '-tp', + type=int, + help='Requested tensor parallelism for inference', + default=1) + parser.add_argument( + "--processes", + "-p", + type=int, + help= + "How many processes to spawn for conversion (default: 1). Set it to a lower value to reduce RAM usage.", + default=1) + parser.add_argument( + "--calibrate-kv-cache", + "-kv", + action="store_true", + help= + "Generate scaling factors for KV cache. Used for storing KV cache in int8." + ) + parser.add_argument( + "--smoothquant", + "-sq", + type=float, + default=None, + help="Set the α parameter (see https://arxiv.org/pdf/2211.10438.pdf)" + " to Smoothquant the model, and output int8 weights." + " A good first try is 0.5. Must be in [0, 1]") + parser.add_argument( + "--model", + default="qwen", + type=str, + help="Specify GPT variants to convert checkpoints correctly", + choices=["qwen", "gpt2", "santacoder", "starcoder"]) + parser.add_argument("--storage-type", + "-t", + type=str, + default="float16", + choices=["float32", "float16", "bfloat16"]) + parser.add_argument("--dataset-cache-dir", + type=str, + default=None, + help="cache dir to load the hugging face dataset") + return ProgArgs(**vars(parser.parse_args(args))) + + +@torch.no_grad() +def smooth_qwen_model(model, scales, alpha, qwen_smoother): + # Smooth the activation and weights with smoother = $\diag{s}$ + for name, module in model.named_modules(): + # if not isinstance(module, QWenBlock): + if not str(type(module)).endswith("QWenBlock'>"): + continue + + # qkv_proj + layer_name = name + ".attn.c_attn" + smoother = smooth_gemm(module.attn.c_attn.weight, + scales[layer_name]["x"], + module.ln_1.weight, + alpha=alpha) + scales[layer_name]["x"] = scales[layer_name]["x"] / smoother + scales[layer_name]["w"] = module.attn.c_attn.weight.abs().max(dim=1)[0] + + # attention dense + layer_name = name + ".attn.c_proj" + smoother3 = smooth_gemm( + module.attn.c_proj.weight, + scales[layer_name]["x"], + None, + alpha=alpha, + ) + qwen_smoother[layer_name] = smoother3.float() + + scales[layer_name]["x"] = scales[layer_name]["x"] / smoother3 + scales[layer_name]["w"] = module.attn.c_proj.weight.abs().max(dim=1)[0] + + # mlp w1 / w2, because then use some input hidden_states as input, so we need to smooth it with same scale + mlp_w1_name = name + ".mlp.w1" + mlp_w2_name = name + ".mlp.w2" + smoother2 = smooth_gemm_mlp(module.mlp.w1.weight, + module.mlp.w2.weight, + scales[mlp_w1_name]["x"], + module.ln_2.weight, + alpha=alpha) + scales[mlp_w1_name]["x"] = scales[mlp_w1_name]["x"] / smoother2 + scales[mlp_w2_name]["x"] = scales[mlp_w2_name]["x"] / smoother2 + scales[mlp_w1_name]["w"] = module.mlp.w1.weight.abs().max(dim=1)[0] + scales[mlp_w2_name]["w"] = module.mlp.w2.weight.abs().max(dim=1)[0] + + # mlp c_proj + layer_name = name + ".mlp.c_proj" + smoother4 = smooth_gemm(module.mlp.c_proj.weight, + scales[layer_name]["x"], + None, + alpha=alpha) + qwen_smoother[layer_name] = smoother4.float() + scales[layer_name]["x"] = scales[layer_name]["x"] / smoother4 + scales[layer_name]["w"] = module.mlp.c_proj.weight.abs().max(dim=1)[0] + + +# SantaCoder separates Q projection from KV projection +def concat_qkv_weight_bias(q, hf_key, hf_model): + kv = hf_model.state_dict()[hf_key.replace("q_attn", "kv_attn")] + return torch.cat([q, kv], dim=-1) + + +# StarCoder uses nn.Linear for these following ops whose weight matrix is transposed compared to transformer.Conv1D +def transpose_weights(hf_name, param): + weight_to_transpose = [ + "attn.c_attn", "attn.c_proj", "mlp.c_proj", "mlp.w1", "mlp.w2" + ] + if any([k in hf_name for k in weight_to_transpose]): + if len(param.shape) == 2: + param = param.transpose(0, 1) + return param + + +def convert_qwen_name(orig_name): + global_weights = { + "transformer.wte.weight": "vocab_embedding.weight", + "transformer.ln_f.weight": "ln_f.weight", + "lm_head.weight": "lm_head.weight" + } + + if orig_name in global_weights: + return global_weights[orig_name] + + _, _, layer_id, *weight_name = orig_name.split(".") + layer_id = int(layer_id) + weight_name = "transformer." + ".".join(weight_name) + + per_layer_weights = { + "transformer.ln_1.weight": "ln_1.weight", + "transformer.ln_2.weight": "ln_2.weight", + "transformer.attn.c_attn.weight": "attention.qkv.weight", + "transformer.attn.c_attn.bias": "attention.qkv.bias", + "transformer.attn.c_proj.weight": "attention.dense.weight", + "transformer.mlp.w1.weight": "mlp.w1.weight", + "transformer.mlp.w2.weight": "mlp.w2.weight", + "transformer.mlp.c_proj.weight": "mlp.c_proj.weight", + } + return f"layers.{layer_id}.{per_layer_weights[weight_name]}" + + +@torch.no_grad() +def hf_qwen_converter(args: ProgArgs): + infer_tp = args.tensor_parallelism + multi_query_mode = True if args.model in ["santacoder", "starcoder" + ] else False + saved_dir = Path(args.out_dir) / f"{infer_tp}-gpu" + saved_dir.mkdir(parents=True, exist_ok=True) + + # load position_embedding from rank 0 + model = AutoModelForCausalLM.from_pretrained( + args.in_file, + device_map= + "auto", # if you gpu memory is not enough, you can set device_map="cpu" + trust_remote_code=True, + torch_dtype=str_dtype_to_torch(args.storage_type), + ).half() # if you gpu memory is not enough, you can set .half() to .float() + model.generation_config = GenerationConfig.from_pretrained( + args.in_file, trust_remote_code=True) + act_range = {} + qwen_smoother = {} + if args.smoothquant is not None or args.calibrate_kv_cache: + os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get( + "TOKENIZERS_PARALLELISM", "false") + from datasets import load_dataset + + # copy from summarize.py + dataset_cnn = load_dataset("ccdv/cnn_dailymail", '3.0.0') + dataset = dataset_cnn["test"] + tokenizer = AutoTokenizer.from_pretrained( + args.in_file, + legacy=False, + padding_side='left', + trust_remote_code=True, + ) + gen_config_path = os.path.join(args.in_file, 'generation_config.json') + with open(gen_config_path, 'r') as f: + gen_config = json.load(f) + chat_format = gen_config['chat_format'] + tokenizer.pad_token_id = tokenizer.im_end_id + # use this prompt to make chat model do summarize + system_prompt = "You are a useful assistant, please directly output the corresponding summary according to the article entered by the user." + act_range = capture_activation_range( + model, + tokenizer, + dataset, + system_prompt=system_prompt, + chat_format=chat_format, + max_input_len=args.max_input_len, + ) + if args.smoothquant is not None: + smooth_qwen_model(model, act_range, args.smoothquant, qwen_smoother) + + config = configparser.ConfigParser() + config["qwen"] = {} + for key in vars(args): + config["qwen"][key] = f"{vars(args)[key]}" + for k, v in vars(model.config).items(): + config["qwen"][k] = f"{v}" + config["qwen"]["storage_dtype"] = args.storage_type + config["qwen"]["multi_query_mode"] = str(multi_query_mode) + with open(saved_dir / "config.ini", 'w') as configfile: + config.write(configfile) + + storage_type = str_dtype_to_torch(args.storage_type) + + global_weights = ["vocab_embedding.weight", "ln_f.weight", "lm_head.weight"] + + int8_outputs = None + if args.calibrate_kv_cache: + int8_outputs = "kv_cache_only" + if args.smoothquant is not None: + int8_outputs = "all" + + starmap_args = [] + for name, param in tqdm( + model.named_parameters(), + desc="convert and save", + total=len(list(model.parameters())), + ncols=80, + ): + if "weight" not in name and "bias" not in name: + continue + converted_name = convert_qwen_name(name) + if name.replace(".weight", "") in qwen_smoother.keys(): + smoother = qwen_smoother[name.replace(".weight", "")] + starmap_arg = ( + 0, + saved_dir, + infer_tp, + f"{converted_name}.smoother".replace(".weight", ""), + smoother, + storage_type, + None, + { + "int8_outputs": int8_outputs, + "multi_query_mode": multi_query_mode, + "local_dim": None, + }, + ) + if args.processes > 1: + starmap_args.append(starmap_arg) + else: + split_and_save_weight(*starmap_arg) + + param = transpose_weights(name, param) + if converted_name in global_weights: + torch_to_numpy(param.to(storage_type).cpu()).tofile( + saved_dir / f"{converted_name}.bin") + else: + if 'q_attn' in name: + param = concat_qkv_weight_bias(param, name, model) + converted_name = converted_name.replace("query", + "query_key_value") + # Needed by QKV projection weight split. With multi_query_mode one does not simply take + # out_dim and divide it by 3 to get local_dim because out_dim = local_dim + 2 * head_size + local_dim = model.transformer.h[ + 0].attn.embed_dim if multi_query_mode else None + starmap_arg = (0, saved_dir, infer_tp, converted_name, + param.to(storage_type), storage_type, + act_range.get(name.replace(".weight", "")), { + "int8_outputs": int8_outputs, + "multi_query_mode": multi_query_mode, + "local_dim": local_dim + }) + if args.processes > 1: + starmap_args.append(starmap_arg) + else: + split_and_save_weight(*starmap_arg) + + if args.processes > 1: + starmap_args = tqdm(starmap_args, desc="saving weights") + with multiprocessing.Pool(args.processes) as pool: + pool.starmap(split_and_save_weight, starmap_args) + + +def run_conversion(args: ProgArgs): + print("\n=============== Arguments ===============") + for key, value in vars(args).items(): + print(f"{key}: {value}") + print("========================================") + hf_qwen_converter(args) + + +if __name__ == "__main__": + torch.multiprocessing.set_start_method("spawn") + run_conversion(ProgArgs.parse()) diff --git a/examples/qwen/requirements.txt b/examples/qwen/requirements.txt new file mode 100644 index 000000000000..69b76164be16 --- /dev/null +++ b/examples/qwen/requirements.txt @@ -0,0 +1,14 @@ +datasets~=2.3.2 +rouge_score~=0.1.2 +transformers~=4.33.1 +transformers-stream-generator +sentencepiece~=0.1.99 +tiktoken +einops + +# optional dependencies +gradio==3.40.1 +mdtex2html +sse_starlette +aiohttp_sse_client +openai diff --git a/examples/qwen/run.py b/examples/qwen/run.py new file mode 100644 index 000000000000..7df77727eff3 --- /dev/null +++ b/examples/qwen/run.py @@ -0,0 +1,315 @@ +# 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 csv +import json +import os +from pathlib import Path + +import numpy as np +import torch +from transformers import AutoTokenizer + +import tensorrt_llm +from tensorrt_llm.quantization import QuantMode +from tensorrt_llm.runtime import GenerationSession, ModelConfig, SamplingConfig +from tensorrt_llm.runtime.generation import Mapping + +from build import get_engine_name # isort:skip + +now_dir = os.path.dirname(os.path.abspath(__file__)) + +MAX_INPUT_LEN = 2048 +MAX_SEQ_LEN = 4096 + + +class QWenForCausalLMGenerationSession(GenerationSession): + + def __init__( + self, + model_config: ModelConfig, + engine_buffer, + mapping: Mapping, + debug_mode=False, + debug_tensors_to_save=None, + cuda_graph_mode=False, + stream: torch.cuda.Stream = None, + global_max_input_length=MAX_INPUT_LEN, + global_max_output_length=MAX_SEQ_LEN, + ): + super().__init__(model_config, + engine_buffer, + mapping, + debug_mode, + debug_tensors_to_save=debug_tensors_to_save, + cuda_graph_mode=cuda_graph_mode, + stream=stream) + self.global_max_input_length = global_max_input_length + self.global_max_output_length = global_max_output_length + + def generate( + self, + input_ids: torch.Tensor, + input_lengths: torch.Tensor, + sampling_config: SamplingConfig, + max_new_tokens: int, + runtime_rank: int = 0, + ): + max_input_length = torch.max(input_lengths).item() + max_new_tokens = min(max_new_tokens, + self.global_max_output_length - max_input_length) + # setup batch_size, max_input_length, max_output_len + self.setup(batch_size=input_lengths.size(0), + max_context_length=max_input_length, + max_new_tokens=max_new_tokens) + output_ids = self.decode(input_ids, input_lengths, sampling_config) + with torch.no_grad(): + torch.cuda.synchronize() + if runtime_rank == 0: + outputs = output_ids[:, 0, :] + return outputs + + +def parse_arguments(): + parser = argparse.ArgumentParser() + parser.add_argument('--max_new_tokens', type=int, default=200) + parser.add_argument('--log_level', type=str, default='error') + parser.add_argument( + '--engine_dir', + type=str, + default="qwen_outputs", + ) + parser.add_argument('--tokenizer_dir', + type=str, + default=".", + help="Directory containing the tokenizer.model.") + default_text = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n你好，请问你叫什么？<|im_end|>\n<|im_start|>assistant\n" + parser.add_argument('--input_text', type=str, default=default_text) + 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) + return parser.parse_args() + + +def get_model(tokenizer_dir, engine_dir, log_level='error'): + # --load the tokenizer and engine # + tensorrt_llm.logger.set_level(log_level) + tokenizer = AutoTokenizer.from_pretrained( + tokenizer_dir, + legacy=False, + trust_remote_code=True, + ) + config_path = os.path.join(engine_dir, 'config.json') + with open(config_path, 'r') as f: + config = json.load(f) + gen_config_path = os.path.join(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['plugin_config']['gpt_attention_plugin'] + remove_input_padding = config['plugin_config']['remove_input_padding'] + dtype = config['builder_config']['precision'] + tp_size = config['builder_config']['tensor_parallel'] + pp_size = config['builder_config']['pipeline_parallel'] + 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['builder_config']['num_heads'] // world_size + hidden_size = config['builder_config']['hidden_size'] // world_size + vocab_size = config['builder_config']['vocab_size'] + num_layers = config['builder_config']['num_layers'] + num_kv_heads = config['builder_config'].get('num_kv_heads', num_heads) + paged_kv_cache = config['plugin_config']['paged_kv_cache'] + tokens_per_block = config['plugin_config']['tokens_per_block'] + quant_mode = QuantMode(config['builder_config']['quant_mode']) + if config['builder_config'].get('multi_query_mode', False): + tensorrt_llm.logger.warning( + "`multi_query_mode` config is deprecated. Please rebuild the engine." + ) + num_kv_heads = 1 + use_custom_all_reduce = config['plugin_config'].get('use_custom_all_reduce', + False) + + 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(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, + paged_kv_cache=paged_kv_cache, + tokens_per_block=tokens_per_block, + remove_input_padding=remove_input_padding, + dtype=dtype, + quant_mode=quant_mode, + use_custom_all_reduce=use_custom_all_reduce) + sampling_config = SamplingConfig( + end_id=eos_token_id, + pad_id=pad_token_id, + num_beams=1, + top_k=top_k, + top_p=top_p, + ) + + engine_name = get_engine_name('qwen', dtype, tp_size, pp_size, runtime_rank) + serialize_path = os.path.join(engine_dir, engine_name) + print(f'Loading engine from {serialize_path}') + return (model_config, sampling_config, runtime_mapping, runtime_rank, + serialize_path, remove_input_padding, tokenizer, eos_token_id, + pad_token_id) + + +def generate( + max_new_tokens: int, + log_level: str = 'error', + engine_dir: str = 'qwen_outputs', + input_text: str = 'Born in north-east France, Soyer trained as a', + input_file: str = None, + output_csv: str = None, + output_npy: str = None, + tokenizer_dir: str = None, + num_beams: int = 1, +): + (model_config, sampling_config, runtime_mapping, runtime_rank, + serialize_path, remove_input_padding, tokenizer, eos_token_id, + pad_token_id) = get_model(tokenizer_dir, engine_dir, log_level) + with open(serialize_path, 'rb') as f: + engine_buffer = f.read() + decoder = QWenForCausalLMGenerationSession( + model_config, + engine_buffer, + runtime_mapping, + ) + + input_tokens = [] + if input_file is None: + input_tokens.append( + tokenizer.encode(input_text, add_special_tokens=False)) + 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_tokens.append(np.array(line, dtype='int32')) + elif input_file.endswith('.npy'): + inputs = np.load(input_file) + for row in inputs: + row = row[row != eos_token_id] + input_tokens.append(row) + else: + print('Input file format not supported.') + raise SystemExit + + input_ids = None + input_lengths = None + if input_file is None: + input_ids = torch.tensor(input_tokens, device="cuda", dtype=torch.int32) + input_lengths = torch.tensor([input_ids.size(1)], + device="cuda", + dtype=torch.int32) + else: + input_lengths = torch.tensor([len(x) for x in input_tokens], + device="cuda", + dtype=torch.int32) + if remove_input_padding: + input_ids = np.concatenate(input_tokens) + input_ids = torch.tensor(input_ids, + device="cuda", + dtype=torch.int32).unsqueeze(0) + else: + input_ids = torch.nested.to_padded_tensor( + torch.nested.nested_tensor(input_tokens, dtype=torch.int32), + eos_token_id).cuda() + + max_input_length = torch.max(input_lengths).item() + max_new_tokens = min(max_new_tokens, MAX_SEQ_LEN - max_input_length) + decoder.setup(batch_size=input_lengths.size(0), + max_context_length=max_input_length, + max_new_tokens=max_new_tokens) + + output_ids = decoder.decode(input_ids, input_lengths, sampling_config) + torch.cuda.synchronize() + + if runtime_rank == 0: + if output_csv is None and output_npy is None: + for b in range(input_lengths.size(0)): + inputs = input_tokens[b] + input_text = tokenizer.decode(inputs) + print(f'Input: \"{input_text}\"') + if num_beams <= 1: + outputs = output_ids[b][0, len(inputs):].tolist() + output_text = tokenizer.decode(outputs, + skip_special_tokens=True) + print(f'Output: \"{output_text}\"') + else: + for beam in range(num_beams): + outputs = output_ids[b][beam, len(inputs):].tolist() + output_text = tokenizer.decode(outputs, + skip_special_tokens=True) + print(f'Output(beam: {beam}): \"{output_text}\"') + + 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) + return + + +if __name__ == '__main__': + args = parse_arguments() + generate(**vars(args)) diff --git a/examples/qwen/smoothquant.py b/examples/qwen/smoothquant.py new file mode 100644 index 000000000000..bdbaddd3435a --- /dev/null +++ b/examples/qwen/smoothquant.py @@ -0,0 +1,209 @@ +# 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. +''' +Utilities for SmoothQuant models +''' + +import functools +import os +import sys +from collections import defaultdict + +import numpy as np +import torch +import torch.nn as nn +from tqdm import tqdm +from transformers.pytorch_utils import Conv1D + +project_dir = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +sys.path.append(project_dir) +from utils.utils import make_context + + +@torch.no_grad() +def apply_smoothing(scales, + gemm_weights, + rmsnorm_weights=None, + dtype=torch.float32, + rmsnorm_1p=False): + if not isinstance(gemm_weights, list): + gemm_weights = [gemm_weights] + + if rmsnorm_weights is not None: + assert rmsnorm_weights.numel() == scales.numel() + rmsnorm_weights.div_(scales).to(dtype) + if rmsnorm_1p: + rmsnorm_weights += (1 / scales) - 1 + + for gemm in gemm_weights: + gemm.mul_(scales.view(1, -1)).to(dtype) + + +@torch.no_grad() +def smooth_gemm(gemm_weights, + act_scales, + rmsnorm_weights=None, + alpha=0.5, + weight_scales=None): + if not isinstance(gemm_weights, list): + gemm_weights = [gemm_weights] + orig_dtype = gemm_weights[0].dtype + + for gemm in gemm_weights: + # gemm_weights are expected to be transposed + assert gemm.shape[1] == act_scales.numel() + + if weight_scales is None: + weight_scales = torch.cat( + [gemm.abs().max(dim=0, keepdim=True)[0] for gemm in gemm_weights], + dim=0) + weight_scales = weight_scales.max(dim=0)[0] + weight_scales.to(float).clamp(min=1e-5) + scales = (act_scales.to(gemm_weights[0].device).to(float).pow(alpha) / + weight_scales.pow(1 - alpha)).clamp(min=1e-5) + + apply_smoothing(scales, gemm_weights, rmsnorm_weights, orig_dtype) + + return scales + + +@torch.no_grad() +def smooth_gemm_mlp(w1_weights, + w2_weights, + act_scales, + rmsnorm_weights=None, + alpha=0.5, + weight_scales=None): + gemm_weights = [] + if not isinstance(w1_weights, list): + w1_weights = [w1_weights] + if not isinstance(w2_weights, list): + w2_weights = [w2_weights] + + for i in range(len(w1_weights)): + gemm_weight = torch.cat([w1_weights[i], w2_weights[i]], dim=0) + gemm_weights.append(gemm_weight) + + orig_dtype = gemm_weights[0].dtype + + for gemm in gemm_weights: + # gemm_weights are expected to be transposed + assert gemm.shape[1] == act_scales.numel() + + if weight_scales is None: + weight_scales = torch.cat( + [gemm.abs().max(dim=0, keepdim=True)[0] for gemm in gemm_weights], + dim=0) + weight_scales = weight_scales.max(dim=0)[0] + weight_scales.to(float).clamp(min=1e-5) + scales = (act_scales.to(gemm_weights[0].device).to(float).pow(alpha) / + weight_scales.pow(1 - alpha)).clamp(min=1e-5) + + apply_smoothing(scales, w1_weights + w2_weights, rmsnorm_weights, + orig_dtype) + + return scales + + +@torch.no_grad() +def smooth_ln_fcs(ln, fcs, act_scales, alpha=0.5): + if not isinstance(fcs, list): + fcs = [fcs] + for fc in fcs: + assert isinstance(fc, nn.Linear) + assert ln.weight.numel() == fc.in_features == act_scales.numel() + + device, dtype = fcs[0].weight.device, fcs[0].weight.dtype + act_scales = act_scales.to(device=device, dtype=dtype) + weight_scales = torch.cat( + [fc.weight.abs().max(dim=0, keepdim=True)[0] for fc in fcs], dim=0) + weight_scales = weight_scales.max(dim=0)[0].clamp(min=1e-5) + + scales = (act_scales.pow(alpha) / + weight_scales.pow(1 - alpha)).clamp(min=1e-5).to(device).to(dtype) + + if ln is not None: + ln.weight.div_(scales) + ln.bias.div_(scales) + + for fc in fcs: + fc.weight.mul_(scales.view(1, -1)) + return scales + + +@torch.no_grad() +def capture_activation_range( + model, + tokenizer, + dataset, + system_prompt, + chat_format, + max_input_len, + num_samples=512, +): + model.eval() + device = next(model.parameters()).device + act_scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) + + def stat_tensor(name, tensor, act_scales, key): + hidden_dim = tensor.shape[-1] + tensor = tensor.view(-1, hidden_dim).abs().detach() + comming_max = torch.max(tensor, dim=0)[0].float() + + if act_scales[name][key] is None: + act_scales[name][key] = comming_max + else: + act_scales[name][key] = torch.max(act_scales[name][key], + comming_max) + + def stat_input_hook(m, x, y, name): + if isinstance(x, tuple): + x = x[0] + stat_tensor(name, x, act_scales, "x") + stat_tensor(name, y, act_scales, "y") + + if act_scales[name]["w"] is None: + act_scales[name]["w"] = m.weight.abs().clip(1e-8, + None).max(dim=1)[0] + + hooks = [] + for name, m in model.named_modules(): + if isinstance(m, nn.Linear) or isinstance(m, Conv1D): + hooks.append( + m.register_forward_hook( + functools.partial(stat_input_hook, name=name))) + num_samples = min(num_samples, len(dataset)) + for i in tqdm(range(num_samples), desc="calibrating model"): + line = dataset[i]["article"] + line = line + ' TL;DR: ' + line = line.strip() + line = line.replace(" n't", "n't") + # use make_content to generate prompt + _, input_id_list = make_context(tokenizer=tokenizer, + query=line, + history=[], + system=system_prompt, + chat_format=chat_format, + max_input_length=max_input_len) + line_encoded = torch.from_numpy(np.array( + input_id_list, dtype=np.int32)).type(torch.int32).unsqueeze(0) + line_encoded = line_encoded.to(device) + model(line_encoded) + + for h in hooks: + h.remove() + + return act_scales diff --git a/examples/llama/summarize.py b/examples/qwen/summarize.py similarity index 58% rename from examples/llama/summarize.py rename to examples/qwen/summarize.py index 9fc883c1ded4..cb47b6a5f9d1 100644 --- a/examples/llama/summarize.py +++ b/examples/qwen/summarize.py @@ -20,69 +20,83 @@ import numpy as np import torch from datasets import load_dataset, load_metric -from transformers import AutoModelForCausalLM, LlamaTokenizer +from run import QWenForCausalLMGenerationSession +from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig +from utils.utils import get_stop_words_ids, make_context import tensorrt_llm import tensorrt_llm.profiler as profiler from tensorrt_llm.logger import logger from tensorrt_llm.quantization import QuantMode +from tensorrt_llm.runtime import ModelConfig from build import get_engine_name # isort:skip +now_dir = os.path.dirname(os.path.abspath(__file__)) -def TRTLLaMA(args, config): +MAX_INPUT_LEN = 2048 +MAX_NEW_TOKENS = 2048 +MAX_SEQ_LEN = 4096 + +TRT_MAX_BATCH_SIZE = 2 +TEMPERATURE = 1.0 +TOP_P = 0.5 +TOP_K = 1 + + +def TRT_QWen(args, config): + use_gpt_attention_plugin = config['plugin_config']['gpt_attention_plugin'] + remove_input_padding = config['plugin_config']['remove_input_padding'] dtype = config['builder_config']['precision'] tp_size = config['builder_config']['tensor_parallel'] pp_size = config['builder_config']['pipeline_parallel'] 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['builder_config']['num_heads'] // tp_size - hidden_size = config['builder_config']['hidden_size'] // tp_size + num_heads = config['builder_config']['num_heads'] // world_size + hidden_size = config['builder_config']['hidden_size'] // world_size vocab_size = config['builder_config']['vocab_size'] num_layers = config['builder_config']['num_layers'] - use_gpt_attention_plugin = bool( - config['plugin_config']['gpt_attention_plugin']) - remove_input_padding = config['plugin_config']['remove_input_padding'] num_kv_heads = config['builder_config'].get('num_kv_heads', num_heads) paged_kv_cache = config['plugin_config']['paged_kv_cache'] tokens_per_block = config['plugin_config']['tokens_per_block'] - use_custom_all_reduce = config['plugin_config'].get('use_custom_all_reduce', - False) - quant_mode = QuantMode(config['builder_config']['quant_mode']) if config['builder_config'].get('multi_query_mode', False): tensorrt_llm.logger.warning( "`multi_query_mode` config is deprecated. Please rebuild the engine." ) num_kv_heads = 1 - num_kv_heads = (num_kv_heads + tp_size - 1) // tp_size - - model_config = tensorrt_llm.runtime.ModelConfig( - vocab_size=vocab_size, - num_layers=num_layers, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - hidden_size=hidden_size, - paged_kv_cache=paged_kv_cache, - tokens_per_block=tokens_per_block, - gpt_attention_plugin=use_gpt_attention_plugin, - remove_input_padding=remove_input_padding, - use_custom_all_reduce=use_custom_all_reduce, - dtype=dtype, - quant_mode=quant_mode) + use_custom_all_reduce = config['plugin_config'].get('use_custom_all_reduce', + False) runtime_rank = tensorrt_llm.mpi_rank() - runtime_mapping = tensorrt_llm.Mapping(world_size, - runtime_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) - engine_name = get_engine_name('llama', dtype, tp_size, pp_size, - runtime_rank) + model_config = ModelConfig(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, + paged_kv_cache=paged_kv_cache, + tokens_per_block=tokens_per_block, + remove_input_padding=remove_input_padding, + dtype=dtype, + quant_mode=quant_mode, + use_custom_all_reduce=use_custom_all_reduce) + + 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) + + engine_name = get_engine_name('qwen', dtype, tp_size, pp_size, runtime_rank) serialize_path = os.path.join(args.engine_dir, engine_name) tensorrt_llm.logger.set_level(args.log_level) @@ -90,9 +104,8 @@ def TRTLLaMA(args, config): profiler.start('load tensorrt_llm engine') with open(serialize_path, 'rb') as f: engine_buffer = f.read() - decoder = tensorrt_llm.runtime.GenerationSession(model_config, - engine_buffer, - runtime_mapping) + decoder = QWenForCausalLMGenerationSession(model_config, engine_buffer, + runtime_mapping) profiler.stop('load tensorrt_llm engine') tensorrt_llm.logger.info( f'Load engine takes: {profiler.elapsed_time_in_sec("load tensorrt_llm engine")} sec' @@ -104,47 +117,63 @@ def main(args): 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 - hf_model_location = args.hf_model_location + test_trt_llm = False + test_hf = False + if args.backend == 'trt_llm': + test_trt_llm = True + elif args.backend == "hf": + test_hf = runtime_rank == 0 # only run hf on rank 0 + else: + raise Exception("unknown backend, only support trt_llm and hf.") profiler.start('load tokenizer') - tokenizer = LlamaTokenizer.from_pretrained(hf_model_location, - legacy=False, - padding_side='left') + tokenizer = AutoTokenizer.from_pretrained( + args.tokenizer_dir, + legacy=False, + padding_side='left', + trust_remote_code=True, + ) profiler.stop('load tokenizer') tensorrt_llm.logger.info( f'Load tokenizer takes: {profiler.elapsed_time_in_sec("load tokenizer")} sec' ) tokenizer.pad_token = tokenizer.eos_token - - dataset_cnn = load_dataset("ccdv/cnn_dailymail", - '3.0.0', - cache_dir=args.dataset_path) + dataset_cnn = load_dataset("ccdv/cnn_dailymail", '3.0.0') + gen_config_path = os.path.join(args.tokenizer_dir, 'generation_config.json') + with open(gen_config_path, 'r') as f: + gen_config = json.load(f) + chat_format = gen_config['chat_format'] max_batch_size = args.batch_size # runtime parameters - # repetition_penalty = 1 - top_k = args.top_k - output_len = 100 - test_token_num = 923 - # top_p = 0.0 - # random_seed = 5 - temperature = 1 + top_p = TOP_K + top_k = TOP_P + temperature = TEMPERATURE + max_new_tokens = MAX_NEW_TOKENS + max_input_len = MAX_INPUT_LEN + max_output_len = MAX_SEQ_LEN num_beams = args.num_beams - pad_id = tokenizer.encode(tokenizer.pad_token, add_special_tokens=False)[0] - end_id = tokenizer.encode(tokenizer.eos_token, add_special_tokens=False)[0] + tokenizer.pad_token_id = pad_id = end_id = tokenizer.im_end_id + # use this prompt to make chat model do summarize + system_prompt = "You are a useful assistant, please directly output the corresponding summary according to the article entered by the user." if test_trt_llm: config_path = os.path.join(args.engine_dir, 'config.json') with open(config_path, 'r') as f: config = json.load(f) - tensorrt_llm_llama = TRTLLaMA(args, config) + + tensorrt_llm_qwen = TRT_QWen(args, config) if test_hf: profiler.start('load HF model') - model = AutoModelForCausalLM.from_pretrained(hf_model_location) + model = AutoModelForCausalLM.from_pretrained( + args.hf_model_dir, + device_map='auto', + trust_remote_code=True, + ) + model.generation_config = GenerationConfig.from_pretrained( + args.hf_model_dir, trust_remote_code=True) profiler.stop('load HF model') tensorrt_llm.logger.info( f'Load HF model takes: {profiler.elapsed_time_in_sec("load HF model")} sec' @@ -155,8 +184,7 @@ def main(args): def summarize_tensorrt_llm(datapoint): batch_size = len(datapoint['article']) - assert batch_size > 0, f"Validation dataset is corrupt (0 samples found). The dataset is loaded from ~/.cache/huggingface/datasets/ccdv___cnn_dailymail" - + assert batch_size > 0 line = copy.copy(datapoint['article']) line_encoded = [] input_lengths = [] @@ -165,20 +193,24 @@ def summarize_tensorrt_llm(datapoint): line[i] = line[i].strip() line[i] = line[i].replace(" n't", "n't") - - input_id = tokenizer.encode(line[i], - return_tensors='pt').type(torch.int32) - input_id = input_id[:, -test_token_num:] + # use make_content to generate prompt + _, input_id_list = make_context( + tokenizer=tokenizer, + query=line[i], + history=[], + system=system_prompt, + max_input_length=max_input_len, + ) + input_id = torch.from_numpy(np.array( + input_id_list, dtype=np.int32)).type(torch.int32).unsqueeze(0) line_encoded.append(input_id) input_lengths.append(input_id.shape[-1]) # do padding, should move outside the profiling to prevent the overhead max_length = max(input_lengths) - if tensorrt_llm_llama.remove_input_padding: - line_encoded = [ - torch.tensor(t, dtype=torch.int32).cuda() for t in line_encoded - ] + if tensorrt_llm_qwen.remove_input_padding: + line_encoded = [torch.IntTensor(t).cuda() for t in line_encoded] else: # do padding, should move outside the profiling to prevent the overhead for i in range(batch_size): @@ -186,28 +218,32 @@ def summarize_tensorrt_llm(datapoint): pad = torch.ones([1, pad_size]).type(torch.int32) * pad_id line_encoded[i] = torch.cat( - [torch.tensor(line_encoded[i], dtype=torch.int32), pad], - axis=-1) + [torch.IntTensor(line_encoded[i]), pad], axis=-1) line_encoded = torch.cat(line_encoded, axis=0).cuda() - input_lengths = torch.tensor(input_lengths, - dtype=torch.int32).cuda() + input_lengths = torch.IntTensor(input_lengths).type( + torch.int32).cuda() sampling_config = tensorrt_llm.runtime.SamplingConfig( - end_id=end_id, pad_id=pad_id, top_k=top_k, num_beams=num_beams) + end_id=end_id, + pad_id=pad_id, + top_k=top_k, + top_p=top_p, + temperature=temperature, + num_beams=num_beams) with torch.no_grad(): - tensorrt_llm_llama.setup(batch_size, - max_context_length=max_length, - max_new_tokens=output_len, - beam_width=num_beams, - max_kv_cache_length=args.max_kv_cache_len) - - if tensorrt_llm_llama.remove_input_padding: - output_ids = tensorrt_llm_llama.decode_batch( + tensorrt_llm_qwen.setup( + batch_size, + max_context_length=max_length, + max_new_tokens=min(max_new_tokens, max_output_len - max_length), + ) + + if tensorrt_llm_qwen.remove_input_padding: + output_ids = tensorrt_llm_qwen.decode_batch( line_encoded, sampling_config) else: - output_ids = tensorrt_llm_llama.decode( + output_ids = tensorrt_llm_qwen.decode( line_encoded, input_lengths, sampling_config, @@ -216,7 +252,7 @@ def summarize_tensorrt_llm(datapoint): torch.cuda.synchronize() # Extract a list of tensors of shape beam_width x output_ids. - if tensorrt_llm_llama.mapping.is_first_pp_rank(): + if tensorrt_llm_qwen.mapping.is_first_pp_rank(): output_beams_list = [ tokenizer.batch_decode(output_ids[batch_idx, :, input_lengths[batch_idx]:], @@ -228,42 +264,71 @@ def summarize_tensorrt_llm(datapoint): def summarize_hf(datapoint): batch_size = len(datapoint['article']) + assert batch_size > 0 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}" ) line = copy.copy(datapoint['article']) - for i in range(batch_size): - line[i] = line[i] + ' TL;DR: ' - line[i] = line[i].strip() - line[i] = line[i].replace(" n't", "n't") - - line_encoded = tokenizer(line, - return_tensors='pt', - padding=True, - truncation=True)["input_ids"].type(torch.int64) + new_line_list = [] + if batch_size > 1: + for i in range(batch_size): + line[i] = line[i] + ' TL;DR: ' + + line[i] = line[i].strip() + line[i] = line[i].replace(" n't", "n't") + # use make_content to generate prompt + raw_text, _ = make_context(tokenizer=tokenizer, + query=line[i], + history=[], + system=system_prompt, + chat_format=chat_format, + max_input_length=max_input_len) + new_line_list.append(raw_text) + line_encoded = tokenizer( + new_line_list, + return_tensors='pt', + padding=True, + truncation=True, + )["input_ids"].type(torch.int64) + else: + line[0] = line[0] + ' TL;DR: ' + line[0] = line[0].strip() + line[0] = line[0].replace(" n't", "n't") + # use make_content to generate prompt + _, input_id_list = make_context(tokenizer=tokenizer, + query=line[0], + history=[], + system=system_prompt, + chat_format=chat_format, + max_input_length=max_input_len) + line_encoded = torch.from_numpy( + np.array(input_id_list, + dtype=np.int64)).type(torch.int64).unsqueeze(0) - line_encoded = line_encoded[:, -test_token_num:] line_encoded = line_encoded.cuda() + stop_words_ids = [] + stop_words_ids.extend(get_stop_words_ids(chat_format, tokenizer)) with torch.no_grad(): - output = model.generate(line_encoded, - max_length=len(line_encoded[0]) + - output_len, - top_k=top_k, - temperature=temperature, - eos_token_id=tokenizer.eos_token_id, - pad_token_id=tokenizer.pad_token_id, - num_beams=num_beams, - num_return_sequences=num_beams, - early_stopping=True) - - tokens_list = output[:, len(line_encoded[0]):].tolist() + output = model.generate( + line_encoded, + max_new_tokens=min(max_new_tokens, + max_output_len - line_encoded.shape[-1]), + top_k=top_k, + top_p=top_p, + do_sample=True, + temperature=temperature, + stop_words_ids=stop_words_ids, + num_beams=num_beams, + num_return_sequences=num_beams, + early_stopping=True) + tokens_list = output[:, line_encoded.shape[-1]:].tolist() output = output.reshape([batch_size, num_beams, -1]) output_lines_list = [ - tokenizer.batch_decode(output[:, i, len(line_encoded[0]):], + tokenizer.batch_decode(output[:, i, line_encoded.shape[-1]:], skip_special_tokens=True) for i in range(num_beams) ] @@ -293,8 +358,10 @@ def summarize_hf(datapoint): logger.info(f"\n Summary : {summary}") logger.info("---------------------------------------------------------") + print("load rouge ...") metric_tensorrt_llm = [load_metric("rouge") for _ in range(num_beams)] metric_hf = [load_metric("rouge") for _ in range(num_beams)] + print("load rouge done") for i in range(num_beams): metric_tensorrt_llm[i].seed = 0 metric_hf[i].seed = 0 @@ -364,8 +431,8 @@ def summarize_hf(datapoint): ) if args.check_accuracy and beam_idx == 0: - assert computed_metrics_tensorrt_llm['rouge1'].mid[ - 2] * 100 > args.tensorrt_llm_rouge1_threshold + assert computed_metrics_tensorrt_llm[ + 'rouge1'] * 100 > args.tensorrt_llm_rouge1_threshold if test_hf: np.random.seed(0) # rouge score use sampling to compute the score logger.info( @@ -381,24 +448,34 @@ def summarize_hf(datapoint): if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('--hf_model_location', - type=str, - default='/workspace/models/llama-models/llama-7b-hf') - parser.add_argument('--test_hf', action='store_true') - parser.add_argument('--test_trt_llm', action='store_true') + + parser.add_argument( + "--backend", + type=str, + choices=["trt_llm", "hf"], + default="hf", + ) + parser.add_argument( + '--hf_model_dir', + type=str, + default=".", + ) + parser.add_argument( + "--tokenizer_dir", + type=str, + default=".", + ) + parser.add_argument( + '--engine_dir', + type=str, + default="qwen_outputs", + ) parser.add_argument('--data_type', type=str, choices=['fp32', 'fp16'], default='fp16') - parser.add_argument('--dataset_path', type=str, default='') - parser.add_argument('--max_kv_cache_len', - type=int, - default=None, - help='The max kv cache length. \ - If the final sequence length exceeds the kv cache length, we will enable cyclic kv cache. \ - If it is set to None, we will use the max sequence length.') + parser.add_argument('--dataset_path', type=str, default="") parser.add_argument('--log_level', type=str, default='info') - parser.add_argument('--engine_dir', type=str, default='llama_outputs') parser.add_argument('--batch_size', type=int, default=1) parser.add_argument('--max_ite', type=int, default=20) parser.add_argument('--check_accuracy', action='store_true') @@ -406,8 +483,10 @@ def summarize_hf(datapoint): type=float, default=15.0) parser.add_argument('--num_beams', type=int, default=1) - parser.add_argument('--top_k', type=int, default=1) - + parser.add_argument("--max_new_tokens", + type=int, + default=100, + help="Maximum number of new tokens to generate.") args = parser.parse_args() main(args) diff --git a/tensorrt_llm/models/internlm/__init__.py b/examples/qwen/utils/__init__.py similarity index 100% rename from tensorrt_llm/models/internlm/__init__.py rename to examples/qwen/utils/__init__.py diff --git a/examples/qwen/utils/convert.py b/examples/qwen/utils/convert.py new file mode 100644 index 000000000000..62cfa810bb00 --- /dev/null +++ b/examples/qwen/utils/convert.py @@ -0,0 +1,304 @@ +# 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. +""" + Utilities for exporting a model to our custom format. +""" + +import numpy as np +import torch + +from tensorrt_llm._utils import torch_to_numpy + + +def cpu_map_location(storage, loc): + return storage.cpu() + + +def gpu_map_location(storage, loc): + if loc.startswith("cuda"): + training_gpu_idx = int(loc.split(":")[1]) + inference_gpu_idx = training_gpu_idx % torch.cuda.device_count() + return storage.cuda(inference_gpu_idx) + elif loc.startswith("cpu"): + return storage.cpu() + else: + raise ValueError(f"Not handled {loc}") + + +def save_val(val, dir, key, tp_num=None): + suffix = "bin" if tp_num is None else f"{tp_num}.bin" + val.tofile(dir / f"model.{key}.{suffix}") + + +def save_split(split_vals, dir, key, i, split_factor): + for j, val in enumerate(split_vals): + save_val(val, dir, key, i * split_factor + j) + + +def generate_int8(weights, act_range, is_qkv=False, multi_query_mode=False): + """ + This function has two purposes: + - compute quantized weights, scaled either per-tensor or per-column + - compute scaling factors + + Depending on the GEMM API (CUTLASS/CUBLAS) the required scaling factors differ. + CUTLASS uses two sets of scaling factors. One for the activation X, one for the weight W. + CUBLAS only has one (we can't do per-row scaling). So we must provide pre-multiplied scaling factor. + + Here is the list of what we need (T means per-tensor, C per-column): + - scale_x_orig_quant puts fp activation into the quantized range (i.e. [-128, 127], for int8). Used before the GEMM. (T) + - scale_y_quant_orig puts quantized activation into the fp range. Used if the GEMM outputs int8. (T) + - scale_w_quant_orig puts weights from quant range to fp range (used with CUTLASS) (T, C) + - scale_y_accum_quant puts the GEMM result (XW) from accumulation range (int32) + to quant range (int8) (used for CUBLAS) (T, C) + + Note that we don't do anything special about row-parallel GEMM. Theoretically, we could have per-GPU scaling factors too, + but then the model would change depending on the number of GPUs used. + + For QKV projection, the behavior is special. Even if we have a single matrix to perform QKV projection, we consider it + as three different matrices: Q, K, and V. So per-tensor actually means one scaling factor for each Q, K and V. + """ + + # compute weight scaling factors for fp->int8 and int8->fp + if is_qkv and not multi_query_mode: + scale_w_orig_quant_t = 127. / torch_to_numpy(act_range["w"].reshape( + 3, -1).max(dim=-1, keepdims=True)[0].cpu()).astype(np.float32) + scale_w_orig_quant_c = 127. / torch_to_numpy(act_range["w"].reshape( + 3, -1).cpu()).astype(np.float32) + elif is_qkv and multi_query_mode: + raise ValueError( + f"Multi-query w/ int8 quant has not been supported yet") + else: + scale_w_orig_quant_t = 127. / torch_to_numpy( + act_range["w"].max().cpu()).astype(np.float32) + scale_w_orig_quant_c = 127. / torch_to_numpy( + act_range["w"].cpu()).astype(np.float32) + scale_w_quant_orig_t = 1.0 / scale_w_orig_quant_t + scale_w_quant_orig_c = 1.0 / scale_w_orig_quant_c + + # compute the rest of needed scaling factors + scale_x_orig_quant_t = np.array(127. / act_range["x"].max().item()) + scale_y_orig_quant_t = np.array(127. / act_range["y"].max().item()) + scale_y_quant_orig_t = np.array(act_range["y"].max().item() / 127.) + scale_y_accum_quant_t = scale_y_orig_quant_t / (scale_x_orig_quant_t * + scale_w_orig_quant_t) + scale_y_accum_quant_c = scale_y_orig_quant_t / (scale_x_orig_quant_t * + scale_w_orig_quant_c) + if is_qkv: + scale_y_accum_quant_t = np.broadcast_to(scale_y_accum_quant_t, + scale_w_orig_quant_c.shape) + scale_w_quant_orig_t = np.broadcast_to(scale_w_quant_orig_t, + scale_w_orig_quant_c.shape) + + to_i8 = lambda x: x.round().clip(-127, 127).astype(np.int8) + return { + "weight.int8": to_i8(weights * scale_w_orig_quant_t), + "weight.int8.col": to_i8(weights * scale_w_orig_quant_c), + "scale_x_orig_quant": scale_x_orig_quant_t.astype(np.float32), + "scale_w_quant_orig": scale_w_quant_orig_t.astype(np.float32), + "scale_w_quant_orig.col": scale_w_quant_orig_c.astype(np.float32), + "scale_y_accum_quant": scale_y_accum_quant_t.astype(np.float32), + "scale_y_accum_quant.col": scale_y_accum_quant_c.astype(np.float32), + "scale_y_quant_orig": scale_y_quant_orig_t.astype(np.float32), + } + + +def write_int8(vals, + dir, + base_key, + split_dim, + tp_rank, + split_factor, + kv_cache_only=False): + if not kv_cache_only: + save_split(np.split(vals["weight.int8"], split_factor, axis=split_dim), + dir, f"{base_key}.weight.int8", tp_rank, split_factor) + save_split( + np.split(vals["weight.int8.col"], split_factor, axis=split_dim), + dir, f"{base_key}.weight.int8.col", tp_rank, split_factor) + + saved_keys_once = ["scale_y_quant_orig"] + if not kv_cache_only: + saved_keys_once += [ + "scale_x_orig_quant", "scale_w_quant_orig", "scale_y_accum_quant" + ] + # per-column scaling factors are loaded per-gpu for ColumnParallel GEMMs (QKV, FC1) + if not kv_cache_only: + if split_dim == -1: + save_split( + np.split(vals["scale_w_quant_orig.col"], + split_factor, + axis=split_dim), dir, + f"{base_key}.scale_w_quant_orig.col", tp_rank, split_factor) + save_split( + np.split(vals["scale_y_accum_quant.col"], + split_factor, + axis=split_dim), dir, + f"{base_key}.scale_y_accum_quant.col", tp_rank, split_factor) + else: + saved_keys_once += [ + "scale_w_quant_orig.col", "scale_y_accum_quant.col" + ] + + if tp_rank == 0: + for save_key in saved_keys_once: + save_val(vals[save_key], dir, f"{base_key}.{save_key}") + + +# Note: in multi_query_mode, only query heads are split between multiple GPUs, while key/value head +# are not split as there is only one head per key/value. +@torch.no_grad() +def split_and_save_weight(tp_rank, saved_dir, split_factor, key, vals, + storage_type, act_range, config): + use_attention_nemo_shape = config.get("use_attention_nemo_shape", False) + split_gated_activation = config.get("split_gated_activation", False) + num_attention_heads = config.get("num_attention_heads", 0) + tp_size = config.get("tp_size", 1) + int8_outputs = config.get("int8_outputs", None) + multi_query_mode = config.get("multi_query_mode", False) + local_dim = config.get("local_dim", None) + + save_int8 = int8_outputs == "all" or int8_outputs == "kv_cache_only" + + if not key.endswith(".smoother"): + if not isinstance(vals, list): + vals = [vals] + + if config.get("transpose_weights", False) and vals[0].ndim == 2: + vals = [val.T for val in vals] + if "layernorm.weight" in key and config.get("apply_layernorm_1p", + False): + vals = [val + 1.0 for val in vals] + vals = [torch_to_numpy(val.cpu().to(storage_type)) for val in vals] + else: + vals = torch_to_numpy(vals.cpu()) + + if "ln_1.weight" in key or "ln_1.bias" in key or \ + "attention.dense.bias" in key or \ + "ln_2.weight" in key or "ln_2.bias" in key or \ + "mlp.c_proj.bias" in key or "ln_f.weight" in key: + # "final_layernorm.weight" in key or "final_layernorm.bias" in key: + + # shared weights, only need to convert the weights of rank 0 + if tp_rank == 0: + save_val(vals[0], saved_dir, key) + + elif "attention.dense.weight" in key or "mlp.c_proj.weight" in key: + cat_dim = 0 + val = np.concatenate(vals, axis=cat_dim) + split_vals = np.split(val, split_factor, axis=cat_dim) + save_split(split_vals, saved_dir, key, tp_rank, split_factor) + if act_range is not None and int8_outputs == "all": + base_key = key.replace(".weight", "") + vals_i8 = generate_int8(val, + act_range, + multi_query_mode=multi_query_mode) + write_int8(vals_i8, saved_dir, base_key, cat_dim, tp_rank, + split_factor) + + elif "mlp.w1.weight" in key or "mlp.w2.weight" in key or "mlp.w1.bias" in key or "mlp.w2.bias" in key: + if split_gated_activation: + splits = [np.split(val, 2, axis=-1) for val in vals] + vals, gates = list(zip(*splits)) + cat_dim = -1 + val = np.concatenate(vals, axis=cat_dim) + split_vals = np.split(val, split_factor, axis=cat_dim) + save_split(split_vals, saved_dir, key, tp_rank, split_factor) + if act_range is not None and int8_outputs == "all": + base_key = key.replace(".weight", "") + vals_i8 = generate_int8(val, + act_range, + multi_query_mode=multi_query_mode) + write_int8(vals_i8, saved_dir, base_key, cat_dim, tp_rank, + split_factor) + + if split_gated_activation: + assert not save_int8 + prefix, dot, suffix = key.rpartition(".") + key = prefix + ".gate" + dot + suffix + + gate = np.concatenate(gates, axis=cat_dim) + split_vals = np.split(gate, split_factor, axis=cat_dim) + save_split(split_vals, saved_dir, key, tp_rank, split_factor) + + elif "attention.qkv.bias" in key: + if local_dim is None: + local_dim = vals[0].shape[-1] // 3 + + if multi_query_mode: + val = vals[0] + # out_feature = local_dim + 2 * head_size; assumes local_dim equals to hidden_dim + b_q, b_kv = np.split(val, [local_dim], axis=-1) + b_q_split = np.split(b_q, split_factor, axis=-1) + split_vals = [np.concatenate((i, b_kv), axis=-1) for i in b_q_split] + else: + if use_attention_nemo_shape: + head_num = num_attention_heads // tp_size + size_per_head = local_dim // num_attention_heads + nemo_shape = (head_num, 3, size_per_head) + vals = [val.reshape(nemo_shape) for val in vals] + vals = [val.transpose(1, 0, 2) for val in vals] + + vals = [val.reshape(3, local_dim) for val in vals] + val = np.concatenate(vals, axis=-1) + split_vals = np.split(val, split_factor, axis=-1) + save_split(split_vals, saved_dir, key, tp_rank, split_factor) + + elif "attention.qkv.weight" in key: + hidden_dim = vals[0].shape[0] + if local_dim is None: + local_dim = vals[0].shape[-1] // 3 + if multi_query_mode: + val = vals[0] + # out_feature = local_dim + 2 * head_size; assumes local_dim equals to hidden_dim + head_size = (val.shape[-1] - local_dim) // 2 + val = val.reshape(hidden_dim, local_dim + 2 * head_size) + w_q, w_kv = np.split(val, [local_dim], axis=-1) + w_q_split = np.split(w_q, split_factor, axis=-1) + split_vals = [np.concatenate((i, w_kv), axis=-1) for i in w_q_split] + else: + if use_attention_nemo_shape: + head_num = num_attention_heads // tp_size + size_per_head = hidden_dim // num_attention_heads + vals = [ + val.reshape(hidden_dim, head_num, 3, size_per_head) + for val in vals + ] + vals = [val.transpose(0, 2, 1, 3) for val in vals] + + vals = [val.reshape(hidden_dim, 3, local_dim) for val in vals] + cat_dim = -1 + val = np.concatenate(vals, axis=cat_dim) + split_vals = np.split(val, split_factor, axis=cat_dim) + save_split(split_vals, saved_dir, key, tp_rank, split_factor) + if save_int8: + base_key = key.replace(".weight", "") + vals_i8 = generate_int8(val, + act_range, + is_qkv=True, + multi_query_mode=multi_query_mode) + write_int8(vals_i8, + saved_dir, + base_key, + cat_dim, + tp_rank, + split_factor, + kv_cache_only=int8_outputs == "kv_cache_only") + + elif "attention.dense.smoother" in key or "mlp.c_proj.smoother" in key: + split_vals = np.split(vals, split_factor, axis=0) + save_split(split_vals, saved_dir, key, tp_rank, split_factor) + else: + print(f"[WARNING] {key} not handled by converter") diff --git a/examples/qwen/utils/utils.py b/examples/qwen/utils/utils.py new file mode 100644 index 000000000000..894c5f2271c1 --- /dev/null +++ b/examples/qwen/utils/utils.py @@ -0,0 +1,134 @@ +# 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. +from typing import List, Tuple + +from transformers import PreTrainedTokenizer + + +def make_context( + tokenizer: PreTrainedTokenizer, + query: str, + history: List[Tuple[str, str]] = None, + system: str = "You are a helpful assistant.", + max_input_length: + int = 2048, # if you want to change this, you need to change the max_input_len in tensorrt_llm_july-release-v1/examples/qwen/build.py + max_window_size: int = 6144, + chat_format: str = "chatml", +): + if history is None: + history = [] + + if chat_format == "chatml": + im_start, im_end = "<|im_start|>", "<|im_end|>" + im_start_tokens = [tokenizer.im_start_id] + im_end_tokens = [tokenizer.im_end_id] + nl_tokens = tokenizer.encode("\n") + + def _tokenize_str(role, content): + return (f"{role}\n{content}", + tokenizer.encode( + role, + allowed_special=set(), + ) + nl_tokens + tokenizer.encode( + content, + allowed_special=set(), + )) + + 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 + + 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}" + ) + + 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 + + tokenizer.encode("assistant") + nl_tokens) + raw_text += f"\n{im_start}user\n{query}{im_end}\n{im_start}assistant\n" + + elif chat_format == "raw": + raw_text = query + context_tokens = tokenizer.encode(raw_text) + else: + raise NotImplementedError(f"Unknown chat format {chat_format!r}") + # truncate to max_input_length, truncate from the front + return raw_text, context_tokens[-max_input_length:] + + +def _decode_chatml(tokens: List[int], + stop_words: List[str], + eod_token_ids: List[int], + tokenizer: PreTrainedTokenizer, + raw_text_len: int, + context_length: int, + verbose: bool = False, + return_end_reason: bool = False, + errors: str = 'replace'): + end_reason = f"Gen length {len(tokens)}" + eod_token_idx = context_length + for eod_token_idx in range(context_length, len(tokens)): + if tokens[eod_token_idx] in eod_token_ids: + end_reason = f"Gen {tokenizer.decode([tokens[eod_token_idx]])!r}" + break + + trim_decode_tokens = tokenizer.decode(tokens[:eod_token_idx], + errors=errors)[raw_text_len:] + if verbose: + print("\nRaw Generate w/o EOD:", + tokenizer.decode(tokens, errors=errors)[raw_text_len:]) + print("\nRaw Generate:", trim_decode_tokens) + print("\nEnd Reason:", end_reason) + for stop_word in stop_words: + trim_decode_tokens = trim_decode_tokens.replace(stop_word, "").strip() + trim_decode_tokens = trim_decode_tokens.strip() + if verbose: + print("\nGenerate:", trim_decode_tokens) + + if return_end_reason: + return trim_decode_tokens, end_reason + else: + return trim_decode_tokens + + +def get_stop_words_ids(chat_format, tokenizer): + if chat_format == "raw": + stop_words_ids = [tokenizer.encode("Human:"), [tokenizer.eod_id]] + elif chat_format == "chatml": + stop_words_ids = [[tokenizer.im_end_id], [tokenizer.im_start_id]] + else: + raise NotImplementedError(f"Unknown chat format {chat_format!r}") + return stop_words_ids diff --git a/examples/qwen/weight.py b/examples/qwen/weight.py new file mode 100644 index 000000000000..31d8e6afb729 --- /dev/null +++ b/examples/qwen/weight.py @@ -0,0 +1,524 @@ +# 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 configparser +import time +from pathlib import Path + +import numpy as np +import torch +from tqdm import tqdm + +import tensorrt_llm +from tensorrt_llm._utils import (str_dtype_to_np, str_dtype_to_torch, + torch_to_numpy) +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models import QWenForCausalLM +from tensorrt_llm.quantization import QuantMode + + +def gen_suffix(rank, use_smooth_quant, quant_per_channel): + suffix = f"{rank}.bin" + if use_smooth_quant: + sq_prefix = "int8." + if quant_per_channel: + sq_prefix += "col." + suffix = sq_prefix + suffix + return suffix + + +def extract_layer_idx(name): + ss = name.split('.') + for s in ss: + if s.isdigit(): + return s + return None + + +def split(v, tp_size, idx, dim=0): + if tp_size == 1: + return v + if len(v.shape) == 1: + return np.ascontiguousarray(np.split(v, tp_size)[idx]) + else: + return np.ascontiguousarray(np.split(v, tp_size, axis=dim)[idx]) + + +def parse_ft_config(ini_file): + qwen_config = configparser.ConfigParser() + qwen_config.read(ini_file) + + vocab_size = qwen_config.getint('qwen', 'vocab_size') + hidden_size = qwen_config.getint('qwen', 'hidden_size') + inter_size = qwen_config.getint('qwen', 'intermediate_size', fallback=None) + num_hidden_layers = qwen_config.getint( + "qwen", + "num_hidden_layers", + fallback=32, + ) + max_position_embeddings = qwen_config.getint("qwen", + "max_position_embeddings", + fallback=8192) + kv_channels = qwen_config.getint('qwen', 'kv_channels', fallback=128) + rotary_pct = qwen_config.getfloat('qwen', 'rotary_pct', fallback=0.0) + rotary_emb_base = qwen_config.getint('qwen', + 'rotary_emb_base', + fallback=10000) + multi_query_mode = qwen_config.getboolean('qwen', + 'multi_query_mode', + fallback=False) + return (vocab_size, hidden_size, inter_size, num_hidden_layers, kv_channels, + rotary_pct, rotary_emb_base, multi_query_mode, + max_position_embeddings) + + +def load_from_ft(tensorrt_llm_qwen: QWenForCausalLM, + dir_path, + mapping=Mapping(), + dtype='float16', + share_embedding_table=False, + parallel_embedding_table=False, + multi_query_mode=False): + tensorrt_llm.logger.info('Loading weights from FT...') + tik = time.time() + quant_mode = getattr(tensorrt_llm_qwen, 'quant_mode', QuantMode(0)) + if quant_mode.is_int8_weight_only(): + plugin_weight_only_quant_type = torch.int8 + elif quant_mode.is_int4_weight_only(): + plugin_weight_only_quant_type = torch.quint4x2 + (vocab_size, hidden_size, inter_size, num_hidden_layers, kv_channels, + rotary_pct, rotary_emb_base, multi_query_mode, + max_position_embeddings) = parse_ft_config(Path(dir_path) / 'config.ini') + np_dtype = str_dtype_to_np(dtype) + + def fromfile(dir_path, name, shape=None, dtype=np.float16): + dtype = np_dtype if dtype is None else dtype + p = dir_path + '/' + name + if Path(p).exists(): + t = np.fromfile(p, dtype=dtype) + if shape is not None: + t = t.reshape(shape) + return t + else: + print(f"Warning: {p} not found.") + return None + + def set_smoothquant_scale_factors( + module, + pre_scale_weight, + dir_path, + basename, + shape, + per_tok_dyn, + per_channel, + is_qkv=False, + rank=None, + ): + suffix = "bin" + if per_channel: + if rank is not None: + suffix = f"{rank}." + suffix + suffix = "col." + suffix + + col_shape = shape if (per_channel or is_qkv) else [1, 1] + if per_tok_dyn: + if pre_scale_weight is not None: + pre_scale_weight.value = np.array([1.0], dtype=np.float32) + t = fromfile(dir_path, f"{basename}scale_w_quant_orig.{suffix}", + col_shape, np.float32) + module.per_channel_scale.value = t + else: + t = fromfile(dir_path, f"{basename}scale_x_orig_quant.bin", [1], + np.float32) + pre_scale_weight.value = t + t = fromfile(dir_path, f"{basename}scale_y_accum_quant.{suffix}", + col_shape, np.float32) + module.per_channel_scale.value = t + t = fromfile(dir_path, f"{basename}scale_y_quant_orig.bin", [1, 1], + np.float32) + module.act_scale.value = t + + def set_smoother(module, dir_path, base_name, shape, rank): + suffix = f"{rank}.bin" + t = fromfile(dir_path, f"{base_name}.smoother.{suffix}", shape, + np.float32) + module.smoother.value = t + + # Determine the quantization mode. + quant_mode = getattr(tensorrt_llm_qwen, "quant_mode", QuantMode(0)) + # Do we use SmoothQuant? + use_smooth_quant = quant_mode.has_act_and_weight_quant() + # Do we use quantization per token? + quant_per_token_dyn = quant_mode.has_per_token_dynamic_scaling() + # Do we use quantization per channel? + quant_per_channel = quant_mode.has_per_channel_scaling() + + # Do we use INT4/INT8 weight-only? + use_weight_only = quant_mode.is_weight_only() + + # Int8 KV cache + use_int8_kv_cache = quant_mode.has_int8_kv_cache() + + # Debug + suffix = gen_suffix(mapping.tp_rank, use_smooth_quant, quant_per_channel) + # The type of weights. + w_type = np_dtype if not use_smooth_quant else np.int8 + + if mapping.is_first_pp_rank(): + tensorrt_llm_qwen.vocab_embedding.weight.value = (fromfile( + dir_path, 'vocab_embedding.weight.bin', [vocab_size, hidden_size])) + + if mapping.is_last_pp_rank(): + tensorrt_llm_qwen.ln_f.weight.value = (fromfile(dir_path, + 'ln_f.weight.bin')) + + lm_head_weight = fromfile(dir_path, 'lm_head.weight.bin', + [vocab_size, hidden_size]) + + if vocab_size % mapping.tp_size != 0: + # padding + vocab_size_padded = tensorrt_llm_qwen.lm_head.out_features * mapping.tp_size + pad_width = vocab_size_padded - vocab_size + lm_head_weight = np.pad(lm_head_weight, ((0, pad_width), (0, 0)), + 'constant', + constant_values=0) + if mapping.is_last_pp_rank(): + tensorrt_llm_qwen.lm_head.weight.value = np.ascontiguousarray( + split(lm_head_weight, mapping.tp_size, mapping.tp_rank)) + + layers_range = list( + range(mapping.pp_rank * tensorrt_llm_qwen.num_layers, + (mapping.pp_rank + 1) * tensorrt_llm_qwen.num_layers, 1)) + + for i in layers_range: + c_attn_out_dim = (3 * hidden_size // + mapping.tp_size) if not multi_query_mode else ( + hidden_size // mapping.tp_size + + (hidden_size // num_hidden_layers) * 2) + + tensorrt_llm_qwen.layers[i].ln_1.weight.value = fromfile( + dir_path, 'model.layers.' + str(i) + '.ln_1.weight.bin') + + dst = tensorrt_llm_qwen.layers[i].ln_2.weight + dst.value = fromfile(dir_path, + 'model.layers.' + str(i) + '.ln_2.weight.bin') + + t = fromfile( + dir_path, + 'model.layers.' + str(i) + '.attention.qkv.weight.' + suffix, + [hidden_size, c_attn_out_dim], w_type) + if t is not None: + dst = tensorrt_llm_qwen.layers[i].attention.qkv.weight + if use_smooth_quant: + dst.value = np.ascontiguousarray(np.transpose(t, [1, 0])) + set_smoothquant_scale_factors( + tensorrt_llm_qwen.layers[i].attention.qkv, + tensorrt_llm_qwen.layers[i].ln_1.scale_to_int, + dir_path, + 'model.layers.' + str(i) + '.attention.qkv.', + [1, c_attn_out_dim], + quant_per_token_dyn, + quant_per_channel, + rank=mapping.tp_rank, + is_qkv=True) + elif use_weight_only: + processed_torch_weights, torch_weight_scales = torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( + torch.tensor(t), plugin_weight_only_quant_type) + dst.value = processed_torch_weights.numpy() + scales = tensorrt_llm_qwen.layers[ + i].attention.qkv.per_channel_scale + scales.value = torch_weight_scales.numpy() + else: + dst.value = np.ascontiguousarray(np.transpose(t, [1, 0])) + + dst = tensorrt_llm_qwen.layers[i].attention.qkv.bias + t = fromfile( + dir_path, 'model.layers.' + str(i) + '.attention.qkv.bias.' + + str(mapping.tp_rank) + '.bin', [c_attn_out_dim]) + dst.value = np.ascontiguousarray(t) + + dst = tensorrt_llm_qwen.layers[i].attention.dense.weight + t = fromfile( + dir_path, + 'model.layers.' + str(i) + '.attention.dense.weight.' + suffix, + [hidden_size // mapping.tp_size, hidden_size], w_type) + if use_smooth_quant: + dst.value = np.ascontiguousarray(np.transpose(t, [1, 0])) + dense_scale = getattr(tensorrt_llm_qwen.layers[i].attention, + "quantization_scaling_factor", None) + set_smoothquant_scale_factors( + tensorrt_llm_qwen.layers[i].attention.dense, + dense_scale, + dir_path, + 'model.layers.' + str(i) + '.attention.dense.', + [1, hidden_size], + quant_per_token_dyn, + quant_per_channel, + ) + set_smoother(tensorrt_llm_qwen.layers[i].attention.dense, dir_path, + 'model.layers.' + str(i) + '.attention.dense', + [1, hidden_size // mapping.tp_size], mapping.tp_rank) + + elif use_weight_only: + processed_torch_weights, torch_weight_scales = torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( + torch.tensor(t), plugin_weight_only_quant_type) + dst.value = processed_torch_weights.numpy() + scales = tensorrt_llm_qwen.layers[ + i].attention.dense.per_channel_scale + scales.value = torch_weight_scales.numpy() + else: + dst.value = np.ascontiguousarray(np.transpose(t, [1, 0])) + + t = fromfile(dir_path, + 'model.layers.' + str(i) + '.mlp.w1.weight.' + suffix, + [hidden_size, inter_size // mapping.tp_size // 2], w_type) + if use_smooth_quant: + tensorrt_llm_qwen.layers[ + i].mlp.gate.weight.value = np.ascontiguousarray( + np.transpose(t, [1, 0])) + set_smoothquant_scale_factors( + tensorrt_llm_qwen.layers[i].mlp.gate, + tensorrt_llm_qwen.layers[i].ln_2.scale_to_int, + dir_path, + 'model.layers.' + str(i) + '.mlp.w1.', + [1, inter_size // mapping.tp_size // 2], + quant_per_token_dyn, + quant_per_channel, + rank=mapping.tp_rank) + elif use_weight_only: + dst = tensorrt_llm_qwen.layers[i].mlp.gate.weight + processed_torch_weights, torch_weight_scales = torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( + torch.tensor(t), plugin_weight_only_quant_type) + dst.value = processed_torch_weights.numpy() + scales = tensorrt_llm_qwen.layers[i].mlp.gate.per_channel_scale + scales.value = torch_weight_scales.numpy() + else: + tensorrt_llm_qwen.layers[ + i].mlp.gate.weight.value = np.ascontiguousarray( + np.transpose(t, [1, 0])) + + t = fromfile(dir_path, + 'model.layers.' + str(i) + '.mlp.w2.weight.' + suffix, + [hidden_size, inter_size // mapping.tp_size // 2], w_type) + if use_smooth_quant: + tensorrt_llm_qwen.layers[ + i].mlp.fc.weight.value = np.ascontiguousarray( + np.transpose(t, [1, 0])) + set_smoothquant_scale_factors( + tensorrt_llm_qwen.layers[i].mlp.fc, + tensorrt_llm_qwen.layers[i].ln_2.scale_to_int, + dir_path, + 'model.layers.' + str(i) + '.mlp.w2.', + [1, inter_size // mapping.tp_size // 2], + quant_per_token_dyn, + quant_per_channel, + rank=mapping.tp_rank) + elif use_weight_only: + dst = tensorrt_llm_qwen.layers[i].mlp.fc.weight + processed_torch_weights, torch_weight_scales = torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( + torch.tensor(t), plugin_weight_only_quant_type) + dst.value = processed_torch_weights.numpy() + scales = tensorrt_llm_qwen.layers[i].mlp.fc.per_channel_scale + scales.value = torch_weight_scales.numpy() + else: + tensorrt_llm_qwen.layers[ + i].mlp.fc.weight.value = np.ascontiguousarray( + np.transpose(t, [1, 0])) + + t = fromfile(dir_path, + 'model.layers.' + str(i) + '.mlp.c_proj.weight.' + suffix, + [inter_size // mapping.tp_size // 2, hidden_size], w_type) + if use_smooth_quant: + tensorrt_llm_qwen.layers[ + i].mlp.proj.weight.value = np.ascontiguousarray( + np.transpose(t, [1, 0])) + proj_scale = getattr(tensorrt_llm_qwen.layers[i].mlp, + "quantization_scaling_factor", None) + set_smoothquant_scale_factors( + tensorrt_llm_qwen.layers[i].mlp.proj, proj_scale, dir_path, + 'model.layers.' + str(i) + '.mlp.c_proj.', [1, hidden_size], + quant_per_token_dyn, quant_per_channel) + set_smoother(tensorrt_llm_qwen.layers[i].mlp.proj, dir_path, + 'model.layers.' + str(i) + '.mlp.c_proj', + [1, inter_size // mapping.tp_size // 2], + mapping.tp_rank) + elif use_weight_only: + dst = tensorrt_llm_qwen.layers[i].mlp.proj.weight + processed_torch_weights, torch_weight_scales = torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( + torch.tensor(t), plugin_weight_only_quant_type) + dst.value = processed_torch_weights.numpy() + scales = tensorrt_llm_qwen.layers[i].mlp.proj.per_channel_scale + scales.value = torch_weight_scales.numpy() + else: + tensorrt_llm_qwen.layers[ + i].mlp.proj.weight.value = np.ascontiguousarray( + np.transpose(t, [1, 0])) + + if use_int8_kv_cache: + t = fromfile( + dir_path, 'model.layers.' + str(i) + + '.attention.qkv.scale_y_quant_orig.bin', [1], np.float32) + tensorrt_llm_qwen.layers[ + i].attention.kv_orig_quant_scale.value = 1.0 / t + tensorrt_llm_qwen.layers[i].attention.kv_quant_orig_scale.value = t + + tok = time.time() + t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) + tensorrt_llm.logger.info(f'Weights loaded. Total time: {t}') + + +def load_from_hf_qwen(tensorrt_llm_qwen: tensorrt_llm.models.QWenForCausalLM, + hf_qwen, + mapping=Mapping(), + max_position_embeddings=8192, + rotary_emb_base=10000, + kv_channels=128, + dtype="float32", + multi_query_mode=False): + tensorrt_llm.logger.info('Loading weights from HF QWen...') + tik = time.time() + + quant_mode = getattr(tensorrt_llm_qwen, 'quant_mode', QuantMode(0)) + if quant_mode.is_int8_weight_only(): + plugin_weight_only_quant_type = torch.int8 + elif quant_mode.is_int4_weight_only(): + plugin_weight_only_quant_type = torch.quint4x2 + use_weight_only = quant_mode.is_weight_only() + + model_params = dict(hf_qwen.named_parameters()) + torch_dtype = str_dtype_to_torch(dtype) + for k, v in tqdm(model_params.items(), + total=len(model_params), + ncols=80, + desc="Converting..."): + if isinstance(v, list): + v = [torch_to_numpy(vv.to(torch_dtype).detach().cpu()) for vv in v] + else: + v = torch_to_numpy(v.to(torch_dtype).detach().cpu()) + if 'transformer.wte.weight' in k: + tensorrt_llm_qwen.vocab_embedding.weight.value = v + elif 'transformer.ln_f.weight' in k: + tensorrt_llm_qwen.ln_f.weight.value = v + elif 'lm_head.weight' in k: + tensorrt_llm_qwen.lm_head.weight.value = np.ascontiguousarray( + split(v, mapping.tp_size, mapping.tp_rank)) + else: + layer_idx = extract_layer_idx(k) + if layer_idx is None: + continue + idx = int(layer_idx) + if idx >= tensorrt_llm_qwen.num_layers: + continue + if 'ln_1.weight' in k: + tensorrt_llm_qwen.layers[idx].ln_1.weight.value = v + elif 'ln_2.weight' in k: + tensorrt_llm_qwen.layers[idx].ln_2.weight.value = v + elif 'attn.c_attn.weight' in k: + dst = tensorrt_llm_qwen.layers[idx].attention.qkv.weight + if multi_query_mode: + assert isinstance(v, list) and len(v) == 3 + wq = split(v[0], mapping.tp_size, mapping.tp_rank) + wk = split(v[1], mapping.tp_size, mapping.tp_rank) + wv = split(v[2], mapping.tp_size, mapping.tp_rank) + split_v = np.concatenate((wq, wk, wv)) + else: + q_emb = v.shape[0] // 3 + model_emb = v.shape[1] + v = v.reshape(3, q_emb, model_emb) + split_v = split(v, mapping.tp_size, mapping.tp_rank, dim=1) + split_v = split_v.reshape(3 * (q_emb // mapping.tp_size), + model_emb) + if use_weight_only: + v = np.ascontiguousarray(split_v.transpose()) + processed_torch_weights, torch_weight_scales = torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( + torch.tensor(v), plugin_weight_only_quant_type) + dst.value = processed_torch_weights.numpy() + scales = tensorrt_llm_qwen.layers[ + idx].attention.qkv.per_channel_scale + scales.value = torch_weight_scales.numpy() + else: + dst.value = np.ascontiguousarray(split_v) + elif 'attn.c_attn.bias' in k: + dst = tensorrt_llm_qwen.layers[idx].attention.qkv.bias + if multi_query_mode: + assert isinstance(v, list) and len(v) == 3 + wq = split(v[0], mapping.tp_size, mapping.tp_rank) + wk = split(v[1], mapping.tp_size, mapping.tp_rank) + wv = split(v[2], mapping.tp_size, mapping.tp_rank) + split_v = np.concatenate((wq, wk, wv)) + else: + q_emb = v.shape[0] // 3 + v = v.reshape(3, q_emb) + split_v = split(v, mapping.tp_size, mapping.tp_rank, dim=1) + split_v = split_v.reshape(3 * (q_emb // mapping.tp_size)) + dst.value = np.ascontiguousarray(split_v) + elif 'attn.c_proj.weight' in k: + dst = tensorrt_llm_qwen.layers[idx].attention.dense.weight + split_v = split(v, mapping.tp_size, mapping.tp_rank, dim=1) + if use_weight_only: + v = np.ascontiguousarray(split_v.transpose()) + processed_torch_weights, torch_weight_scales = torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( + torch.tensor(v), plugin_weight_only_quant_type) + dst.value = processed_torch_weights.numpy() + scales = tensorrt_llm_qwen.layers[ + idx].attention.dense.per_channel_scale + scales.value = torch_weight_scales.numpy() + else: + dst.value = np.ascontiguousarray(split_v) + elif 'mlp.w1.weight' in k: + dst = tensorrt_llm_qwen.layers[idx].mlp.gate.weight + split_v = split(v, mapping.tp_size, mapping.tp_rank, dim=0) + if use_weight_only: + v = np.ascontiguousarray(split_v.transpose()) + processed_torch_weights, torch_weight_scales = torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( + torch.tensor(v), plugin_weight_only_quant_type) + dst.value = processed_torch_weights.numpy() + scales = tensorrt_llm_qwen.layers[ + idx].mlp.gate.per_channel_scale + scales.value = torch_weight_scales.numpy() + else: + dst.value = np.ascontiguousarray(split_v) + elif 'mlp.w2.weight' in k: + dst = tensorrt_llm_qwen.layers[idx].mlp.fc.weight + split_v = split(v, mapping.tp_size, mapping.tp_rank, dim=0) + if use_weight_only: + v = np.ascontiguousarray(split_v.transpose()) + processed_torch_weights, torch_weight_scales = torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( + torch.tensor(v), plugin_weight_only_quant_type) + dst.value = processed_torch_weights.numpy() + scales = tensorrt_llm_qwen.layers[ + idx].mlp.fc.per_channel_scale + scales.value = torch_weight_scales.numpy() + else: + dst.value = np.ascontiguousarray(split_v) + elif 'mlp.c_proj.weight' in k: + dst = tensorrt_llm_qwen.layers[idx].mlp.proj.weight + split_v = split(v, mapping.tp_size, mapping.tp_rank, dim=1) + if use_weight_only: + v = np.ascontiguousarray(split_v.transpose()) + processed_torch_weights, torch_weight_scales = torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix( + torch.tensor(v), plugin_weight_only_quant_type) + dst.value = processed_torch_weights.numpy() + scales = tensorrt_llm_qwen.layers[ + idx].mlp.proj.per_channel_scale + scales.value = torch_weight_scales.numpy() + else: + dst.value = np.ascontiguousarray(split_v) + else: + print("unknown key: ", k) + + tok = time.time() + t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) + tensorrt_llm.logger.info(f'Weights loaded. Total time: {t}') + return diff --git a/examples/summarize.py b/examples/summarize.py new file mode 100644 index 000000000000..78f742727f28 --- /dev/null +++ b/examples/summarize.py @@ -0,0 +1,560 @@ +# 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 +from pathlib import Path + +import evaluate +import numpy as np +import torch +from datasets import load_dataset +from transformers import (AutoModel, AutoModelForCausalLM, + AutoModelForSeq2SeqLM, AutoTokenizer, T5Tokenizer) + +import tensorrt_llm +import tensorrt_llm.profiler as profiler +from tensorrt_llm.logger import logger +from tensorrt_llm.runtime import ModelRunner +from tensorrt_llm.tools.ppl import ppl + +DEFAULT_HF_MODEL_DIRS = { + 'baichuan': 'baichuan-inc/Baichuan-13B-Chat', + 'bloom': 'bigscience/bloom-560m', + 'chatglm_6b': 'THUDM/chatglm-6b', + 'chatglm2_6b': 'THUDM/chatglm2-6b', + 'chatglm2_6b_32k': 'THUDM/chatglm2-6b-32k', + 'chatglm3_6b': 'THUDM/chatglm3-6b', + 'chatglm3_6b_base': 'THUDM/chatglm3-6b-base', + 'chatglm3_6b_32k': 'THUDM/chatglm3-6b-32k', + 'falcon': 'tiiuae/falcon-rw-1b', + 'glm_10b': 'THUDM/glm-10b', + 'gpt': 'gpt2-medium', + 'gptj': 'EleutherAI/gpt-j-6b', + 'gptneox': 'EleutherAI/gpt-neox-20b', + 'internlm': 'internlm/internlm-chat-7b', + 'llama': 'meta-llama/Llama-2-7b-hf', + 'opt': 'facebook/opt-350m', +} + +DTYPE_STR_MAPPING = { + 'fp32': torch.float32, + 'fp16': torch.float16, + 'bf16': torch.bfloat16, + 'float32': torch.float32, + 'float16': torch.float16, + 'bfloat16': torch.bfloat16, +} + + +def read_model_name_from_config(config_path: Path): + with open(config_path, 'r') as f: + config = json.load(f) + return config['builder_config']['name'] + + +def main(args): + runtime_rank = tensorrt_llm.mpi_rank() + logger.set_level(args.log_level) + + model_name = read_model_name_from_config( + Path(args.engine_dir) / "config.json") + if args.hf_model_dir is None: + args.hf_model_dir = DEFAULT_HF_MODEL_DIRS[model_name] + if args.tokenizer_dir is None: + args.tokenizer_dir = args.hf_model_dir + + test_hf = args.test_hf and runtime_rank == 0 # only run hf on rank 0 + test_trt_llm = args.test_trt_llm + profiler.start('load tokenizer') + if args.vocab_file is None: + # Should set both padding_side and truncation_side to be 'left' + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_dir, + legacy=False, + padding_side='left', + truncation_side='left', + trust_remote_code=True) + else: + # From gpt-next + tokenizer = T5Tokenizer(vocab_file=args.vocab_file, + padding_side='left', + truncation_side='left') + profiler.stop('load tokenizer') + logger.info( + f'Load tokenizer takes: {profiler.elapsed_time_in_sec("load tokenizer")} sec' + ) + if not model_name.startswith('chatglm'): + tokenizer.pad_token = tokenizer.eos_token + if model_name == 'falcon' and tokenizer.pad_token_id is None: + tokenizer.pad_token_id = tokenizer.eos_token_id + + if args.eval_task == 'code_completion': + dataset_name = "openai_humaneval" + dataset_revision = None + dataset_input_key = 'prompt' + dataset_output_key = 'canonical_solution' + elif args.eval_task == 'summarize': + dataset_name = "ccdv/cnn_dailymail" + dataset_revision = "3.0.0" + dataset_input_key = 'article' + dataset_output_key = 'highlights' + dataset = load_dataset(dataset_name, + dataset_revision, + cache_dir=args.dataset_path) + + max_batch_size = args.batch_size + + # runtime parameters + # repetition_penalty = 1 + top_k = args.top_k + output_len = args.output_len + # TODO: The below lines are used to be compatible with the original code; may need fix + test_token_num = 800 if model_name.startswith('chatglm') else 923 + # top_p = 0.0 + # random_seed = 5 + temperature = 1 + num_beams = args.num_beams + length_penalty = args.length_penalty + + # TODO: The below lines are used to be compatible with the original code; may need fix + if model_name == 'falcon': + pad_id = tokenizer.pad_token_id + end_id = tokenizer.eos_token_id + else: + pad_id = tokenizer.encode(tokenizer.pad_token, + add_special_tokens=False)[0] + end_id = tokenizer.encode(tokenizer.eos_token, + add_special_tokens=False)[0] + + if test_trt_llm: + runner = ModelRunner.from_dir(args.engine_dir, + rank=runtime_rank, + debug_mode=args.debug_mode) + assert not (args.eval_ppl and not runner.session.gather_all_token_logits), \ + "PPL evaluation requires engine built with gather_all_token_logits enabled" + + if test_hf: + profiler.start('load HF model') + torch_dtype = DTYPE_STR_MAPPING[args.data_type] + if model_name.startswith('chatglm'): + auto_model_cls = AutoModel + elif model_name.startswith('glm'): + auto_model_cls = AutoModelForSeq2SeqLM + else: + auto_model_cls = AutoModelForCausalLM + model = auto_model_cls.from_pretrained( + args.hf_model_dir, + trust_remote_code=True, + torch_dtype=torch_dtype, + device_map='auto' if args.hf_device_map_auto else None) + if not args.hf_device_map_auto: + model.cuda() + profiler.stop('load HF model') + logger.info( + f'Load HF model takes: {profiler.elapsed_time_in_sec("load HF model")} sec' + ) + + 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') + + def _prepare_inputs(batch_input_texts, + eval_task='summarize', + add_special_tokens=True): + 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") + input_ids = tokenizer.encode(curr_text, + return_tensors='pt', + add_special_tokens=add_special_tokens, + truncation=True, + max_length=test_token_num) + # TODO: The below lines are used to be compatible with the original code; may need fix + if model_name.startswith(('chatglm2', 'chatglm3')): + input_ids = tokenizer.encode(curr_text, return_tensors='pt') + input_ids = input_ids[:, :test_token_num] + + 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): + 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) + input_lengths = [x.size(1) for x in batch_input_ids] + + with torch.no_grad(): + outputs = runner.generate( + batch_input_ids, + max_new_tokens=output_len, + max_kv_cache_length=args.max_kv_cache_length, + end_id=end_id, + pad_id=pad_id, + top_k=top_k, + num_beams=num_beams, + length_penalty=length_penalty, + output_sequence_lengths=True, + return_dict=True) + torch.cuda.synchronize() + + # Extract a list of tensors of shape beam_width x output_ids. + if runner.session.mapping.is_first_pp_rank(): + 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(batch_size) + ] + output_ids_list = [ + output_ids[batch_idx, :, input_lengths[batch_idx]:] + for batch_idx in range(batch_size) + ] + + ppls = [] + 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 = torch.stack( + outputs['generation_logits'][1:], dim=1) + for bidx in range(batch_size): + # [batch, beam, step] + curr_len = seq_lengths[bidx, 0] + curr_ctx_len = input_lengths[bidx] + curr_gen_len = curr_len - curr_ctx_len + + curr_ids = output_ids[bidx, 0, 1:curr_len] + curr_logits = torch.cat([ + context_logits[bidx], + generation_logits[bidx, :curr_gen_len - 1] + ], + dim=0) + curr_ppl = ppl(curr_logits, curr_ids) + ppls.append(curr_ppl) + logger.debug( + f"TensorRT-LLM PPL: {curr_ppl:.3f} | Generation length: {curr_gen_len}" + ) + + return output_beams_list, output_ids_list, ppls + return [], [], [] + + def eval_hf(datapoint, + eval_task='summarize', + eval_ppl=False, + add_special_tokens=True): + 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) + input_lengths = [x.size(1) 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.squeeze(0)]) + for x, pad in zip(batch_input_ids, paddings) + ] + batch_input_ids = torch.stack(batch_input_ids) + batch_input_ids = batch_input_ids.cuda() + + with torch.no_grad(): + outputs = model.generate(batch_input_ids, + max_new_tokens=output_len, + top_k=top_k, + temperature=temperature, + eos_token_id=tokenizer.eos_token_id, + pad_token_id=tokenizer.pad_token_id, + num_beams=num_beams, + num_return_sequences=num_beams, + early_stopping=True, + length_penalty=length_penalty, + output_scores=True, + return_dict_in_generate=True) + 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[:, len(batch_input_ids[0]):].tolist() + output_ids = output_ids.reshape([batch_size, num_beams, -1]) + output_lines_list = [ + tokenizer.batch_decode(output_ids[:, i, + len(batch_input_ids[0]):], + skip_special_tokens=True) + for i in range(num_beams) + ] + + ppls = [] + if eval_ppl and batch_size == 1: + # Only for batch size of 1 + seq_lens = [output_ids.size(-1) for _ in range(batch_size)] + context_logits = context_outputs['logits'] + # Remove the first generation logits which are same to last context logits + generation_logits = torch.stack(outputs['scores'][1:], dim=1) + + ppls = [] + for bidx in range(batch_size): + curr_len = seq_lens[bidx] + curr_ctx_len = input_lengths[bidx] + curr_gen_len = curr_len - curr_ctx_len + + curr_ids = output_ids[bidx, 0, 1:curr_len] + curr_logits = torch.cat([ + context_logits[bidx], + generation_logits[bidx, :curr_gen_len - 1] + ], + dim=0) + curr_ppl = ppl(curr_logits, curr_ids) + ppls.append(curr_ppl) + logger.debug( + f"HF PPL: {curr_ppl:.3f} | Generation length: {curr_gen_len}" + ) + + return output_lines_list, tokens_list, ppls + + if test_trt_llm: + datapoint = dataset['test'][0:1] + output, *_ = eval_trt_llm(datapoint, + eval_task=args.eval_task, + eval_ppl=args.eval_ppl, + add_special_tokens=args.add_special_tokens) + if runtime_rank == 0: + 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( + "---------------------------------------------------------") + + if test_hf: + datapoint = dataset['test'][0:1] + output, *_ = eval_hf(datapoint, + eval_task=args.eval_task, + eval_ppl=args.eval_ppl, + add_special_tokens=args.add_special_tokens) + 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("---------------------------------------------------------") + + # TODO: Add random_seed flag in gptj + metric_tensorrt_llm = [evaluate.load("rouge") for _ in range(num_beams)] + metric_hf = [evaluate.load("rouge") for _ in range(num_beams)] + for i in range(num_beams): + metric_tensorrt_llm[i].seed = 0 + metric_hf[i].seed = 0 + ppls_trt_llm, ppls_hf = [], [] + + ite_count = 0 + data_point_idx = 0 + while (data_point_idx < len(dataset['test'])) 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['test'][data_point_idx:(data_point_idx + + max_batch_size)] + + if test_trt_llm: + profiler.start('tensorrt_llm') + output_tensorrt_llm, _, curr_ppls_trt_llm = eval_trt_llm( + datapoint, + eval_task=args.eval_task, + eval_ppl=args.eval_ppl, + add_special_tokens=args.add_special_tokens) + profiler.stop('tensorrt_llm') + + if test_hf: + profiler.start('hf') + output_hf, _, curr_ppls_hf = eval_hf( + datapoint, + eval_task=args.eval_task, + eval_ppl=args.eval_ppl, + add_special_tokens=args.add_special_tokens) + profiler.stop('hf') + + if runtime_rank == 0: + if test_trt_llm: + for batch_idx in range(len(output_tensorrt_llm)): + for beam_idx in range(num_beams): + metric_tensorrt_llm[beam_idx].add_batch( + predictions=[ + output_tensorrt_llm[batch_idx][beam_idx] + ], + references=[ + datapoint[dataset_output_key][batch_idx] + ]) + if output_dir is not None: + # yapf: disable + for i in range(len(output_tensorrt_llm[0])): + for beam_idx in range(num_beams): + 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') + # yapf: enable + ppls_trt_llm.extend(curr_ppls_trt_llm) + if test_hf: + for beam_idx in range(num_beams): + 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 output_dir is not None: + # yapf: disable + for i in range(len(output_hf[0])): + for beam_idx in range(num_beams): + 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') + # yapf: enable + ppls_hf.extend(curr_ppls_hf) + + logger.debug('-' * 100) + logger.debug(f"Input : {datapoint[dataset_input_key]}") + if test_trt_llm: + logger.debug(f'TensorRT-LLM Output: {output_tensorrt_llm}') + if test_hf: + logger.debug(f'HF Output: {output_hf}') + logger.debug(f"Reference : {datapoint[dataset_output_key]}") + + data_point_idx += max_batch_size + ite_count += 1 + + if runtime_rank == 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)' + ) + for beam_idx in range(num_beams): + logger.info(f"TensorRT-LLM beam {beam_idx} result") + 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: + assert computed_metrics_tensorrt_llm[ + 'rouge1'] * 100 > args.tensorrt_llm_rouge1_threshold + if args.eval_ppl: + logger.info(f" Per-token perplexity: {np.mean(ppls_trt_llm)}") + 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)' + ) + for beam_idx in range(num_beams): + logger.info(f"HF beam {beam_idx} result") + computed_metrics_hf = metric_hf[beam_idx].compute() + 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)}") + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--hf_model_dir', type=str, default=None) + parser.add_argument( + '--tokenizer_dir', + default=None, + help='tokenizer path; defaults to hf_model_dir if left unspecified') + parser.add_argument('--vocab_file') + parser.add_argument('--test_hf', action='store_true') + parser.add_argument('--test_trt_llm', action='store_true') + parser.add_argument( + '--data_type', + type=str, + choices=['fp32', 'fp16', 'bf16', 'float32', 'float16', 'bfloat16'], + default='fp16') + parser.add_argument('--dataset_path', type=str, default='') + parser.add_argument('--log_level', type=str, default='info') + parser.add_argument('--engine_dir', type=str, default='engine_outputs') + 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_kv_cache_length', + type=int, + default=None, + help='The max kv cache length. \ + If the final sequence length exceeds the kv cache length, we will enable cyclic kv cache. \ + If it is set to None, we will use the max sequence length.') + parser.add_argument('--check_accuracy', action='store_true') + parser.add_argument('--tensorrt_llm_rouge1_threshold', + type=float, + default=15.0) + parser.add_argument('--num_beams', type=int, default=1) + parser.add_argument('--top_k', type=int, default=1) + parser.add_argument('--eval_task', + type=str, + default='summarize', + choices=['summarize', 'code_completion']) + parser.add_argument('--length_penalty', type=float, default=1.0) + parser.add_argument('--eval_ppl', action='store_true') + parser.add_argument('--debug_mode', + default=False, + action='store_true', + help="Whether or not to turn on the debug mode") + 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_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( + '--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.") + + args = parser.parse_args() + + main(args) diff --git a/requirements-dev.txt b/requirements-dev.txt index 02536864b26b..4edf2b4ef7af 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,5 +1,6 @@ torch transformers==4.33.1 +typing-extensions==4.8.0 diffusers==0.15.0 accelerate==0.20.3 colored @@ -10,6 +11,7 @@ mpi4py numpy cuda-python==12.2.0 mypy +pybind11-stubgen pytest-cov pytest-xdist pytest-forked diff --git a/scripts/build_wheel.py b/scripts/build_wheel.py index a93098d424c4..13025e66c718 100755 --- a/scripts/build_wheel.py +++ b/scripts/build_wheel.py @@ -52,7 +52,8 @@ def main(build_type: str = "Release", cpp_only: bool = False, install: bool = False, skip_building_wheel: bool = False, - python_bindings: bool = False): + python_bindings: bool = False, + benchmarks: bool = False): project_dir = Path(__file__).parent.resolve().parent os.chdir(project_dir) build_run = partial(run, shell=True, check=True) @@ -145,6 +146,7 @@ def main(build_type: str = "Release", th_common_lib = "" if cpp_only else "th_common" build_pybind = "ON" if python_bindings else "OFF" bindings_lib = "bindings" if python_bindings else "" + benchmarks_lib = "benchmarks" if benchmarks else "" with working_directory(build_dir): cmake_def_args = " ".join(cmake_def_args) @@ -155,7 +157,7 @@ def main(build_type: str = "Release", ) build_run( f'cmake --build . --config {build_type} --parallel {job_count} ' - f'--target tensorrt_llm tensorrt_llm_static nvinfer_plugin_tensorrt_llm {th_common_lib} {bindings_lib}' + f'--target tensorrt_llm tensorrt_llm_static nvinfer_plugin_tensorrt_llm {th_common_lib} {bindings_lib} {benchmarks_lib}' f'{" ".join(extra_make_targets)}') if cpp_only: @@ -192,6 +194,10 @@ def main(build_type: str = "Release", pybind_lib ) == 1, f"Exactly one pybind library should be present: {pybind_lib}" copy(pybind_lib[0], pkg_dir) + build_run(f"{sys.executable} -m pip install pybind11-stubgen") + build_run( + f"cd {pkg_dir} && {sys.executable} -m pybind11_stubgen -o . bindings" + ) if dist_dir is None: dist_dir = project_dir / "build" @@ -265,5 +271,8 @@ def main(build_type: str = "Release", "-p", action="store_true", help="Build the python bindings for the C++ runtime.") + parser.add_argument("--benchmarks", + action="store_true", + help="Build the benchmarks for the C++ runtime.") args = parser.parse_args() main(**vars(args)) diff --git a/setup.py b/setup.py index 1644c422fac2..63171df7eaf2 100644 --- a/setup.py +++ b/setup.py @@ -64,7 +64,7 @@ def is_pure(self): (['libs/th_common.dll', 'libs/nvinfer_plugin_tensorrt_llm.dll'] if platform.system() == "Windows" else [ 'libs/libth_common.so', 'libs/libnvinfer_plugin_tensorrt_llm.so', - 'bindings.*.so' + 'bindings.*.so', 'bindings.pyi' ]) + ['tools/plugin_gen/templates/*'], }, python_requires=">=3.7, <4", diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index 3540466e7b29..30ae90a5fe31 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -17,6 +17,7 @@ import math import struct from functools import partial +from pathlib import Path, PosixPath import numpy as np import tensorrt as trt @@ -30,8 +31,8 @@ def torch_to_numpy(x: torch.Tensor): assert isinstance(x, torch.Tensor), \ f'x must be a torch.Tensor object, but got {type(x)}.' if x.dtype != torch.bfloat16: - return x.cpu().numpy() - return x.view(torch.int16).cpu().numpy().view(np_bfloat16) + return x.detach().cpu().numpy() + return x.view(torch.int16).detach().cpu().numpy().view(np_bfloat16) def numpy_to_torch(x): @@ -40,6 +41,21 @@ def numpy_to_torch(x): return torch.tensor(x.view(np.int16)).view(torch.bfloat16) +def numpy_to_dtype(x, dtype: str): + if x.dtype == np_bfloat16: + # BF16 --> non-BF16 or BF16 + if dtype != 'bfloat16': + torch_to_numpy(numpy_to_torch(x).to(str_dtype_to_torch(dtype))) + else: + return x + else: + # non-BF16 types --> non-BF16 or BF16 + if dtype != 'bfloat16': + return x.astype(str_dtype_to_np(dtype)) + else: + return torch_to_numpy(torch.from_numpy(x).to(torch.bfloat16)) + + fp32_array = partial(np.array, dtype=np.float32) fp16_array = partial(np.array, dtype=np.float16) int32_array = partial(np.array, dtype=np.int32) @@ -115,14 +131,12 @@ def str_dtype_to_trt(dtype): np.dtype('int32'): trt.int32, np.dtype('float16'): trt.float16, np.dtype('float32'): trt.float32, + np_bfloat16: trt.bfloat16, + np.bool_: trt.bool, } def np_dtype_to_trt(dtype): - if trt_version() >= '7.0' and dtype == np.bool_: - return trt.bool - if trt_version() >= '9.0' and dtype == np_bfloat16: - return trt.bfloat16 ret = _np_to_trt_dtype_dict.get(dtype) assert ret is not None, f'Unsupported dtype: {dtype}' return ret @@ -134,12 +148,11 @@ def np_dtype_to_trt(dtype): trt.float16: np.float16, trt.float32: np.float32, trt.bool: np.bool_, + trt.bfloat16: np_bfloat16, } def trt_dtype_to_np(dtype): - if trt_version() >= '9.0' and dtype == trt.bfloat16: - return np_bfloat16 ret = _trt_to_np_dtype_dict.get(dtype) assert ret is not None, f'Unsupported dtype: {dtype}' return ret @@ -162,12 +175,11 @@ def torch_dtype_to_np(dtype): trt.float32: torch.float32, trt.int32: torch.int32, trt.int8: torch.int8, + trt.bfloat16: torch.bfloat16 } def trt_dtype_to_torch(dtype): - if trt_version() >= '9.0' and dtype == trt.bfloat16: - return torch.bfloat16 ret = _trt_to_torch_dtype_dict.get(dtype) assert ret is not None, f'Unsupported dtype: {dtype}' return ret @@ -242,3 +254,18 @@ def numpy_fp32_to_bf16(src): bytes = struct.pack(' Tensor: ''' Add an activation function. @@ -2939,7 +2946,8 @@ def bert_attention(tensor: Tensor, q_scaling: float, relative_attention: bool = False, relative_attention_bias: Tensor = None, - max_distance: int = 0) -> Tuple[Tensor]: + max_distance: int = 0, + max_input_length: Tensor = None) -> Tuple[Tensor]: ''' Add an operation that performs the multi-head attention in BERT. @@ -2990,6 +2998,9 @@ def bert_attention(tensor: Tensor, Implicit mode is only enabled when passing in non-zero positive max_distance value. See relative attention bias in docs/gpt_attention.md + max_input_length: Tensor = None + The maximum input sequence length represented by Tensor shape. Requires for remove_input_padding to pre-define plugin workspace size. + Returns: The tensor produced by that layer. ''' @@ -3025,14 +3036,22 @@ def bert_attention(tensor: Tensor, max_distance = trt.PluginField("max_distance", np.array(max_distance, dtype=np.int32), trt.PluginFieldType.INT32) + remove_padding = trt.PluginField( + "remove_padding", + np.array(np.int8(default_net().plugin_config.remove_input_padding), + dtype=np.int8), trt.PluginFieldType.INT8) pfc = trt.PluginFieldCollection([ nheads, head_size, q_scaling, enable_qk_half_accum, context_fmha_type, - pf_type, do_relative_attention, max_distance + pf_type, do_relative_attention, max_distance, remove_padding ]) attn_plug = attn_plg_creator.create_plugin("padding_attn", pfc) plug_inputs = [tensor, input_lengths] + if max_input_length is not None: + # for remove padding mode + plug_inputs += [max_input_length] if relative_attention_bias is not None: + # for relative attention mode plug_inputs += [relative_attention_bias] plug_inputs = [i.trt_tensor for i in plug_inputs] @@ -3198,6 +3217,7 @@ def gpt_attention( * tensorrt_llm.layers.AttentionMaskType.padding for BERT, * tensorrt_llm.layers.AttentionMaskType.causal for GPT, * tensorrt_llm.layers.AttentionMaskType.bidirectional for ChatGLM-6B, + * tensorrt_llm.layers.AttentionMaskType.bidirectionalglm for GLM-10B, alibi_slopes: Tensor The ALiBi slopes. The ALiBi bias is computed on-the-fly in the kernel @@ -3874,3 +3894,109 @@ def non_gated_version(activation): if is_gated_activation(activation): return GATED_ACT_2_ACT[activation] return activation + + +def lora_plugin( + input: Tensor = None, + in_hidden_size: int = 0, + out_hidden_size: int = 0, + host_request_types: Tensor = None, + transa: bool = False, + transb: bool = False, + host_context_lengths: Tensor = None, # for pad-free input mode + max_context_length: int = 0, + max_low_rank: int = 0, + lora_ranks: Tensor = None, + lora_weights_pointers: Tensor = None, +): + ''' + Parameters: + lora_ids : cpu Tensor = None + A tensor that contains the lora ids of different inputs. + + in_hidden_size/out_hidden_size : int + the lora computation workflow is + [M, in_hidden_size] -> [M, low_rank] -> [M, out_hidden_size] + + host_request_types : Tensor = None + The tensor on the host that indicates if a request is in context or + generation phase. Its shape is [batch_size]. See Inflight Batching + in docs/gpt_attention.md, + + transa : bool + Is the first input transposed? Set to 'True' if you want the first + input to be transposed, 'False' otherwise. + + transb : bool + Is the second input transposed? Set to 'True' if you want the + second input to be transposed, 'False' otherwise. + + host_context_lengths: cpu Tensor = None + A host tensor that contains the lengths of the different inputs, + + max_context_length : int + Maximum length during context phase, used to determine the workspace size. + + max_low_rank : int + Maximum low_rank, used to determine the workspace size. + + lora_ranks : cpu Tensor with shape [batch_size] + The low_rank of each request + + lora_weights_pointers : cpu int64 Tensor with shape [batch_size, 2] + The weights pointers of each request. Consist of in_pointer and out_pointer. + + Return: + The tensor produced by that layer. + + ''' + assert host_context_lengths is not None or not default_net( + ).plugin_config.remove_input_padding + + trt.get_plugin_registry().plugin_creator_list + in_hidden_size = trt.PluginField("in_hidden_size", + np.array(in_hidden_size, dtype=np.int32), + trt.PluginFieldType.INT32) + out_hidden_size = trt.PluginField("out_hidden_size", + np.array(out_hidden_size, dtype=np.int32), + trt.PluginFieldType.INT32) + transa = 1 if transa else 0 + transa = trt.PluginField("transa", np.array(transa, dtype=np.int32), + trt.PluginFieldType.INT32) + transb = 1 if transb else 0 + transb = trt.PluginField("transb", np.array(transb, dtype=np.int32), + trt.PluginFieldType.INT32) + + plg_creator = trt.get_plugin_registry().get_plugin_creator( + 'Lora', '1', TRT_LLM_PLUGIN_NAMESPACE) + assert plg_creator is not None + + p_dtype = default_net().plugin_config.lora_plugin + pf_type = trt.PluginField( + "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), + trt.PluginFieldType.INT32) + remove_input_padding = trt.PluginField( + "remove_input_padding", + np.array(np.int8(default_net().plugin_config.remove_input_padding), + dtype=np.int8), trt.PluginFieldType.INT8) + max_context_length_filed = trt.PluginField( + "max_context_length", np.array(max_context_length, dtype=np.int32), + trt.PluginFieldType.INT32) + max_low_rank_filed = trt.PluginField("max_low_rank", + np.array(max_low_rank, dtype=np.int32), + trt.PluginFieldType.INT32) + + pfc = trt.PluginFieldCollection([ + in_hidden_size, out_hidden_size, transa, transb, pf_type, + remove_input_padding, max_context_length_filed, max_low_rank_filed + ]) + lora_plug = plg_creator.create_plugin("lora", pfc) + + plug_inputs = [input, host_request_types, lora_ranks, lora_weights_pointers] + if default_net().plugin_config.remove_input_padding: + plug_inputs += [host_context_lengths] + + plug_inputs = [i.trt_tensor for i in plug_inputs] + layer = default_trtnet().add_plugin_v2(plug_inputs, lora_plug) + + return _create_tensor(layer.get_output(0), layer) diff --git a/tensorrt_llm/graph_rewriting.py b/tensorrt_llm/graph_rewriting.py index dc12168c5783..fabe13acae13 100644 --- a/tensorrt_llm/graph_rewriting.py +++ b/tensorrt_llm/graph_rewriting.py @@ -544,9 +544,10 @@ def __enter__(self): FLayerInfoMemo.cur_flayer = self.layer def __exit__(self, exc_type, exc_val, exc_tb): - assert self.layer.layer_name != "", f"FLayer {self.layer.layer_kind} without a plugin name detected" - FLayerInfoMemo.instance().add(self.layer.layer_name, self.layer) FLayerInfoMemo.cur_flayer = None + if exc_type is None: + assert self.layer.layer_name != "", f"FLayer {self.layer.layer_kind} without a plugin name detected" + FLayerInfoMemo.instance().add(self.layer.layer_name, self.layer) def record_signature(f): diff --git a/tensorrt_llm/layers/__init__.py b/tensorrt_llm/layers/__init__.py index 911020981c12..2df45dd12fd9 100644 --- a/tensorrt_llm/layers/__init__.py +++ b/tensorrt_llm/layers/__init__.py @@ -20,6 +20,7 @@ from .conv import Conv2d, ConvTranspose2d from .embedding import Embedding, PromptTuningEmbedding from .linear import ColumnLinear, Linear, RowLinear +from .lora import Lora, LoraParams from .mlp import MLP, FusedGatedMLP, GatedMLP from .normalization import GroupNorm, LayerNorm, RmsNorm from .pooling import AvgPool2d @@ -47,4 +48,6 @@ 'Cast', 'AttentionParams', 'KeyValueCacheParams', + 'Lora', + 'LoraParams', ] diff --git a/tensorrt_llm/layers/attention.py b/tensorrt_llm/layers/attention.py index 7745a8b24a77..8ec1ad3cab18 100644 --- a/tensorrt_llm/layers/attention.py +++ b/tensorrt_llm/layers/attention.py @@ -31,6 +31,7 @@ from ..quantization import QuantMode from ..quantization.layers import FP8Linear, FP8RowLinear from .linear import ColumnLinear, RowLinear +from .lora import Lora class RopeEmbeddingUtils: @@ -150,6 +151,7 @@ def apply_rotary_pos_emb_chatglm( num_attention_heads, attention_head_size, max_position_embeddings, + rotary_embedding_scale, ) -> Tensor: half_head_size = attention_head_size // 2 @@ -175,6 +177,7 @@ def apply_rotary_pos_emb_chatglm( embedding_weight = RopeEmbeddingUtils.create_sinusoidal_positions( max_position_embeddings, half_head_size) + embedding_weight /= rotary_embedding_scale embedding_weight = np.split(embedding_weight.squeeze(0), 2, axis=1) embedding_weight = np.concatenate( [ @@ -324,6 +327,7 @@ def __init__( max_position_embeddings=1024, num_layers=1, apply_query_key_layer_scaling=False, + attention_head_size=None, attention_mask_type=AttentionMaskType.padding, bias=True, dtype=None, @@ -348,7 +352,7 @@ def __init__( self.cross_attention = cross_attention self.attention_mask_type = attention_mask_type - self.attention_head_size = hidden_size // num_attention_heads + self.attention_head_size = hidden_size // num_attention_heads if attention_head_size is None else attention_head_size assert num_attention_heads % tp_size == 0, \ "num_attention_heads must be divisible by tp_size" self.num_attention_heads = num_attention_heads // tp_size @@ -421,15 +425,16 @@ def __init__( self.use_fp8_qdq = self.quant_mode.has_fp8_qdq() if self.use_fp8_qdq: - self.qkv = FP8Linear(hidden_size, - hidden_size + - (2 * tp_size * self.num_attention_kv_heads * - self.attention_head_size), - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False) + self.qkv = FP8Linear( + hidden_size, + tp_size * self.num_attention_heads * self.attention_head_size + + (2 * tp_size * self.num_attention_kv_heads * + self.attention_head_size), + bias=bias, + dtype=dtype, + tp_group=tp_group, + tp_size=tp_size, + gather_output=False) self.dense = FP8RowLinear(hidden_size, hidden_size, bias=dense_bias, @@ -438,16 +443,20 @@ def __init__( tp_size=tp_size, instance_id=instance_id) else: - self.qkv = ColumnLinear(hidden_size, - hidden_size + - (2 * tp_size * self.num_attention_kv_heads * - self.attention_head_size), - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False) - self.dense = RowLinear(hidden_size, + # out dim is not necessarily hidden_size + kv specific size (in MQA/GQA), but num_heads * heads_size + # example: d_model != num_heads * head_size in Flan-T5 + self.qkv = ColumnLinear( + hidden_size, + tp_size * self.num_attention_heads * self.attention_head_size + + (2 * tp_size * self.num_attention_kv_heads * + self.attention_head_size), + bias=bias, + dtype=dtype, + tp_group=tp_group, + tp_size=tp_size, + gather_output=False) + self.dense = RowLinear(tp_size * self.num_attention_heads * + self.attention_head_size, hidden_size, bias=dense_bias, dtype=dtype, @@ -461,18 +470,25 @@ def __init__( tp_size, num_buckets), dtype=dtype) - def forward( - self, - hidden_states: Tensor, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - encoder_output: Optional[Tensor] = None, - workspace=None, - position_embedding=None, - norm_before_bmm1=False, - ): + self.qkv_lora = Lora( + in_hidden_size=hidden_size, + out_hidden_size=hidden_size + + (2 * tp_size * self.num_attention_kv_heads * + self.attention_head_size), + max_low_rank=hidden_size, + ) + + def forward(self, + hidden_states: Tensor, + attention_mask=None, + use_cache=False, + kv_cache_params=None, + attention_params=None, + encoder_output: Optional[Tensor] = None, + workspace=None, + position_embedding=None, + norm_before_bmm1=False, + lora_params=None): assert isinstance(hidden_states, Tensor) @@ -491,6 +507,15 @@ def forward( qkv = self.qkv(hidden_states) + if default_net().plugin_config.lora_plugin: + qkv = qkv + self.qkv_lora( + hidden_states, + host_request_types=attention_params.host_request_types, + host_context_lengths=attention_params.host_context_lengths, + max_context_length=attention_params.max_context_length, + lora_ranks=lora_params.lora_ranks, + lora_weights_pointers=lora_params.lora_weights_pointers_list[0]) + if self.position_embedding_type == PositionEmbeddingType.chatglm: qkv = RopeEmbeddingUtils.apply_rotary_pos_emb_chatglm( qkv, @@ -498,7 +523,10 @@ def forward( self.num_attention_heads, self.attention_head_size, self.max_position_embeddings, + self.rotary_embedding_scale, ) + self.rotary_embedding_scale_type = RotaryScalingType.none + self.rotary_embedding_scale = 1.0 paged_kv_cache = default_net().plugin_config.paged_kv_cache @@ -529,7 +557,8 @@ def forward( if default_net().plugin_config.gpt_attention_plugin: assert self.attention_mask_type in [ - AttentionMaskType.causal, AttentionMaskType.bidirectional + AttentionMaskType.causal, AttentionMaskType.bidirectional, + AttentionMaskType.bidirectionalglm ], 'Plugin only support masked MHA.' kv_orig_quant_scale = self.kv_orig_quant_scale.value if self.quant_mode.has_kv_cache_quant( ) else None @@ -847,9 +876,10 @@ class BertAttention(Module): def __init__(self, hidden_size, num_attention_heads, - num_kv_heads=None, max_position_embeddings=1024, num_layers=1, + attention_head_size=None, + num_kv_heads=None, q_scaling=1.0, apply_query_key_layer_scaling=False, bias=True, @@ -862,7 +892,7 @@ def __init__(self, num_buckets=0): super().__init__() - self.attention_head_size = hidden_size // num_attention_heads + self.attention_head_size = hidden_size // num_attention_heads if attention_head_size is None else attention_head_size self.num_attention_heads = num_attention_heads // tp_size self.num_attention_kv_heads = ( num_kv_heads + tp_size - 1 @@ -886,16 +916,20 @@ def __init__(self, self.relative_attention = relative_attention self.max_distance = max_distance - self.qkv = ColumnLinear(hidden_size, - hidden_size + - (2 * tp_size * self.num_attention_kv_heads * - self.attention_head_size), - bias=bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=False) - self.dense = RowLinear(hidden_size, + # out dim is not necessarily hidden_size + kv specific size (in MQA/GQA), but num_heads * heads_size + # example: d_model != num_heads * head_size in Flan-T5 + self.qkv = ColumnLinear( + hidden_size, + tp_size * self.num_attention_heads * self.attention_head_size + + (2 * tp_size * self.num_attention_kv_heads * + self.attention_head_size), + bias=bias, + dtype=dtype, + tp_group=tp_group, + tp_size=tp_size, + gather_output=False) + self.dense = RowLinear(tp_size * self.num_attention_heads * + self.attention_head_size, hidden_size, bias=bias, dtype=dtype, @@ -911,7 +945,9 @@ def __init__(self, def forward(self, hidden_states: Tensor, attention_mask=None, - input_lengths=None): + input_lengths=None, + workspace=None, + max_input_length=None): assert isinstance(hidden_states, Tensor) qkv = self.qkv(hidden_states) @@ -928,7 +964,8 @@ def forward(self, relative_attention=self.relative_attention, max_distance=self.max_distance, relative_attention_bias=self.rel_attn_table.value - if self.relative_attention else None) + if self.relative_attention else None, + max_input_length=max_input_length) else: # plain TRT mode def transpose_for_scores(x): @@ -958,6 +995,6 @@ def transpose_for_scores(x): concat([shape(context, 0), shape(context, 1), self.hidden_size])) - context = self.dense(context) + context = self.dense(context, workspace) return context diff --git a/tensorrt_llm/layers/lora.py b/tensorrt_llm/layers/lora.py new file mode 100644 index 000000000000..d6deb83a721a --- /dev/null +++ b/tensorrt_llm/layers/lora.py @@ -0,0 +1,66 @@ +# 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. + +from typing import List + +from .._common import default_net +from ..functional import Tensor, lora_plugin +from ..module import Module + + +class Lora(Module): + + def __init__(self, + in_hidden_size: int = 0, + out_hidden_size: int = 0, + max_low_rank: int = 0) -> None: + super().__init__() + + self.in_hidden_size = in_hidden_size + self.out_hidden_size = out_hidden_size + self.max_low_rank = max_low_rank + + def forward(self, + x, + host_request_types=None, + host_context_lengths=None, + max_context_length: int = 0, + lora_ranks=None, + lora_weights_pointers=None): + if default_net().plugin_config.lora_plugin: + x = lora_plugin(x, + in_hidden_size=self.in_hidden_size, + out_hidden_size=self.out_hidden_size, + host_request_types=host_request_types, + transb=True, + host_context_lengths=host_context_lengths, + max_context_length=max_context_length, + max_low_rank=self.max_low_rank, + lora_ranks=lora_ranks, + lora_weights_pointers=lora_weights_pointers) + else: + assert False, "Not support lora without plugin" + + return x + + +class LoraParams(object): + + def __init__(self, + lora_ranks: Tensor = None, + lora_weights_pointers_list: List[Tensor] = None): + + self.lora_ranks = lora_ranks + self.lora_weights_pointers_list = lora_weights_pointers_list diff --git a/tensorrt_llm/mapping.py b/tensorrt_llm/mapping.py index 3c47fb314649..0454f0f4dcc8 100644 --- a/tensorrt_llm/mapping.py +++ b/tensorrt_llm/mapping.py @@ -64,6 +64,9 @@ def __init__(self, self.tp_group = self.tp_groups[self.pp_rank] self.pp_group = self.pp_groups[self.tp_rank] + def has_tp(self): + return self.tp_size > 1 + def is_last_pp_rank(self): return self.pp_rank == self.pp_size - 1 diff --git a/tensorrt_llm/models/__init__.py b/tensorrt_llm/models/__init__.py index 193e7a2c927b..4d309fddc875 100755 --- a/tensorrt_llm/models/__init__.py +++ b/tensorrt_llm/models/__init__.py @@ -16,14 +16,16 @@ from .bert.model import BertForQuestionAnswering, BertModel from .bloom.model import BloomForCausalLM, BloomModel from .chatglm.model import ChatGLMHeadModel, ChatGLMModel +from .enc_dec.model import DecoderModel, EncoderModel from .falcon.model import FalconForCausalLM, FalconModel from .gpt.model import GPTLMHeadModel, GPTModel from .gptj.model import GPTJForCausalLM, GPTJModel from .gptneox.model import GPTNeoXForCausalLM, GPTNeoXModel -from .internlm.model import InternLMForCausalLM, InternLMModel from .llama.model import LLaMAForCausalLM, LLaMAModel from .opt.model import OPTLMHeadModel, OPTModel -from .quantized.quant import quantize_model # noqa +from .qwen.model import QWenForCausalLM + +from .quantized.quant import quantize_model # noqa # isort:skip __all__ = [ 'BertModel', @@ -46,6 +48,7 @@ 'ChatGLMHeadModel', 'ChatGLMModel', 'BaichuanForCausalLM', - 'InternLMForCausalLM', - 'InternLMModel', + 'QWenForCausalLM', + 'EncoderModel', + 'DecoderModel', ] diff --git a/tensorrt_llm/models/bloom/model.py b/tensorrt_llm/models/bloom/model.py index c12eed21d1e0..accde2137409 100644 --- a/tensorrt_llm/models/bloom/model.py +++ b/tensorrt_llm/models/bloom/model.py @@ -253,6 +253,7 @@ def __init__(self, elif quant_mode.has_fp8_kv_cache(): self._kv_dtype = str_dtype_to_trt('fp8') + self.mapping = mapping self.quant_mode = quant_mode self._num_layers = num_layers diff --git a/tensorrt_llm/models/chatglm/model.py b/tensorrt_llm/models/chatglm/model.py index 9eea959fc44c..226e1f757e3e 100644 --- a/tensorrt_llm/models/chatglm/model.py +++ b/tensorrt_llm/models/chatglm/model.py @@ -18,8 +18,8 @@ from ..._common import default_net from ..._utils import pad_vocab_size, str_dtype_to_trt -from ...functional import (PositionEmbeddingType, Tensor, - gather_last_token_logits) +from ...functional import (PositionEmbeddingType, Tensor, concat, + gather_last_token_logits, shape) from ...layers import (MLP, Attention, AttentionMaskType, AttentionParams, ColumnLinear, Embedding, KeyValueCacheParams, LayerNorm, RmsNorm) @@ -33,15 +33,33 @@ def __init__(self, layer_id, args): super().__init__() - self.model_version = args.model_version + self.model_name = args.model_name self.use_cache = args.use_cache + rotary_embedding_scaling = None - if self.model_version == "1": + if self.model_name in ["chatglm_6b"]: self.alpha = (2 * args.num_layers)**0.5 self.norm = LayerNorm - else: + attention_mask_type = AttentionMaskType.bidirectional + position_embedding_type = PositionEmbeddingType.chatglm + elif args.model_name in [ + "chatglm2_6b", "chatglm2_6b_32k", "chatglm3_6b", + "chatglm3_6b_base", "chatglm3_6b_32k" + ]: self.apply_residual_connection_post_layernorm = args.apply_residual_connection_post_layernorm self.norm = RmsNorm if args.rmsnorm else LayerNorm + attention_mask_type = AttentionMaskType.causal + position_embedding_type = PositionEmbeddingType.rope_gptj + if args.model_name in ["chatglm2_6b_32k", "chatglm3_6b_32k"]: + rotary_embedding_scaling = { + "type": "linear", + "factor": args.rotary_embedding_scaling + } + elif args.model_name in ["glm_10b"]: + self.apply_residual_connection_post_layernorm = args.apply_residual_connection_post_layernorm + self.norm = LayerNorm + attention_mask_type = AttentionMaskType.bidirectionalglm + position_embedding_type = PositionEmbeddingType.learned_absolute self.pre_norm = self.norm( normalized_shape=args.hidden_size, @@ -57,14 +75,12 @@ def __init__(self, layer_id, args): max_position_embeddings=args.max_seq_length, num_layers=args.num_layers, apply_query_key_layer_scaling=args.apply_query_key_layer_scaling, - attention_mask_type=AttentionMaskType.bidirectional - if args.model_version == "1" else AttentionMaskType.causal, + attention_mask_type=attention_mask_type, bias=args.qkv_bias, dtype=args.dtype, - position_embedding_type=PositionEmbeddingType.chatglm - if args.model_version == "1" else PositionEmbeddingType.rope_gptj, + position_embedding_type=position_embedding_type, rotary_embedding_base=10000.0, - rotary_embedding_scaling=None, + rotary_embedding_scaling=rotary_embedding_scaling, use_int8_kv_cache=args.quant_mode.has_int8_kv_cache(), rotary_embedding_percentage=0.5, tp_group=args.mapping.tp_group, @@ -123,7 +139,7 @@ def forward( if self.use_cache: attention_output, presents = attention_output - if self.model_version == "1": + if self.model_name in ["chatglm_6b"]: residual = norm_output norm_input = residual * self.alpha + attention_output @@ -136,7 +152,10 @@ def forward( output = residual * self.alpha + mlp_output - else: + elif self.model_name in [ + "chatglm2_6b", "chatglm2_6b_32k", "chatglm3_6b", + "chatglm3_6b_base", "chatglm3_6b_32k", "glm_10b" + ]: residual = norm_output if self.apply_residual_connection_post_layernorm else hidden_states norm_input = residual + attention_output @@ -158,7 +177,15 @@ def __init__(self, args): super().__init__() - self.norm = LayerNorm if args.model_version == "1" else RmsNorm + self.model_name = args.model_name + + if args.model_name in ["chatglm_6b", "glm_10b"]: + self.norm = LayerNorm + elif args.model_name in [ + "chatglm2_6b", "chatglm2_6b_32k", "chatglm3_6b", + "chatglm3_6b_base", "chatglm3_6b_32k" + ]: + self.norm = RmsNorm self.use_cache = args.use_cache self.embedding = Embedding( @@ -172,6 +199,28 @@ def __init__(self, args): instance_id=args.num_layers * 2, ) + if args.model_name in ["glm_10b"]: + self.position_embeddings = Embedding( + args.max_seq_length + 1, + args.hidden_size, + dtype=args.dtype, + tp_size=1, #args.mapping.tp_size, + tp_group=None, #args.mapping.tp_group, + sharding_dim=0, + tp_rank=0, #args.mapping.rank, + instance_id=args.num_layers * 2, + ) + self.block_embeddings = Embedding( + args.max_seq_length + 1, + args.hidden_size, + dtype=args.dtype, + tp_size=1, #args.mapping.tp_size, + tp_group=None, #args.mapping.tp_group, + sharding_dim=0, + tp_rank=0, #args.mapping.rank, + instance_id=args.num_layers * 2, + ) + self.layers = ModuleList( ChatGLMDecoderLayer(i, args) for i in range(args.num_layers)) @@ -192,6 +241,21 @@ def forward( hidden_states = self.embedding(input_ids) + if self.model_name in ["glm_10b"]: + position_ids_list = position_ids.split(1, dim=1) + position_embedding = self.position_embeddings(position_ids_list[0]) + block_embedding = self.block_embeddings(position_ids_list[1]) + position_embedding = position_embedding + block_embedding + + position_embedding = position_embedding.view( + concat([ + shape(input_ids, 0), + shape(input_ids, 1), + 4096, + ])) + + hidden_states = hidden_states + position_embedding + kv_cache_params.fill_none_tensor_list(len(self.layers)) if self.use_cache: @@ -230,30 +294,39 @@ class ChatGLMHeadModel(ChatGLMModel, GenerationMixin): def __init__(self, **args): if "args" not in args.keys(): - argNamespace = argparse.Namespace() + new_args = argparse.Namespace() for key, value in args.items(): - argNamespace.__setattr__(key, value) - assert "model_version" in args.keys(), "model_version not set" + new_args.__setattr__(key, value) + assert "model_name" in args.keys(), "model_name not set" # Other default values - argNamespace.norm_epsilon = 1.0e-5 - argNamespace.tokens_per_block = 64 - argNamespace.use_cache = True - if argNamespace.model_version == "1": - argNamespace.ffn_hidden_size = 16384 - argNamespace.linear_bias = True - argNamespace.max_seq_length = min( - 2048, argNamespace.max_position_embeddings) - argNamespace.num_kv_heads = 32 - argNamespace.qkv_bias = True - else: - argNamespace.apply_residual_connection_post_layernorm = False - argNamespace.ffn_hidden_size = 13696 - argNamespace.linear_bias = False - argNamespace.num_kv_heads = 2 - argNamespace.qkv_bias = True - argNamespace.rmsnorm = True - - args = argNamespace + new_args.norm_epsilon = 1.0e-5 + new_args.tokens_per_block = 64 + new_args.use_cache = True + if new_args.model_name in ["chatglm_6b"]: + new_args.ffn_hidden_size = 16384 + new_args.linear_bias = True + new_args.max_seq_length = min(2048, + new_args.max_position_embeddings) + new_args.num_kv_heads = 32 + new_args.qkv_bias = True + elif new_args.model_name in ["glm_10b"]: + new_args.ffn_hidden_size = 16384 + new_args.linear_bias = True + new_args.max_seq_length = min(1024, + new_args.max_position_embeddings) + new_args.num_kv_heads = 32 + new_args.qkv_bias = True + elif new_args.model_name in [ + "chatglm2_6b", "chatglm2_6b_32k", "chatglm3_6b", + "chatglm3_6b_base", "chatglm3_6b_32k" + ]: + new_args.apply_residual_connection_post_layernorm = False + new_args.ffn_hidden_size = 13696 + new_args.linear_bias = False + new_args.num_kv_heads = 2 + new_args.qkv_bias = True + new_args.rmsnorm = True + args = new_args else: args = args["args"] @@ -270,21 +343,21 @@ def init(self, args): self.kv_dtype = args.dtype self.dtype = self.kv_dtype + if isinstance(args.logits_dtype, str): + self.logits_dtype = str_dtype_to_trt(args.logits_dtype) + else: + assert isinstance(args.logits_dtype, trt.DataType) + self.logits_dtype = args.logits_dtype + if args.quant_mode.has_int8_kv_cache(): self.kv_dtype = str_dtype_to_trt('int8') elif args.quant_mode.has_fp8_kv_cache(): self.kv_dtype = str_dtype_to_trt('fp8') - if isinstance(args.logits_dtype, str): - self._logits_dtype = str_dtype_to_trt(args.logits_dtype) - else: - assert isinstance(args.logits_dtype, trt.DataType) - self._logits_dtype = args.logits_dtype - self.hidden_size = args.hidden_size self.mapping = args.mapping self.max_num_tokens = args.max_output_len + args.max_input_len - self.model_version = args.model_version + self.model_name = args.model_name self.num_heads = args.num_heads self.num_kv_heads = args.num_kv_heads self.num_layers = args.num_layers @@ -325,7 +398,7 @@ def forward( default_net().plugin_config.remove_input_padding) lm_logits = self.lm_head(hidden_states) - lm_logits.mark_output('logits', self._logits_dtype) + lm_logits.mark_output('logits', self.logits_dtype) if self.use_cache and default_net( ).plugin_config.paged_kv_cache == False: @@ -372,7 +445,7 @@ def prepare_inputs( mapping=self.mapping, max_num_tokens=self.max_num_tokens, prompt_embedding_table_size=0, - is_chatglm6b=(self.model_version == "1"), + position_encoding_2d=(self.model_name in ["chatglm_6b", "glm_10b"]), ) return (model_inputs['input_ids'], model_inputs['position_ids'], diff --git a/tensorrt_llm/models/enc_dec/model.py b/tensorrt_llm/models/enc_dec/model.py index 9f485a30bb68..efbfcaee4bdd 100644 --- a/tensorrt_llm/models/enc_dec/model.py +++ b/tensorrt_llm/models/enc_dec/model.py @@ -2,19 +2,20 @@ from collections import OrderedDict from typing import Optional -import numpy as np import tensorrt as trt from tensorrt_llm._common import default_net from tensorrt_llm._utils import str_dtype_to_trt from tensorrt_llm.functional import (LayerNormPositionType, LayerNormType, - PositionEmbeddingType, Tensor, assertion, - concat, constant, expand, expand_mask, - gather_last_token_logits, shape, slice) + MLPType, PositionEmbeddingType, Tensor, + assertion, gather_last_token_logits, recv, + send, shape) from tensorrt_llm.layers import (MLP, Attention, AttentionMaskType, AttentionParams, BertAttention, ColumnLinear, - Embedding, GroupNorm, KeyValueCacheParams, - LayerNorm, RmsNorm) + Embedding, FusedGatedMLP, GatedMLP, GroupNorm, + KeyValueCacheParams, LayerNorm, RmsNorm) +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.generation_mixin import GenerationMixin from tensorrt_llm.module import Module, ModuleList layernorm_map = { @@ -23,40 +24,65 @@ LayerNormType.GroupNorm: GroupNorm, } +mlp_map = { + MLPType.MLP: MLP, + MLPType.GatedMLP: GatedMLP, + MLPType.FusedGatedMLP: FusedGatedMLP, +} + class EncDecEmbedding(Module): - def __init__( - self, - vocab_size, - hidden_size, - max_position_embeddings=None, - has_position_embedding=False, - type_vocab_size=None, - has_embedding_layernorm=False, - has_embedding_scale=False, - layernorm_eps=1e-5, - layernorm_type=LayerNormType.LayerNorm, - dtype=None, - ): + def __init__(self, + vocab_size, + hidden_size, + max_position_embeddings=None, + has_position_embedding=False, + type_vocab_size=None, + has_embedding_layernorm=False, + has_embedding_scale=False, + layernorm_eps=1e-5, + layernorm_type=LayerNormType.LayerNorm, + dtype=None, + use_parallel_embedding=False, + embedding_sharding_dim=0, + mapping=Mapping()): super().__init__() self.layernorm_type = layernorm_type ln_type = layernorm_map[layernorm_type] - self.vocab_embedding = Embedding(vocab_size, hidden_size, dtype=dtype) + self.vocab_embedding = Embedding( + vocab_size, + hidden_size, + dtype=dtype, + tp_size=mapping.tp_size if use_parallel_embedding else 1, + tp_group=mapping.tp_group if use_parallel_embedding else None, + sharding_dim=embedding_sharding_dim, + tp_rank=mapping.tp_rank) + self.position_embedding = None self.max_position_embeddings = max_position_embeddings if has_position_embedding: - self.position_embedding = Embedding(max_position_embeddings, - hidden_size, - dtype=dtype) + self.position_embedding = Embedding( + max_position_embeddings, + hidden_size, + dtype=dtype, + tp_size=mapping.tp_size if use_parallel_embedding else 1, + tp_group=mapping.tp_group if use_parallel_embedding else None, + sharding_dim=embedding_sharding_dim, + tp_rank=mapping.tp_rank) self.token_type_embedding = None if type_vocab_size: - self.token_type_embedding = Embedding(type_vocab_size, - hidden_size, - dtype=dtype) + self.token_type_embedding = Embedding( + type_vocab_size, + hidden_size, + dtype=dtype, + tp_size=mapping.tp_size if use_parallel_embedding else 1, + tp_group=mapping.tp_group if use_parallel_embedding else None, + sharding_dim=embedding_sharding_dim, + tp_rank=mapping.tp_rank) # e.g. BART true, T5 false self.embedding_layernorm = None @@ -74,30 +100,8 @@ def __init__( # we just need to shrink its position embedding table by [offset:] during weight loading def forward(self, input_ids, position_ids=None, token_type_ids=None): - seq_len_2d = concat([1, shape(input_ids, 1)]) - - if self.position_embedding: - position_ids_buffer = constant( - np.expand_dims( - np.arange(self.max_position_embeddings).astype(np.int32), - 0)) - if position_ids is None: - # slice - position_ids = slice(position_ids_buffer, - starts=[0, 0], - sizes=seq_len_2d) - position_ids = expand(position_ids, shape(input_ids)) - - if self.token_type_embedding: - token_type_ids_buffer = constant( - np.expand_dims( - np.zeros(self.max_position_embeddings).astype(np.int32), 0)) - if token_type_ids is None: - # slice - token_type_ids = slice(token_type_ids_buffer, - starts=[0, 0], - sizes=seq_len_2d) - token_type_ids = expand(token_type_ids, shape(input_ids)) + # position_ids and token_type_ids are provided inputs + # and should not be formulated determinisitically x = self.vocab_embedding(input_ids) * self.embedding_scale if self.position_embedding: @@ -117,6 +121,7 @@ def __init__(self, ffn_hidden_size, num_attention_heads, num_kv_heads, + head_size, max_position_embeddings=None, q_scaling=1.0, has_attention_qkvo_bias=False, @@ -125,8 +130,8 @@ def __init__(self, layernorm_type=LayerNormType.LayerNorm, layernorm_eps=1e-5, hidden_act="relu", - tp_group=None, - tp_size=1, + mlp_type=MLPType.MLP, + mapping=Mapping(), dtype=None, residual_scaling=1.0, relative_attention=False, @@ -147,12 +152,14 @@ def __init__(self, self.attention = BertAttention( hidden_size, num_attention_heads, + attention_head_size=head_size, num_kv_heads=num_kv_heads, max_position_embeddings=max_position_embeddings, q_scaling=q_scaling, bias=has_attention_qkvo_bias, - tp_group=tp_group, - tp_size=tp_size, + tp_group=mapping.tp_group, + tp_size=mapping.tp_size, + tp_rank=mapping.tp_rank, dtype=dtype, relative_attention=relative_attention, max_distance=max_distance, @@ -162,15 +169,19 @@ def __init__(self, eps=layernorm_eps, dtype=dtype) - self.mlp = MLP( + # T5/BART MLP, Flan-T5 GatedMLP + self.mlp_type = mlp_type + mlp_f = mlp_map[mlp_type] + self.mlp = mlp_f( hidden_size=hidden_size, ffn_hidden_size=ffn_hidden_size, hidden_act=hidden_act, bias=has_mlp_bias, - tp_group=tp_group, - tp_size=tp_size, + tp_group=mapping.tp_group, + tp_size=mapping.tp_size, dtype=dtype, ) + self.mlp_layernorm = ln_type(normalized_shape=hidden_size, eps=layernorm_eps, dtype=dtype) @@ -180,7 +191,9 @@ def __init__(self, def forward(self, hidden_states: Tensor, attention_mask=None, - input_lengths=None): + input_lengths=None, + all_reduce_workspace=None, + max_input_length=None): assert isinstance(hidden_states, Tensor) # self attention @@ -189,11 +202,11 @@ def forward(self, if self.layernorm_position == LayerNormPositionType.pre_layernorm: hidden_states = self.attention_layernorm(hidden_states) - attention_output = self.attention( - hidden_states, - attention_mask=attention_mask, - input_lengths=input_lengths, - ) + attention_output = self.attention(hidden_states, + attention_mask=attention_mask, + input_lengths=input_lengths, + workspace=all_reduce_workspace, + max_input_length=max_input_length) hidden_states = residual + attention_output @@ -206,7 +219,7 @@ def forward(self, if self.layernorm_position == LayerNormPositionType.pre_layernorm: hidden_states = self.mlp_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) + hidden_states = self.mlp(hidden_states, workspace=all_reduce_workspace) hidden_states = residual + hidden_states @@ -223,6 +236,7 @@ def __init__(self, ffn_hidden_size, num_attention_heads, num_kv_heads, + head_size, max_position_embeddings=None, q_scaling=1.0, has_attention_qkvo_bias=False, @@ -231,8 +245,8 @@ def __init__(self, layernorm_type=LayerNormType.LayerNorm, layernorm_eps=1e-5, hidden_act="relu", - tp_group=None, - tp_size=1, + mlp_type=MLPType.MLP, + mapping=Mapping(), dtype=None, residual_scaling=1.0, relative_attention=False, @@ -253,13 +267,15 @@ def __init__(self, self.self_attention = Attention( hidden_size, num_attention_heads, + attention_head_size=head_size, num_kv_heads=num_kv_heads, max_position_embeddings=max_position_embeddings, q_scaling=q_scaling, bias=has_attention_qkvo_bias, attention_mask_type=AttentionMaskType.causal, - tp_group=tp_group, - tp_size=tp_size, + tp_group=mapping.tp_group, + tp_size=mapping.tp_size, + tp_rank=mapping.tp_rank, dtype=dtype, cross_attention=False, relative_attention=relative_attention, @@ -272,16 +288,23 @@ def __init__(self, eps=layernorm_eps, dtype=dtype) + # self attn uses MMHA, mask is always causal triangular + # cross attn has two scenarios: + # - in context phase, all ones mask, same as padding type + # - in generation phase, same causal triangular mask as MMHA + # - context phase special handling is done in plugin by resetting mask type self.cross_attention = Attention( hidden_size, num_attention_heads, - num_kv_heads=num_attention_heads, + attention_head_size=head_size, + num_kv_heads=num_kv_heads, max_position_embeddings=max_position_embeddings, q_scaling=q_scaling, bias=has_attention_qkvo_bias, attention_mask_type=AttentionMaskType.causal, - tp_group=tp_group, - tp_size=tp_size, + tp_group=mapping.tp_group, + tp_size=mapping.tp_size, + tp_rank=mapping.tp_rank, dtype=dtype, cross_attention=True, relative_attention= @@ -294,13 +317,16 @@ def __init__(self, eps=layernorm_eps, dtype=dtype) - self.mlp = MLP( + # T5/BART MLP, Flan-T5 GatedMLP + self.mlp_type = mlp_type + mlp_f = mlp_map[mlp_type] + self.mlp = mlp_f( hidden_size=hidden_size, ffn_hidden_size=ffn_hidden_size, hidden_act=hidden_act, bias=has_mlp_bias, - tp_group=tp_group, - tp_size=tp_size, + tp_group=mapping.tp_group, + tp_size=mapping.tp_size, dtype=dtype, ) @@ -310,15 +336,14 @@ def __init__(self, self.residual_scaling = residual_scaling - def forward( - self, - hidden_states: Tensor, - encoder_output: Optional[Tensor] = None, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - ): + def forward(self, + hidden_states: Tensor, + encoder_output: Optional[Tensor] = None, + attention_mask=None, + use_cache=False, + kv_cache_params=None, + attention_params=None, + all_reduce_workspace=None): assert isinstance(hidden_states, Tensor) if encoder_output: @@ -336,7 +361,7 @@ def forward( use_cache=use_cache, kv_cache_params=kv_cache_params, attention_params=attention_params, - ) + workspace=all_reduce_workspace) if use_cache: attention_output, presents_self = attention_output @@ -359,7 +384,7 @@ def forward( use_cache=use_cache, kv_cache_params=kv_cache_params, attention_params=attention_params, - ) + workspace=all_reduce_workspace) if use_cache: attention_output, presents_cross = attention_output @@ -375,7 +400,7 @@ def forward( if self.layernorm_position == LayerNormPositionType.pre_layernorm: hidden_states = self.mlp_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) + hidden_states = self.mlp(hidden_states, workspace=all_reduce_workspace) hidden_states = residual + hidden_states @@ -387,7 +412,7 @@ def forward( return hidden_states -class EncoderModel(Module): +class EncoderModel(Module, GenerationMixin): def __init__(self, num_layers, @@ -396,6 +421,7 @@ def __init__(self, ffn_hidden_size, vocab_size, dtype, + head_size=None, num_kv_heads=None, max_position_embeddings=None, has_position_embedding=False, @@ -413,10 +439,13 @@ def __init__(self, layernorm_position=LayerNormPositionType.pre_layernorm, layernorm_type=LayerNormType.LayerNorm, hidden_act="relu", - tp_group=None, - tp_size=1, - residual_scaling=1.0): + mlp_type=MLPType.MLP, + residual_scaling=1.0, + use_parallel_embedding=False, + embedding_sharding_dim=0, + mapping=Mapping()): super().__init__() + self.mapping = mapping self.has_position_embedding = has_position_embedding self.has_token_type_embedding = type_vocab_size is not None @@ -429,69 +458,101 @@ def __init__(self, self.has_attention_qkvo_bias = has_attention_qkvo_bias self.has_mlp_bias = has_mlp_bias + # e.g. BART false, T5 true + self.has_model_final_layernorm = has_model_final_layernorm + if isinstance(dtype, str): self._dtype = str_dtype_to_trt(dtype) else: assert isinstance(dtype, trt.DataType) self._dtype = dtype - self.num_layers = num_layers + self.total_num_layers = num_layers + self.num_layers = num_layers // self.mapping.pp_size - self.embedding = EncDecEmbedding( - vocab_size, - hidden_size, - max_position_embeddings=max_position_embeddings, - has_position_embedding=has_position_embedding, - type_vocab_size=type_vocab_size, - has_embedding_layernorm=has_embedding_layernorm, - has_embedding_scale=has_embedding_scale, - layernorm_eps=layernorm_eps, - layernorm_type=layernorm_type, - dtype=dtype) - - self.encoder_layers = ModuleList([ - EncoderLayer( - hidden_size=hidden_size, - ffn_hidden_size=ffn_hidden_size, - num_attention_heads=num_heads, - num_kv_heads=num_kv_heads if num_kv_heads else num_heads, + self.hidden_size = hidden_size + self.num_heads = num_heads + if num_kv_heads is None or num_kv_heads <= 0: + num_kv_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_size = self.hidden_size // self.num_heads if head_size is None else head_size + + if self.mapping.is_first_pp_rank(): + self.embedding = EncDecEmbedding( + vocab_size, + hidden_size, max_position_embeddings=max_position_embeddings, - q_scaling=q_scaling, - has_attention_qkvo_bias=has_attention_qkvo_bias, - has_mlp_bias=has_mlp_bias, - layernorm_position=layernorm_position, + has_position_embedding=has_position_embedding, + type_vocab_size=type_vocab_size, + has_embedding_layernorm=has_embedding_layernorm, + has_embedding_scale=has_embedding_scale, layernorm_eps=layernorm_eps, layernorm_type=layernorm_type, - hidden_act=hidden_act, - tp_group=tp_group, - tp_size=tp_size, dtype=dtype, - residual_scaling=residual_scaling, - relative_attention=relative_attention, - max_distance=max_distance, - num_buckets=num_buckets) for _ in range(num_layers) + use_parallel_embedding=use_parallel_embedding, + embedding_sharding_dim=embedding_sharding_dim, + mapping=self.mapping) + + self.encoder_layers = ModuleList([ + EncoderLayer(hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + num_attention_heads=num_heads, + num_kv_heads=num_kv_heads, + head_size=self.head_size, + max_position_embeddings=max_position_embeddings, + q_scaling=q_scaling, + has_attention_qkvo_bias=has_attention_qkvo_bias, + has_mlp_bias=has_mlp_bias, + layernorm_position=layernorm_position, + layernorm_eps=layernorm_eps, + layernorm_type=layernorm_type, + hidden_act=hidden_act, + mlp_type=mlp_type, + mapping=self.mapping, + dtype=dtype, + residual_scaling=residual_scaling, + relative_attention=relative_attention, + max_distance=max_distance, + num_buckets=num_buckets) for _ in + self.get_transformer_layers(self.mapping, self.total_num_layers) ]) - # e.g. BART false, T5 true - if has_model_final_layernorm: - self.final_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) + if self.mapping.is_last_pp_rank(): + if self.has_model_final_layernorm: + self.final_layernorm = ln_type(normalized_shape=hidden_size, + eps=layernorm_eps, + dtype=dtype) def forward(self, input_ids: Tensor, input_lengths=None, position_ids=None, - token_type_ids=None): - hidden_states = self.embedding(input_ids, position_ids, token_type_ids) - for layer_idx, encoder_layer in enumerate(self.encoder_layers): - hidden_states = encoder_layer(hidden_states=hidden_states, - input_lengths=input_lengths) - - if self.final_layernorm: - hidden_states = self.final_layernorm(hidden_states) - - hidden_states.mark_output('encoder_output', self._dtype) + token_type_ids=None, + hidden_states=None, + all_reduce_workspace=None, + max_input_length=None): + + # In PP, layer 0 has ids as inputs, all other layers have hidden_states as inputs + if self.mapping.is_first_pp_rank(): + hidden_states = self.embedding(input_ids, position_ids, + token_type_ids) + else: + hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) + + for encoder_layer in self.encoder_layers: + hidden_states = encoder_layer( + hidden_states=hidden_states, + input_lengths=input_lengths, + all_reduce_workspace=all_reduce_workspace, + max_input_length=max_input_length) + + if self.mapping.is_last_pp_rank(): + if self.final_layernorm: + hidden_states = self.final_layernorm(hidden_states) + hidden_states.mark_output('encoder_output', self._dtype) + else: + hidden_states = send(hidden_states, self.mapping.next_pp_rank()) + hidden_states.mark_output('hidden_states_output', self._dtype) return hidden_states @@ -502,6 +563,9 @@ def prepare_inputs(self, max_batch_size, max_input_len): @return: a list contains values which can be fed into the self.forward() ''' + num_heads = self.num_heads + head_size = self.head_size + bs_range = [1, (max_batch_size + 1) // 2, max_batch_size] inlen_range = [1, (max_input_len + 1) // 2, max_input_len] num_tokens_range = [ @@ -510,56 +574,95 @@ def prepare_inputs(self, max_batch_size, max_input_len): max_input_len * max_batch_size, ] - position_ids, token_type_ids = None, None + input_ids, position_ids, token_type_ids, hidden_states = None, None, None, None remove_input_padding = default_net().plugin_config.remove_input_padding + use_custom_all_reduce = default_net( + ).plugin_config.use_custom_all_reduce + if remove_input_padding: - input_ids = Tensor( - name="input_ids", - dtype=trt.int32, - shape=[1, -1], - dim_range=OrderedDict([("batch_size_fake", [1]), - ("num_tokens", [num_tokens_range])]), - ) - if self.has_position_embedding: - position_ids = Tensor( - name='position_ids', + if self.mapping.is_first_pp_rank(): + input_ids = Tensor( + name="input_ids", dtype=trt.int32, shape=[1, -1], dim_range=OrderedDict([('batch_size_fake', [1]), - ('num_tokens', [num_tokens_range])]), - ) - if self.has_token_type_embedding: - token_type_ids = Tensor( - name='token_type_ids', - dtype=trt.int32, - shape=[1, -1], - dim_range=OrderedDict([('batch_size_fake', [1]), - ('num_tokens', [num_tokens_range])]), + ("num_tokens", [num_tokens_range])]), ) + if self.has_position_embedding: + position_ids = Tensor( + name='position_ids', + dtype=trt.int32, + shape=[1, -1], + dim_range=OrderedDict([('batch_size_fake', [1]), + ('num_tokens', + [num_tokens_range])]), + ) + if self.has_token_type_embedding: + token_type_ids = Tensor( + name='token_type_ids', + dtype=trt.int32, + shape=[1, -1], + dim_range=OrderedDict([('batch_size_fake', [1]), + ('num_tokens', + [num_tokens_range])]), + ) + else: + hidden_states = Tensor(name='hidden_states_input', + dtype=self._dtype, + shape=[1, -1, head_size * num_heads], + dim_range=OrderedDict([ + ('batch_size_fake', [1]), + ('num_tokens', [num_tokens_range]), + ('hidden_size', + [head_size * num_heads]), + ])) else: - input_ids = Tensor( - name="input_ids", - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([("batch_size", [bs_range]), - ("input_len", [inlen_range])]), - ) - if self.has_position_embedding: - position_ids = Tensor( - name='position_ids', + if self.mapping.is_first_pp_rank(): + input_ids = Tensor( + name="input_ids", dtype=trt.int32, shape=[-1, -1], - dim_range=OrderedDict([('batch_size', [bs_range]), - ('input_len', [inlen_range])]), - ) - if self.has_token_type_embedding: - token_type_ids = Tensor( - name='token_type_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([('batch_size', [bs_range]), - ('input_len', [inlen_range])]), + dim_range=OrderedDict([("batch_size", [bs_range]), + ("input_len", [inlen_range])]), ) + if self.has_position_embedding: + position_ids = Tensor( + name='position_ids', + dtype=trt.int32, + shape=[-1, -1], + dim_range=OrderedDict([('batch_size', [bs_range]), + ('input_len', [inlen_range])]), + ) + if self.has_token_type_embedding: + token_type_ids = Tensor( + name='token_type_ids', + dtype=trt.int32, + shape=[-1, -1], + dim_range=OrderedDict([('batch_size', [bs_range]), + ('input_len', [inlen_range])]), + ) + else: + hidden_states = Tensor(name='hidden_states_input', + dtype=self._dtype, + shape=[-1, -1, head_size * num_heads], + dim_range=OrderedDict([ + ('batch_size', [bs_range]), + ('input_len', [inlen_range]), + ('hidden_size', + [head_size * num_heads]), + ])) + + all_reduce_workspace = None + if use_custom_all_reduce and self.mapping.tp_size > 1: + # 3 (= buffer + signals_in + signals_out) + workspace_size = 3 * self.mapping.tp_size + all_reduce_workspace = Tensor(name='all_reduce_workspace', + dtype=trt.int64, + shape=[workspace_size], + dim_range=OrderedDict([ + ('all_reduce_size', + [workspace_size]) + ])) input_lengths = Tensor( name="input_lengths", @@ -567,11 +670,18 @@ def prepare_inputs(self, max_batch_size, max_input_len): shape=[-1], dim_range=OrderedDict([("batch_size", [bs_range])]), ) + max_input_length = Tensor( + name="max_input_length", + dtype=trt.int32, + shape=[-1], + dim_range=OrderedDict([("max_input_length", [inlen_range])]), + ) - return (input_ids, input_lengths, position_ids, token_type_ids) + return (input_ids, input_lengths, position_ids, token_type_ids, + hidden_states, all_reduce_workspace, max_input_length) -class DecoderModel(Module): +class DecoderModel(Module, GenerationMixin): def __init__(self, num_layers, @@ -583,7 +693,10 @@ def __init__(self, vocab_size, dtype, logits_dtype='float32', + head_size=None, + encoder_head_size=None, num_kv_heads=None, + encoder_num_kv_heads=None, max_position_embeddings=None, has_position_embedding=False, relative_attention=False, @@ -600,11 +713,14 @@ def __init__(self, layernorm_position=LayerNormPositionType.pre_layernorm, layernorm_type=LayerNormType.LayerNorm, hidden_act="relu", + mlp_type=MLPType.MLP, has_lm_head_bias=False, - tp_group=None, - tp_size=1, - residual_scaling=1.0): + residual_scaling=1.0, + use_parallel_embedding=False, + embedding_sharding_dim=0, + mapping=Mapping()): super().__init__() + self.mapping = mapping self.has_position_embedding = has_position_embedding self.has_token_type_embedding = type_vocab_size is not None @@ -617,6 +733,9 @@ def __init__(self, self.has_attention_qkvo_bias = has_attention_qkvo_bias self.has_mlp_bias = has_mlp_bias + # e.g. BART false, T5 true + self.has_model_final_layernorm = has_model_final_layernorm + if isinstance(dtype, str): self._dtype = str_dtype_to_trt(dtype) else: @@ -632,98 +751,117 @@ def __init__(self, assert isinstance(logits_dtype, trt.DataType) self._logits_dtype = logits_dtype - self.num_layers = num_layers + self.total_num_layers = num_layers + self.num_layers = num_layers // self.mapping.pp_size + self.hidden_size = hidden_size self.num_heads = num_heads + if num_kv_heads is None or num_kv_heads <= 0: + num_kv_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_size = self.hidden_size // self.num_heads if head_size is None else head_size + self.encoder_hidden_size = encoder_hidden_size self.encoder_num_heads = encoder_num_heads - self.tp_size = tp_size + if encoder_num_kv_heads is None or encoder_num_kv_heads <= 0: + encoder_num_kv_heads = encoder_num_heads + self.encoder_num_kv_heads = encoder_num_kv_heads + self.encoder_head_size = self.encoder_hidden_size // self.num_heads if encoder_head_size is None else encoder_head_size self.has_position_embedding = has_position_embedding self.has_token_type_embedding = type_vocab_size is not None - self.embedding = EncDecEmbedding( - vocab_size, - hidden_size, - max_position_embeddings=max_position_embeddings, - has_position_embedding=has_position_embedding, - type_vocab_size=type_vocab_size, - has_embedding_layernorm=has_embedding_layernorm, - has_embedding_scale=has_embedding_scale, - layernorm_eps=layernorm_eps, - layernorm_type=layernorm_type, - dtype=dtype) - - self.decoder_layers = ModuleList([ - DecoderLayer( - hidden_size=hidden_size, - ffn_hidden_size=ffn_hidden_size, - num_attention_heads=num_heads, - num_kv_heads=num_kv_heads if num_kv_heads else num_heads, + if self.mapping.is_first_pp_rank(): + self.embedding = EncDecEmbedding( + vocab_size, + hidden_size, max_position_embeddings=max_position_embeddings, - q_scaling=q_scaling, - has_attention_qkvo_bias=has_attention_qkvo_bias, - has_mlp_bias=has_mlp_bias, - layernorm_position=layernorm_position, + has_position_embedding=has_position_embedding, + type_vocab_size=type_vocab_size, + has_embedding_layernorm=has_embedding_layernorm, + has_embedding_scale=has_embedding_scale, layernorm_eps=layernorm_eps, layernorm_type=layernorm_type, - hidden_act=hidden_act, - tp_group=tp_group, - tp_size=tp_size, dtype=dtype, - residual_scaling=residual_scaling, - relative_attention=relative_attention, - max_distance=max_distance, - num_buckets=num_buckets) for _ in range(num_layers) + use_parallel_embedding=use_parallel_embedding, + embedding_sharding_dim=embedding_sharding_dim, + mapping=self.mapping) + + self.decoder_layers = ModuleList([ + DecoderLayer(hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + num_attention_heads=num_heads, + num_kv_heads=self.num_kv_heads, + head_size=self.head_size, + max_position_embeddings=max_position_embeddings, + q_scaling=q_scaling, + has_attention_qkvo_bias=has_attention_qkvo_bias, + has_mlp_bias=has_mlp_bias, + layernorm_position=layernorm_position, + layernorm_eps=layernorm_eps, + layernorm_type=layernorm_type, + hidden_act=hidden_act, + mlp_type=mlp_type, + mapping=self.mapping, + dtype=dtype, + residual_scaling=residual_scaling, + relative_attention=relative_attention, + max_distance=max_distance, + num_buckets=num_buckets) for _ in + self.get_transformer_layers(self.mapping, self.total_num_layers) ]) - # e.g. BART false, T5 true - if has_model_final_layernorm: - self.final_layernorm = ln_type(normalized_shape=hidden_size, - eps=layernorm_eps, - dtype=dtype) + if self.mapping.is_last_pp_rank(): + if self.has_model_final_layernorm: + self.final_layernorm = ln_type(normalized_shape=hidden_size, + eps=layernorm_eps, + dtype=dtype) - self.lm_head = ColumnLinear( - hidden_size, - vocab_size, - bias=has_lm_head_bias, - dtype=dtype, - tp_group=tp_group, - tp_size=tp_size, - gather_output=True, - ) + self.lm_head = ColumnLinear( + hidden_size, + vocab_size, + bias=has_lm_head_bias, + dtype=dtype, + tp_group=mapping.tp_group, + tp_size=mapping.tp_size, + gather_output=True, + ) - def forward( - self, - decoder_input_ids: Tensor, - encoder_output: Tensor, - position_ids=None, - token_type_ids=None, - use_cache=False, - attention_mask=None, - last_token_ids=None, - kv_cache_params=None, - attention_params=None, - ): - assert last_token_ids is not None, "Expecting last token ids to be not None" - assert isinstance(decoder_input_ids, Tensor) + def forward(self, + decoder_input_ids: Tensor, + encoder_output: Tensor, + position_ids=None, + token_type_ids=None, + use_cache=False, + attention_mask=None, + last_token_ids=None, + kv_cache_params=None, + attention_params=None, + hidden_states=None, + all_reduce_workspace=None): + if self.mapping.is_first_pp_rank(): + assert isinstance(decoder_input_ids, Tensor) + else: + assert isinstance(hidden_states, Tensor) + + if self.mapping.is_last_pp_rank(): + assert last_token_ids is not None, "Expecting last token ids to be not None" - hidden_states = self.embedding(decoder_input_ids, position_ids, - token_type_ids) + # In PP, layer 0 has ids as inputs, all other layers have hidden_states as inputs + if self.mapping.is_first_pp_rank(): + hidden_states = self.embedding(decoder_input_ids, position_ids, + token_type_ids) + else: + hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - past_key_value = kv_cache_params.past_key_value - if past_key_value is None: - past_key_value = tuple([None] * len(self.decoder_layers)) + kv_cache_params.fill_none_tensor_list(len(self.decoder_layers)) if use_cache: presents = [] - if attention_mask is not None: - attention_mask = expand_mask(attention_mask, - shape(decoder_input_ids, -1)) - - for decoder_layer, past in zip(self.decoder_layers, - kv_cache_params.past_key_value): + for decoder_layer, past, max_kv_cache_length in zip( + self.decoder_layers, kv_cache_params.past_key_value, + kv_cache_params.host_max_kv_cache_lengths): hidden_states = decoder_layer( hidden_states, encoder_output=encoder_output, @@ -733,9 +871,10 @@ def forward( past_key_value=past, host_past_key_value_lengths=kv_cache_params. host_past_key_value_lengths, + host_max_kv_cache_lengths=max_kv_cache_length, cache_indirection=kv_cache_params.cache_indirection), attention_params=attention_params, - ) + all_reduce_workspace=all_reduce_workspace) if use_cache: presents_self, presents_cross = hidden_states[1], hidden_states[ @@ -743,39 +882,49 @@ def forward( presents.append((presents_self, presents_cross)) hidden_states = hidden_states[0] - if self.final_layernorm: - hidden_states = self.final_layernorm(hidden_states) + if self.mapping.is_last_pp_rank(): + if self.final_layernorm: + hidden_states = self.final_layernorm(hidden_states) - # [bs, seq, hidden_size] -> [bs, hidden_size] - hidden_states = gather_last_token_logits( - hidden_states, last_token_ids, - default_net().plugin_config.remove_input_padding) + # [bs, seq, hidden_size] or [1, num_tokens, hidden_size] -> [bs, hidden_size] + hidden_states = gather_last_token_logits( + hidden_states, last_token_ids, + default_net().plugin_config.remove_input_padding) - # Rescale output before projecting on vocab (for T5) - # See https://github.com/huggingface/transformers/blob/0b192de1f353b0e04dad4813e02e2c672de077be/src/transformers/models/t5/modeling_t5.py#L1769-L1772 - # Note: this is specific for T5, to make it more generic, one can pass in a config: - # self.config.tie_word_embeddings - default to be True for T5 - hidden_states = hidden_states * (self.hidden_size**-0.5) + # Rescale output before projecting on vocab (for T5) + # See https://github.com/huggingface/transformers/blob/0b192de1f353b0e04dad4813e02e2c672de077be/src/transformers/models/t5/modeling_t5.py#L1769-L1772 + # Note: this is specific for T5, to make it more generic, one can pass in a config: + # self.config.tie_word_embeddings - default to be True for T5 + hidden_states = hidden_states * (self.hidden_size**-0.5) - # [bs, hidden_size] -> [bs, vocab_size] - lm_logits = self.lm_head(hidden_states) - lm_logits.mark_output('logits', self._logits_dtype) - - if use_cache: - for i, present in enumerate(presents): + # [bs, hidden_size] -> [bs, vocab_size] + lm_logits = self.lm_head(hidden_states) + lm_logits.mark_output('logits', self._logits_dtype) + else: + hidden_states = send(hidden_states, self.mapping.next_pp_rank()) + hidden_states.mark_output('hidden_states_output', self._dtype) + + if use_cache and default_net().plugin_config.paged_kv_cache == False: + for i, present in zip( + self.get_transformer_layers(self.mapping, + self.total_num_layers), + presents): present[0].mark_output(f'present_key_value_{i}', self._kv_dtype) present[1].mark_output(f'cross_present_key_value_{i}', self._kv_dtype) - return (lm_logits, tuple(presents)) - - return lm_logits + if self.mapping.is_last_pp_rank(): + return (lm_logits, tuple(presents)) + return (hidden_states, tuple(presents)) + else: + if self.mapping.is_last_pp_rank(): + return lm_logits + return hidden_states def prepare_inputs( self, - num_layers, max_batch_size, max_beam_width, - max_input_len, + max_decoder_input_len, max_new_tokens, max_encoder_input_len, ): @@ -786,29 +935,42 @@ def prepare_inputs( ''' # Prepare inputs - max_output_len = max_input_len + max_new_tokens + max_output_len = max_decoder_input_len + max_new_tokens + + num_heads = self.num_heads + head_size = self.head_size + num_kv_heads = (self.num_kv_heads + self.mapping.tp_size - + 1) // self.mapping.tp_size + + self.encoder_num_heads + encoder_head_size = self.encoder_head_size + encoder_num_kv_heads = (self.encoder_num_kv_heads + self.mapping.tp_size + - 1) // self.mapping.tp_size - head_size = self.hidden_size // self.num_heads - num_heads = self.num_heads // self.tp_size - encoder_head_size = self.encoder_hidden_size // self.encoder_num_heads bb_range = [ 1, (max_batch_size * max_beam_width + 1) // 2, max_batch_size * max_beam_width ] bs_range = [1, (max_batch_size + 1) // 2, max_batch_size] beam_width_range = [1, (max_beam_width + 1) // 2, max_beam_width] - inlen_range = [1, 1, max_input_len - ] # context phase >= 1, generation phase = 1 + inlen_range = [ + 1, 1, max_decoder_input_len + ] # context phase >= 1 (if forced_input_ids), generation phase = 1 encoder_inlen_range = [ 1, (max_encoder_input_len + 1) // 2, max_encoder_input_len ] mask_len_range = [1, (max_output_len + 1) // 2 + 1, max_output_len + 1] max_output_len_range = [0, (max_output_len + 1) // 2, max_output_len] - num_tokens_range = [ + encoder_num_tokens_range = [ + 1, + (max_encoder_input_len * max_batch_size + 1) // 2, + max_encoder_input_len * max_batch_size, + ] + decoder_num_tokens_range = [ 1, max_batch_size * max_beam_width, - max(max_input_len * max_batch_size, + max(max_decoder_input_len * max_batch_size, max_beam_width * max_batch_size), ] @@ -824,58 +986,124 @@ def prepare_inputs( use_gpt_attention_plugin = default_net( ).plugin_config.gpt_attention_plugin remove_input_padding = default_net().plugin_config.remove_input_padding + use_custom_all_reduce = default_net( + ).plugin_config.use_custom_all_reduce - position_ids = None - token_type_ids = None + input_ids, position_ids, token_type_ids, hidden_states = None, None, None, None if remove_input_padding: - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[1, -1], - dim_range=OrderedDict([ - ('batch_size_fake', [1]), - ('num_tokens', [num_tokens_range]), - ])) - if self.has_position_embedding: - position_ids = Tensor(name='position_ids', - dtype=trt.int32, - shape=[1, -1], - dim_range=OrderedDict([ - ('batch_size_fake', [1]), - ('num_tokens', [num_tokens_range]), - ])) - if self.has_token_type_embedding: - token_type_ids = Tensor( - name='token_type_ids', - dtype=trt.int32, - shape=[1, -1], - dim_range=OrderedDict([('batch_size_fake', [1]), - ('num_tokens', [num_tokens_range])]), - ) + if self.mapping.is_first_pp_rank(): + input_ids = Tensor(name='input_ids', + dtype=trt.int32, + shape=[1, -1], + dim_range=OrderedDict([ + ('batch_size_fake', [1]), + ('decoder_num_tokens', + [decoder_num_tokens_range]), + ])) + if self.has_position_embedding: + position_ids = Tensor(name='position_ids', + dtype=trt.int32, + shape=[1, -1], + dim_range=OrderedDict([ + ('batch_size_fake', [1]), + ('decoder_num_tokens', + [decoder_num_tokens_range]), + ])) + if self.has_token_type_embedding: + token_type_ids = Tensor( + name='token_type_ids', + dtype=trt.int32, + shape=[1, -1], + dim_range=OrderedDict([('batch_size_fake', [1]), + ('decoder_num_tokens', + [decoder_num_tokens_range])]), + ) + else: + hidden_states = Tensor(name='hidden_states_input', + dtype=self._dtype, + shape=[1, -1, head_size * num_heads], + dim_range=OrderedDict([ + ('batch_size_fake', [1]), + ('decoder_num_tokens', + [decoder_num_tokens_range]), + ('hidden_size', + [head_size * num_heads]), + ])) else: - input_ids = Tensor(name='input_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('input_len', [inlen_range]), - ])) - if self.has_position_embedding: - position_ids = Tensor(name='position_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([ - ('batch_size_beam_width', [bb_range]), - ('input_len', [inlen_range]), - ])) - if self.has_token_type_embedding: - token_type_ids = Tensor( - name='token_type_ids', - dtype=trt.int32, - shape=[-1, -1], - dim_range=OrderedDict([('batch_size_beam_width', [bb_range + if self.mapping.is_first_pp_rank(): + input_ids = Tensor(name='input_ids', + dtype=trt.int32, + shape=[-1, -1], + dim_range=OrderedDict([ + ('batch_size_beam_width', [bb_range]), + ('input_len', [inlen_range]), + ])) + if self.has_position_embedding: + position_ids = Tensor(name='position_ids', + dtype=trt.int32, + shape=[-1, -1], + dim_range=OrderedDict([ + ('batch_size_beam_width', + [bb_range]), + ('input_len', [inlen_range]), + ])) + if self.has_token_type_embedding: + token_type_ids = Tensor( + name='token_type_ids', + dtype=trt.int32, + shape=[-1, -1], + dim_range=OrderedDict([('batch_size_beam_width', + [bb_range]), + ('input_len', [inlen_range])]), + ) + else: + hidden_states = Tensor(name='hidden_states_input', + dtype=self._dtype, + shape=[-1, -1, head_size * num_heads], + dim_range=OrderedDict([ + ('batch_size_beam_width', [bb_range ]), - ('input_len', [inlen_range])]), - ) + ('input_len', [inlen_range]), + ('hidden_size', + [head_size * num_heads]), + ])) + + encoder_input_lengths = Tensor( + name="encoder_input_lengths", + dtype=trt.int32, + shape=[-1], + dim_range=OrderedDict([("batch_size_beam_width", [bb_range])]), + ) + encoder_max_input_length = Tensor( + name="encoder_max_input_length", + dtype=trt.int32, + shape=[-1], + dim_range=OrderedDict([("encoder_max_input_length", + [encoder_inlen_range])]), + ) + encoder_output = None + if remove_input_padding: + encoder_output = Tensor( + name="encoder_output", + dtype=self._dtype, + shape=[-1, -1, self.encoder_hidden_size], + dim_range=OrderedDict([ + ("batch_size_fake", [1]), + ("encoder_num_tokens", [encoder_num_tokens_range]), + ("encoder_hidden_size", [self.encoder_hidden_size]), + ]), + ) + else: + encoder_output = Tensor( + name="encoder_output", + dtype=self._dtype, + shape=[-1, -1, self.encoder_hidden_size], + dim_range=OrderedDict([ + ("batch_size", [bs_range]), + ("encoder_input_len", [encoder_input_len_range]), + ("encoder_hidden_size", [self.encoder_hidden_size]), + ]), + ) if use_gpt_attention_plugin: host_past_key_value_lengths = Tensor( @@ -919,25 +1147,16 @@ def prepare_inputs( [bb_range]) ])) - encoder_input_lengths = Tensor( - name="encoder_input_lengths", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("batch_size", [bs_range])]), - ) - encoder_max_input_length = Tensor( - name="encoder_max_input_length", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("encoder_max_input_length", - [encoder_inlen_range])]), - ) - last_token_ids = Tensor( - name="last_token_ids", - dtype=trt.int32, - shape=[-1], - dim_range=OrderedDict([("batch_size_last_token_ids", [bb_range])]), - ) + last_token_ids = None + if self.mapping.is_last_pp_rank(): + last_token_ids = Tensor( + name="last_token_ids", + dtype=trt.int32, + shape=[-1], + dim_range=OrderedDict([("batch_size_last_token_ids", [bb_range]) + ]), + ) + if not use_gpt_attention_plugin: attention_mask = Tensor( name='attention_mask', @@ -960,58 +1179,76 @@ def prepare_inputs( ]), ) - encoder_output = Tensor( - name="encoder_output", - dtype=self._dtype, - shape=[-1, -1, self.encoder_hidden_size], - dim_range=OrderedDict([ - ("batch_size", [bs_range]), - ("encoder_input_len", [encoder_input_len_range]), - ("encoder_hidden_size", [self.encoder_hidden_size]), - ]), - ) + all_reduce_workspace = None + if use_custom_all_reduce and self.mapping.tp_size > 1: + # 3 (= buffer + signals_in + signals_out) + workspace_size = 3 * self.mapping.tp_size + all_reduce_workspace = Tensor(name='all_reduce_workspace', + dtype=trt.int64, + shape=[workspace_size], + dim_range=OrderedDict([ + ('all_reduce_size', + [workspace_size]) + ])) + + layers_range = self.get_transformer_layers(self.mapping, + self.total_num_layers) + + if use_gpt_attention_plugin: + host_max_kv_cache_lengths = [] + for i in layers_range: + host_kv_cache_length_tensor = Tensor( + name=f'host_max_kv_cache_length_{i}', + dtype=trt.int32, + shape=[1], + dim_range=OrderedDict([('scalar', [1])])) + host_max_kv_cache_lengths.append(host_kv_cache_length_tensor) - for i in range(num_layers): + for i in layers_range: kv_dim_range = OrderedDict([ ('batch_size_beam_width', [bb_range]), ('kv', [2]), - ('num_heads', [num_heads]), + ('num_heads', [num_kv_heads]), ('past_key_len', [max_output_len_range]), ('head_size', [head_size]), ]) kv = Tensor(name=f'past_key_value_{i}', dtype=self._kv_dtype, - shape=[-1, 2, num_heads, -1, head_size], + shape=[-1, 2, num_kv_heads, -1, head_size], dim_range=kv_dim_range) cross_kv_dim_range = OrderedDict([ ('batch_size_beam_width', [bb_range]), ('kv', [2]), - ('cross_num_heads', [self.encoder_num_heads]), + ('cross_num_heads', [encoder_num_kv_heads]), ('cross_past_key_len', [encoder_input_len_range]), ('cross_head_size', [encoder_head_size]), ]) cross_kv = Tensor( name=f'cross_past_key_value_{i}', dtype=self._kv_dtype, - shape=[-1, 2, self.encoder_num_heads, -1, encoder_head_size], + shape=[-1, 2, encoder_num_kv_heads, -1, encoder_head_size], dim_range=cross_kv_dim_range) past_key_value.append((kv, cross_kv)) # TODO: Remove this when TRT fix the named dimension if not remove_input_padding: - assertion(shape(input_ids, 0) == shape(kv, 0), 'batch size') + assertion( + shape( + input_ids if self.mapping.is_first_pp_rank() else + hidden_states, 0) == shape(kv, 0), 'batch size') kv_cache_params = KeyValueCacheParams( past_key_value=past_key_value, host_past_key_value_lengths=host_past_key_value_lengths, + host_max_kv_cache_lengths=host_max_kv_cache_lengths, cache_indirection=cache_indirection) attention_params = AttentionParams( sequence_length=sequence_length, context_lengths=context_lengths, host_context_lengths=host_context_lengths, - max_context_length=max_input_len, + max_context_length=max_decoder_input_len, host_request_types=host_request_types, encoder_input_lengths=encoder_input_lengths, encoder_max_input_length=encoder_max_input_length, @@ -1019,4 +1256,4 @@ def prepare_inputs( return (input_ids, encoder_output, position_ids, token_type_ids, True, attention_mask, last_token_ids, kv_cache_params, - attention_params) + attention_params, hidden_states, all_reduce_workspace) diff --git a/tensorrt_llm/models/generation_mixin.py b/tensorrt_llm/models/generation_mixin.py index 06962b95bf12..5d6e4175f32d 100644 --- a/tensorrt_llm/models/generation_mixin.py +++ b/tensorrt_llm/models/generation_mixin.py @@ -51,7 +51,8 @@ def prepare_basic_inputs(self, mapping=Mapping(), max_num_tokens=None, prompt_embedding_table_size: int = 0, - is_chatglm6b=False): + position_encoding_2d=False, + use_lora_plugin: bool = False): max_len = max_input_len + max_new_tokens @@ -134,7 +135,7 @@ def prepare_basic_inputs(self, [1, 1] if enable_two_optimization_profiles else [1]), ('num_tokens', num_tokens_range), ])) - if is_chatglm6b: + if position_encoding_2d: position_ids = Tensor( name='position_ids', dtype=trt.int32, @@ -184,7 +185,7 @@ def prepare_basic_inputs(self, ('batch_size_beam_width', bb_range), ('input_len', inlen_range), ])) - if is_chatglm6b: + if position_encoding_2d: position_ids = Tensor( name='position_ids', dtype=trt.int32, @@ -381,7 +382,7 @@ def prepare_basic_inputs(self, if use_gpt_attention_plugin: host_max_kv_cache_lengths = [] - for i in range(num_layers): + for i in layers_range: host_kv_cache_length_tensor = Tensor( name=f'host_max_kv_cache_length_{i}', dtype=trt.int32, @@ -467,6 +468,29 @@ def prepare_basic_inputs(self, [1, 1] if enable_two_optimization_profiles else [1]) ])) + lora_weights_pointers_list = None + lora_ranks = None + if use_lora_plugin: + lora_weights_pointers_list = [] + for i in layers_range: + lora_weights_pointers = Tensor( + name=f'lora_weights_pointers_{i}', + dtype=trt.int64, + shape=[-1, 2], + dim_range=OrderedDict([ + ('batch_size_beam_width', bb_range), + ('in_out', + [2, 2] if enable_two_optimization_profiles else [2]), + ])) + lora_weights_pointers_list.append(lora_weights_pointers) + + lora_ranks = Tensor( + name='lora_ranks', + dtype=trt.int32, + shape=[-1], + dim_range=OrderedDict([('batch_size_beam_width', bb_range)]), + ) + return { 'input_ids': input_ids, 'hidden_states_input': hidden_states, @@ -486,4 +510,6 @@ def prepare_basic_inputs(self, 'tasks': tasks, 'prompt_vocab_size': prompt_vocab_size, 'all_reduce_workspace': all_reduce_workspace, + 'lora_ranks': lora_ranks, + 'lora_weights_pointers_list': lora_weights_pointers_list, } diff --git a/tensorrt_llm/models/gpt/model.py b/tensorrt_llm/models/gpt/model.py index be2678390d54..19fcb7daf00f 100644 --- a/tensorrt_llm/models/gpt/model.py +++ b/tensorrt_llm/models/gpt/model.py @@ -23,7 +23,8 @@ is_gated_activation, non_gated_version) from ...layers import (MLP, Attention, AttentionMaskType, AttentionParams, ColumnLinear, Embedding, GatedMLP, KeyValueCacheParams, - LayerNorm, PositionEmbeddingType, PromptTuningEmbedding) + LayerNorm, LoraParams, PositionEmbeddingType, + PromptTuningEmbedding) from ...mapping import Mapping from ...module import Module, ModuleList from ...quantization import QuantMode @@ -110,9 +111,11 @@ def __init__(self, position_embedding_type=PositionEmbeddingType.learned_absolute, quant_mode=QuantMode(0), rotary_embedding_percentage=1.0, + rotary_base=10000.0, + rotary_scaling=None, inter_size=None, bias=True, - multi_query_mode=False, + num_kv_heads=None, tp_group=None, tp_size=1, tp_rank=0, @@ -135,7 +138,7 @@ def __init__(self, self.attention = Attention( hidden_size, num_attention_heads, - 1 if multi_query_mode else num_attention_heads, + num_kv_heads, max_position_embeddings, num_layers, apply_query_key_layer_scaling, @@ -143,6 +146,8 @@ def __init__(self, attention_mask_type=attention_mask_type, position_embedding_type=position_embedding_type, rotary_embedding_percentage=rotary_embedding_percentage, + rotary_embedding_base=rotary_base, + rotary_embedding_scaling=rotary_scaling, bias=bias, tp_group=tp_group, tp_size=tp_size, @@ -172,7 +177,8 @@ def forward(self, use_cache=False, kv_cache_params=None, attention_params=None, - workspace=None): + workspace=None, + lora_params=None): assert isinstance(hidden_states, Tensor) @@ -185,7 +191,8 @@ def forward(self, use_cache=use_cache, kv_cache_params=kv_cache_params, attention_params=attention_params, - workspace=workspace) + workspace=workspace, + lora_params=lora_params) if use_cache: attention_output, presents = attention_output @@ -218,10 +225,12 @@ def __init__(self, apply_query_key_layer_scaling=False, position_embedding_type=PositionEmbeddingType.learned_absolute, rotary_embedding_percentage=1.0, + rotary_base=10000.0, + rotary_scaling=None, inter_size=None, bias=True, quant_mode=QuantMode(0), - multi_query_mode=False, + num_kv_heads=None, use_prompt_tuning=False, use_parallel_embedding=False, embedding_sharding_dim=0): @@ -255,7 +264,9 @@ def __init__(self, hidden_act=hidden_act, position_embedding_type=position_embedding_type, rotary_embedding_percentage=rotary_embedding_percentage, - multi_query_mode=multi_query_mode, + rotary_base=rotary_base, + rotary_scaling=rotary_scaling, + num_kv_heads=num_kv_heads, tp_group=mapping.tp_group, tp_size=mapping.tp_size, tp_rank=mapping.tp_rank, @@ -278,7 +289,8 @@ def forward(self, prompt_embedding_table=None, prompt_tasks=None, prompt_vocab_size=None, - workspace=None): + workspace=None, + lora_params=None): hidden_states = self.embedding(input_ids, position_ids, @@ -292,10 +304,18 @@ def forward(self, if use_cache: presents = [] - for layer, past, pointer, max_kv_cache_length in zip( - self.layers, kv_cache_params.past_key_value, - kv_cache_params.kv_cache_block_pointers, - kv_cache_params.host_max_kv_cache_lengths): + for layer_idx, (layer, past, pointer, max_kv_cache_length) in enumerate( + zip(self.layers, kv_cache_params.past_key_value, + kv_cache_params.kv_cache_block_pointers, + kv_cache_params.host_max_kv_cache_lengths)): + lora_param = None + if lora_params.lora_ranks is not None: + lora_param = LoraParams( + lora_ranks=lora_params.lora_ranks, + lora_weights_pointers_list=[ + lora_params.lora_weights_pointers_list[layer_idx] + ]) + hidden_states = layer( hidden_states, use_cache=use_cache, @@ -308,7 +328,8 @@ def forward(self, kv_cache_block_pointers=[pointer], cache_indirection=kv_cache_params.cache_indirection), attention_params=attention_params, - workspace=workspace) + workspace=workspace, + lora_params=lora_param) if use_cache: presents.append(hidden_states[1]) @@ -336,10 +357,12 @@ def __init__(self, apply_query_key_layer_scaling=False, position_embedding_type=PositionEmbeddingType.learned_absolute, rotary_embedding_percentage=1.0, + rotary_base=10000.0, + rotary_scaling=None, inter_size=None, bias=True, quant_mode=QuantMode(0), - multi_query_mode=False, + num_kv_heads=None, use_prompt_tuning=False, use_parallel_embedding=False, embedding_sharding_dim=0, @@ -376,14 +399,30 @@ def __init__(self, self._hidden_size = hidden_size self._vocab_size = vocab_size self._tp_size = mapping.tp_size - self._multi_query_mode = multi_query_mode - - super().__init__(num_layers, num_heads, hidden_size, vocab_size, - hidden_act, max_position_embeddings, dtype, mapping, - apply_query_key_layer_scaling, position_embedding_type, - rotary_embedding_percentage, inter_size, bias, - quant_mode, multi_query_mode, use_prompt_tuning, - use_parallel_embedding, embedding_sharding_dim) + self._num_kv_heads = num_kv_heads if num_kv_heads else num_heads + + super().__init__( + num_layers=num_layers, + num_heads=num_heads, + hidden_size=hidden_size, + vocab_size=vocab_size, + hidden_act=hidden_act, + max_position_embeddings=max_position_embeddings, + dtype=dtype, + mapping=mapping, + apply_query_key_layer_scaling=apply_query_key_layer_scaling, + position_embedding_type=position_embedding_type, + rotary_embedding_percentage=rotary_embedding_percentage, + rotary_base=rotary_base, + rotary_scaling=rotary_scaling, + inter_size=inter_size, + bias=bias, + quant_mode=quant_mode, + num_kv_heads=num_kv_heads, + use_prompt_tuning=use_prompt_tuning, + use_parallel_embedding=use_parallel_embedding, + embedding_sharding_dim=embedding_sharding_dim, + ) vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) share_weight = None @@ -409,13 +448,15 @@ def forward(self, prompt_embedding_table=None, prompt_tasks=None, prompt_vocab_size=None, - workspace=None): + workspace=None, + lora_params=None): hidden_states = super().forward(input_ids, position_ids, use_cache, attention_mask, kv_cache_params, attention_params, prompt_embedding_table, prompt_tasks, - prompt_vocab_size, workspace) + prompt_vocab_size, workspace, + lora_params) if use_cache: hidden_states, presents = hidden_states @@ -454,7 +495,7 @@ def prepare_inputs(self, # Prepare inputs head_size = self._hidden_size // self._num_heads - num_heads_kv = 1 if self._multi_query_mode else self._num_heads + num_heads_kv = self._num_kv_heads remove_input_padding = default_net().plugin_config.remove_input_padding use_gpt_attention_plugin = default_net( ).plugin_config.gpt_attention_plugin @@ -463,6 +504,7 @@ def prepare_inputs(self, tokens_per_block = default_net().plugin_config.tokens_per_block use_custom_all_reduce = default_net( ).plugin_config.use_custom_all_reduce + use_lora_plugin = default_net().plugin_config.lora_plugin model_inputs = self.prepare_basic_inputs( max_batch_size=max_batch_size, @@ -484,7 +526,8 @@ def prepare_inputs(self, gather_all_token_logits=gather_all_token_logits, mapping=self.mapping, max_num_tokens=max_num_tokens, - prompt_embedding_table_size=prompt_embedding_table_size) + prompt_embedding_table_size=prompt_embedding_table_size, + use_lora_plugin=use_lora_plugin) return (model_inputs['input_ids'], model_inputs['position_ids'], True, model_inputs['last_token_ids'], model_inputs['attention_mask'], @@ -506,4 +549,6 @@ def prepare_inputs(self, host_request_types=model_inputs['host_request_types']), model_inputs['prompt_embedding_table'], model_inputs['tasks'], model_inputs['prompt_vocab_size'], - model_inputs['all_reduce_workspace']) + model_inputs['all_reduce_workspace'], + LoraParams(model_inputs['lora_ranks'], + model_inputs['lora_weights_pointers_list'])) diff --git a/tensorrt_llm/models/internlm/model.py b/tensorrt_llm/models/internlm/model.py deleted file mode 100644 index 2324757e4245..000000000000 --- a/tensorrt_llm/models/internlm/model.py +++ /dev/null @@ -1,427 +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 tensorrt as trt - -from ..._common import default_net -from ..._utils import pad_vocab_size, str_dtype_to_trt -from ...functional import gather_last_token_logits, recv, send -from ...layers import (Attention, AttentionMaskType, AttentionParams, - ColumnLinear, Embedding, GatedMLP, KeyValueCacheParams, - PositionEmbeddingType, RmsNorm) -from ...mapping import Mapping -from ...module import Module, ModuleList -from ...quantization import QuantMode -from ..generation_mixin import GenerationMixin - - -class InternLMDecoderLayer(Module): - - def __init__(self, - layer_id, - hidden_size, - num_attention_heads, - num_kv_heads=None, - max_position_embeddings=2048, - dtype=None, - attention_mask_type=AttentionMaskType.causal, - hidden_act='silu', - attn_bias=True, - position_embedding_type=PositionEmbeddingType.rope_gpt_neox, - rotary_base=10000.0, - rotary_scaling=None, - mlp_hidden_size=None, - tp_group=None, - tp_size=1, - quant_mode=QuantMode(0), - rms_norm_eps=1e-06): - super().__init__() - self._layer_id = layer_id # useful for debugging - # used for quantizing model - self.hidden_size = hidden_size - self.num_attention_heads = num_attention_heads - self.num_kv_heads = num_kv_heads - self.max_position_embeddings = max_position_embeddings - self.dtype = dtype - self.hidden_act = hidden_act - self.tp_group = tp_group - self.tp_size = tp_size - self.mlp_hidden_size = mlp_hidden_size - self.attention_mask_type = attention_mask_type - self.position_embedding_type = position_embedding_type - self.input_layernorm = RmsNorm(normalized_shape=hidden_size, - eps=rms_norm_eps, - dtype=dtype) - - self.attention = Attention( - hidden_size, - num_attention_heads, - num_kv_heads, - max_position_embeddings, - dtype=dtype, - attention_mask_type=AttentionMaskType.causal, - bias=attn_bias, - position_embedding_type=position_embedding_type, - rotary_embedding_base=rotary_base, - rotary_embedding_scaling=rotary_scaling, - tp_group=tp_group, - tp_size=tp_size, - use_int8_kv_cache=quant_mode.has_int8_kv_cache(), - quant_mode=quant_mode, - instance_id=2 * layer_id, - ) - if not mlp_hidden_size: - self.mlp_hidden_size = hidden_size * 4 - self.mlp = GatedMLP(hidden_size=hidden_size, - ffn_hidden_size=self.mlp_hidden_size, - hidden_act=hidden_act, - dtype=dtype, - bias=False, - tp_group=tp_group, - tp_size=tp_size, - quant_mode=quant_mode, - instance_id=2 * layer_id + 1) - self.post_layernorm = RmsNorm(normalized_shape=hidden_size, - eps=rms_norm_eps, - dtype=dtype) - - def forward(self, - hidden_states, - attention_mask=None, - use_cache=False, - kv_cache_params=None, - attention_params=None, - all_reduce_workspace=None): - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - if self._layer_id == 0: - self.register_network_output(f"norm0", hidden_states) - - attention_output = self.attention(hidden_states, - attention_mask=attention_mask, - use_cache=use_cache, - kv_cache_params=kv_cache_params, - attention_params=attention_params, - workspace=all_reduce_workspace) - - if use_cache: - attention_output, presents = attention_output - if self._layer_id == 0: - self.register_network_output(f"attn", attention_output) - - hidden_states = residual + attention_output - - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - if self._layer_id == 0: - self.register_network_output(f"norm1", hidden_states) - - hidden_states = self.mlp(hidden_states, all_reduce_workspace) - if self._layer_id == 0: - self.register_network_output(f"mlp", hidden_states) - - hidden_states = residual + hidden_states - if use_cache: - return (hidden_states, presents) - return hidden_states - - -class InternLMModel(Module): - - def __init__(self, - num_layers, - num_heads, - num_kv_heads, - hidden_size, - vocab_size, - hidden_act, - attn_bias, - max_position_embeddings, - dtype, - mlp_hidden_size=None, - position_embedding_type=PositionEmbeddingType.rope_gpt_neox, - rotary_base=10000.0, - rotary_scaling=None, - mapping=Mapping(), - quant_mode=QuantMode(0), - use_parallel_embedding=False, - embedding_sharding_dim=0, - rms_norm_eps=1e-06): - super().__init__() - self.mapping = mapping - - if self.mapping.is_first_pp_rank(): - self.vocab_embedding = Embedding( - num_embeddings=vocab_size, - embedding_dim=hidden_size, - dtype=dtype, - tp_size=mapping.tp_size if use_parallel_embedding else 1, - tp_group=mapping.tp_group if use_parallel_embedding else None, - sharding_dim=embedding_sharding_dim, - tp_rank=mapping.tp_rank) - - self.layers = ModuleList([ - InternLMDecoderLayer( - layer_id=i, - hidden_size=hidden_size, - num_attention_heads=num_heads, - num_kv_heads=num_kv_heads, - max_position_embeddings=max_position_embeddings, - dtype=dtype, - hidden_act=hidden_act, - attn_bias=attn_bias, - mlp_hidden_size=mlp_hidden_size, - position_embedding_type=position_embedding_type, - rotary_base=rotary_base, - rotary_scaling=rotary_scaling, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - quant_mode=quant_mode, - rms_norm_eps=rms_norm_eps) - for i in self.get_transformer_layers(self.mapping, num_layers) - ]) - - if self.mapping.is_last_pp_rank(): - self.ln_f = RmsNorm(normalized_shape=hidden_size, - eps=rms_norm_eps, - dtype=dtype) - - def forward(self, - input_ids, - position_ids=None, - use_cache=False, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - all_reduce_workspace=None): - - if kv_cache_params.past_key_value is None: - tuple([None] * len(self.layers)) - - if use_cache: - presents = [] - - if self.mapping.is_first_pp_rank(): - hidden_states = self.vocab_embedding(input_ids) - else: - hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) - self.register_network_output(f"embd", hidden_states) - - for layer, past, pointer in zip( - self.layers, kv_cache_params.past_key_value, - kv_cache_params.kv_cache_block_pointers): - hidden_states = layer( - hidden_states, - use_cache=use_cache, - attention_mask=attention_mask, - kv_cache_params=KeyValueCacheParams( - past_key_value=[past], - host_past_key_value_lengths=kv_cache_params. - host_past_key_value_lengths, - kv_cache_block_pointers=[pointer], - cache_indirection=kv_cache_params.cache_indirection), - attention_params=attention_params, - all_reduce_workspace=all_reduce_workspace) - - if use_cache: - presents.append(hidden_states[1]) - hidden_states = hidden_states[0] - - if self.mapping.is_last_pp_rank(): - hidden_states = self.ln_f(hidden_states) - else: - hidden_states = send(hidden_states, self.mapping.next_pp_rank()) - - if use_cache: - return (hidden_states, tuple(presents)) - return hidden_states - - -class InternLMForCausalLM(InternLMModel, GenerationMixin): - - def __init__(self, - num_layers, - num_heads, - num_kv_heads, - hidden_size, - vocab_size, - hidden_act, - attn_bias, - max_position_embeddings, - dtype, - logits_dtype="float32", - mlp_hidden_size=None, - position_embedding_type=PositionEmbeddingType.rope_gpt_neox, - rotary_base=10000.0, - rotary_scaling=None, - mapping=Mapping(), - quant_mode=QuantMode(0), - use_parallel_embedding=False, - embedding_sharding_dim=0, - rms_norm_eps=1e-06): - - if isinstance(dtype, str): - self.dtype = str_dtype_to_trt(dtype) - else: - assert isinstance(dtype, trt.DataType) - self.dtype = dtype - - if isinstance(logits_dtype, str): - self.logits_dtype = str_dtype_to_trt(logits_dtype) - else: - assert isinstance(logits_dtype, trt.DataType) - self.logits_dtype = logits_dtype - - self.num_layers = num_layers - self.num_heads = num_heads - if num_kv_heads is None or num_kv_heads <= 0: - num_kv_heads = num_heads - self.num_kv_heads = num_kv_heads - self.hidden_size = hidden_size - self.attn_bias = attn_bias - self.vocab_size = vocab_size - self.tp_size = mapping.tp_size - - self.kv_dtype = self.dtype - if quant_mode.has_int8_kv_cache(): - self.kv_dtype = str_dtype_to_trt('int8') - elif quant_mode.has_fp8_kv_cache(): - self.kv_dtype = str_dtype_to_trt('fp8') - - self.quant_mode = quant_mode - self.use_parallel_embedding = use_parallel_embedding - self.embedding_sharding_dim = embedding_sharding_dim - - super().__init__(num_layers, num_heads, num_kv_heads, hidden_size, - vocab_size, hidden_act, attn_bias, - max_position_embeddings, dtype, mlp_hidden_size, - position_embedding_type, rotary_base, rotary_scaling, - mapping, quant_mode, use_parallel_embedding, - embedding_sharding_dim, rms_norm_eps) - - vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) - if self.mapping.is_last_pp_rank(): - self.lm_head = ColumnLinear(hidden_size, - vocab_size_padded, - bias=False, - dtype=dtype, - tp_group=mapping.tp_group, - tp_size=mapping.tp_size, - gather_output=True) - - def forward(self, - input_ids, - position_ids=None, - use_cache=False, - last_token_ids=None, - attention_mask=None, - kv_cache_params=None, - attention_params=None, - hidden_states=None, - all_reduce_workspace=None): - hidden_states = super().forward(input_ids, position_ids, use_cache, - attention_mask, kv_cache_params, - attention_params, hidden_states, - all_reduce_workspace) - - if use_cache: - hidden_states, presents = hidden_states - - if self.mapping.is_last_pp_rank(): - hidden_states = gather_last_token_logits( - hidden_states, last_token_ids, - default_net().plugin_config.remove_input_padding) - - # [batch_size, hidden_size] -> [batch_size, vocab_size] - lm_logits = self.lm_head(hidden_states) - lm_logits.mark_output('logits', self.logits_dtype) - else: - hidden_states.mark_output('hidden_states_output', self.dtype) - - if use_cache and default_net().plugin_config.paged_kv_cache == False: - for i, present in zip( - self.get_transformer_layers(self.mapping, self.num_layers), - presents): - present.mark_output(f'present_key_value_{i}', self.kv_dtype) - if self.mapping.is_last_pp_rank(): - return (lm_logits, presents) - return (hidden_states, presents) - else: - if self.mapping.is_last_pp_rank(): - return lm_logits - return hidden_states - - def prepare_inputs(self, - max_batch_size, - max_input_len, - max_new_tokens, - use_cache, - max_beam_width, - max_num_tokens: int = None): - '''@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() - ''' - - # Prepare inputs - head_size = self.hidden_size // self.num_heads - remove_input_padding = default_net().plugin_config.remove_input_padding - use_gpt_attention_plugin = default_net( - ).plugin_config.gpt_attention_plugin - use_gemm_plugin = default_net().plugin_config.gemm_plugin - paged_kv_cache = default_net().plugin_config.paged_kv_cache - tokens_per_block = default_net().plugin_config.tokens_per_block - use_custom_all_reduce = default_net( - ).plugin_config.use_custom_all_reduce - - model_inputs = self.prepare_basic_inputs( - max_batch_size, - max_beam_width, - max_input_len, - max_new_tokens, - self.num_kv_heads, - head_size, - self.num_layers, - self.kv_dtype, - remove_input_padding=remove_input_padding, - use_gpt_attention_plugin=use_gpt_attention_plugin, - use_gemm_plugin=use_gemm_plugin, - use_custom_all_reduce=use_custom_all_reduce, - paged_kv_cache=paged_kv_cache, - tokens_per_block=tokens_per_block, - dtype=self.dtype, - num_heads=self.num_heads, - mapping=self.mapping, - max_num_tokens=max_num_tokens) - - return (model_inputs['input_ids'], model_inputs['position_ids'], True, - model_inputs['last_token_ids'], model_inputs['attention_mask'], - KeyValueCacheParams( - past_key_value=model_inputs['past_key_value'], - host_past_key_value_lengths=model_inputs[ - 'host_past_key_value_lengths'], - kv_cache_block_pointers=model_inputs[ - 'kv_cache_block_pointers_list'], - cache_indirection=model_inputs['cache_indirection'], - ), - AttentionParams( - sequence_length=model_inputs['sequence_length'], - context_lengths=model_inputs['context_lengths'], - host_context_lengths=model_inputs['host_context_lengths'], - max_context_length=max_input_len, - host_request_types=model_inputs['host_request_types']), - model_inputs['hidden_states_input'], - model_inputs['all_reduce_workspace']) diff --git a/tensorrt_llm/models/llama/model.py b/tensorrt_llm/models/llama/model.py index 548aa947cf4f..a91ebd41ce43 100644 --- a/tensorrt_llm/models/llama/model.py +++ b/tensorrt_llm/models/llama/model.py @@ -48,6 +48,8 @@ def __init__(self, tp_size=1, quant_mode=QuantMode(0), rms_norm_eps=1e-06, + attn_bias=False, + mlp_bias=False, use_fused_mlp=False): super().__init__() self._layer_id = layer_id # useful for debugging @@ -74,7 +76,7 @@ def __init__(self, max_position_embeddings, dtype=dtype, attention_mask_type=AttentionMaskType.causal, - bias=False, + bias=attn_bias, position_embedding_type=position_embedding_type, rotary_embedding_base=rotary_base, rotary_embedding_scaling=rotary_scaling, @@ -91,7 +93,7 @@ def __init__(self, ffn_hidden_size=self.mlp_hidden_size, hidden_act=hidden_act, dtype=dtype, - bias=False, + bias=mlp_bias, tp_group=tp_group, tp_size=tp_size, quant_mode=quant_mode, @@ -162,6 +164,8 @@ def __init__(self, embedding_sharding_dim=0, rms_norm_eps=1e-06, use_fused_mlp=False, + attn_bias=False, + mlp_bias=False, use_prompt_tuning: bool = False): super().__init__() self.mapping = mapping @@ -197,6 +201,8 @@ def __init__(self, tp_size=mapping.tp_size, quant_mode=quant_mode, rms_norm_eps=rms_norm_eps, + attn_bias=attn_bias, + mlp_bias=mlp_bias, use_fused_mlp=use_fused_mlp) for i in self.get_transformer_layers(self.mapping, num_layers) ]) @@ -292,6 +298,8 @@ def __init__(self, embedding_sharding_dim=0, rms_norm_eps=1e-06, use_fused_mlp=False, + attn_bias=False, + mlp_bias=False, use_prompt_tuning: bool = False): if isinstance(dtype, str): @@ -330,7 +338,8 @@ def __init__(self, mlp_hidden_size, position_embedding_type, rotary_base, rotary_scaling, mapping, quant_mode, use_parallel_embedding, embedding_sharding_dim, - rms_norm_eps, use_fused_mlp, use_prompt_tuning) + rms_norm_eps, use_fused_mlp, attn_bias, mlp_bias, + use_prompt_tuning) vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) if self.mapping.is_last_pp_rank(): @@ -391,16 +400,15 @@ def forward( return lm_logits return hidden_states - def prepare_inputs( - self, - max_batch_size, - max_input_len, - max_new_tokens, - use_cache, - max_beam_width, - max_num_tokens: int = None, - prompt_embedding_table_size: int = 0, - ): + def prepare_inputs(self, + max_batch_size, + max_input_len, + max_new_tokens, + use_cache, + max_beam_width, + max_num_tokens: int = None, + prompt_embedding_table_size: int = 0, + gather_all_token_logits: bool = False): '''@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. @@ -438,6 +446,7 @@ def prepare_inputs( mapping=self.mapping, max_num_tokens=max_num_tokens, prompt_embedding_table_size=prompt_embedding_table_size, + gather_all_token_logits=gather_all_token_logits, ) return ( diff --git a/tensorrt_llm/models/quantized/ammo.py b/tensorrt_llm/models/quantized/ammo.py index fecdd34eb909..b9df697ee1b6 100644 --- a/tensorrt_llm/models/quantized/ammo.py +++ b/tensorrt_llm/models/quantized/ammo.py @@ -97,6 +97,10 @@ def quantize_and_export(model: torch.nn.Module, model_type = "gpt2" elif "Falcon" in model_cls_name or "RW" in model_cls_name: model_type = "falcon" + elif "ChatGLM" in model_cls_name: + model_type = "chatglm" + elif "MPT" in model_cls_name: + model_type = "mpt" else: raise NotImplementedError( f"Deploying quantized model {model_cls_name} is not supported") diff --git a/tensorrt_llm/models/quantized/quant.py b/tensorrt_llm/models/quantized/quant.py index 246eee0421c6..11ef2f68550e 100644 --- a/tensorrt_llm/models/quantized/quant.py +++ b/tensorrt_llm/models/quantized/quant.py @@ -18,8 +18,8 @@ from ...layers import ColumnLinear, RowLinear from ...models import (BaichuanForCausalLM, BloomForCausalLM, FalconForCausalLM, - GPTJForCausalLM, GPTLMHeadModel, InternLMForCausalLM, - LLaMAForCausalLM) + GPTJForCausalLM, GPTLMHeadModel, LLaMAForCausalLM, + QWenForCausalLM) from ...module import Module from ...quantization import QuantMode from ...quantization.layers import FP8Linear, FP8RowLinear @@ -210,10 +210,52 @@ def _smooth_quantize_internlm(model, quant_mode): return model +def _smooth_quantize_qwen(model, quant_mode): + assert quant_mode.has_act_and_weight_quant() + for layer in model.layers: + assert hasattr(layer, "ln_1"), "The layer has no ln_1" + layer.ln_1 = SmoothQuantRmsNorm(normalized_shape=layer.hidden_size, + dtype=layer.dtype, + quant_mode=quant_mode) + assert hasattr(layer, "attention"), "The layer has no attention" + layer.attention = SmoothQuantAttention( + layer.hidden_size, + layer.num_attention_heads, + max_position_embeddings=layer.max_position_embeddings, + num_layers=layer.num_layers, + apply_query_key_layer_scaling=layer.apply_query_key_layer_scaling, + attention_mask_type=layer.attention_mask_type, + bias=layer.bias, + qkv_bias_only=True, + dtype=layer.dtype, + position_embedding_type=layer.position_embedding_type, + tp_group=layer.tp_group, + tp_size=layer.tp_size, + quant_mode=quant_mode) + assert hasattr(layer, "mlp"), "The layer has no mlp" + layer.mlp = SmoothQuantGatedMLP(hidden_size=layer.hidden_size, + ffn_hidden_size=layer.mlp_hidden_size // + 2, + hidden_act=layer.hidden_act, + dtype=layer.dtype, + bias=layer.bias, + tp_group=layer.tp_group, + tp_size=layer.tp_size, + quant_mode=quant_mode) + assert hasattr(layer, "ln_2"), "The layer has no ln_2" + layer.ln_2 = SmoothQuantRmsNorm(normalized_shape=layer.hidden_size, + dtype=layer.dtype, + quant_mode=quant_mode) + + setattr(model, 'quant_mode', quant_mode) + return model + + def _smooth_quantize(model, quant_mode): assert isinstance(model, GPTLMHeadModel) or isinstance(model, LLaMAForCausalLM) \ - or isinstance(model, BloomForCausalLM) or isinstance(model, BaichuanForCausalLM) or isinstance(model, InternLMForCausalLM), \ - "Only GPTLMHeadModel, LLaMAForCausalLM BloomForCausalLM, InternLMForCausalLM and BaichuanForCausalLM are well tested now" + or isinstance(model, BloomForCausalLM) or isinstance(model, BaichuanForCausalLM) \ + or isinstance(model, QWenForCausalLM), \ + "Only GPTLMHeadModel, LLaMAForCausalLM BloomForCausalLM and BaichuanForCausalLM are well tested now" if isinstance(model, GPTLMHeadModel): return _smooth_quantize_gpt(model, quant_mode) elif isinstance(model, LLaMAForCausalLM): @@ -222,8 +264,8 @@ def _smooth_quantize(model, quant_mode): return _smooth_quantize_bloom(model, quant_mode) elif isinstance(model, BaichuanForCausalLM): return _smooth_quantize_baichuan(model, quant_mode) - elif isinstance(model, InternLMForCausalLM): - return _smooth_quantize_internlm(model, quant_mode) + elif isinstance(model, QWenForCausalLM): + return _smooth_quantize_qwen(model, quant_mode) else: assert False, f"Model {type(model).__name__} is not supported by SmoothQuant yet" diff --git a/tensorrt_llm/models/qwen/__init__.py b/tensorrt_llm/models/qwen/__init__.py new file mode 100644 index 000000000000..2a36ca922710 --- /dev/null +++ b/tensorrt_llm/models/qwen/__init__.py @@ -0,0 +1,14 @@ +# 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. diff --git a/tensorrt_llm/models/qwen/model.py b/tensorrt_llm/models/qwen/model.py new file mode 100644 index 000000000000..ba13e9c3c889 --- /dev/null +++ b/tensorrt_llm/models/qwen/model.py @@ -0,0 +1,640 @@ +# 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 math + +import tensorrt as trt + +from ..._common import default_net +from ..._utils import pad_vocab_size, str_dtype_to_trt +from ...functional import (RotaryScalingType, Tensor, gather_last_token_logits, + gpt_attention, partial, recv, send, unary) +from ...layers import (AttentionMaskType, AttentionParams, ColumnLinear, + Embedding, GatedMLP, KeyValueCacheParams, + PositionEmbeddingType, RmsNorm, RowLinear) +from ...mapping import Mapping +from ...module import Module, ModuleList +from ...parameter import Parameter +from ...quantization import QuantMode +from ...quantization.layers import FP8Linear, FP8RowLinear +from ..generation_mixin import GenerationMixin + +log = partial(unary, op=trt.UnaryOperation.LOG) +ceil = partial(unary, op=trt.UnaryOperation.CEIL) + + +class QWenAttention(Module): + + def __init__( + self, + hidden_size, + num_attention_heads, + max_position_embeddings, + seq_length, # 2048 + num_kv_heads=None, + num_layers=1, + apply_query_key_layer_scaling=False, + attention_mask_type=AttentionMaskType.causal, + bias=True, + dtype=None, + position_embedding_type=PositionEmbeddingType.rope_gpt_neox, + rotary_embedding_base=10000.0, + rotary_embedding_scaling=None, + neox_rotary_style=False, + use_int8_kv_cache=False, + rotary_embedding_percentage=1.0, + tp_group=None, + tp_size=1, + quant_mode: QuantMode = QuantMode(0), + q_scaling=1.0, + cross_attention=False, + relative_attention=False, + max_distance=0, + num_buckets=0, + instance_id: int = 0, + use_dynamic_ntk=True, + use_logn_attn=True, + ): + super().__init__() + self.cross_attention = cross_attention + self.seq_length = seq_length + self.hidden_size = hidden_size + self.num_heads = num_attention_heads + + self.attention_mask_type = attention_mask_type + self.bias = bias + self.attention_head_size = hidden_size // num_attention_heads + self.num_attention_heads = num_attention_heads // tp_size + self.num_attention_kv_heads = ( + num_kv_heads + tp_size - 1 + ) // tp_size if num_kv_heads is not None else self.num_attention_heads + self.hidden_size = hidden_size // tp_size + self.max_position_embeddings = max_position_embeddings + + self.num_layers = num_layers + self.apply_query_key_layer_scaling = apply_query_key_layer_scaling + self.norm_factor = math.sqrt(self.attention_head_size) + self.q_scaling = q_scaling + if self.apply_query_key_layer_scaling: + self.norm_factor *= self.num_layers + self.q_scaling *= self.num_layers + + self.position_embedding_type = position_embedding_type + + self.relative_attention = relative_attention + self.max_distance = max_distance + + self.rotary_embedding_base = rotary_embedding_base + self.rotary_embedding_scale_type = RotaryScalingType.none + self.rotary_embedding_scale = 1.0 + if rotary_embedding_scaling is not None: + assert rotary_embedding_scaling["type"] in ["linear", "dynamic"] + self.rotary_embedding_scale_type = RotaryScalingType.linear if rotary_embedding_scaling[ + "type"] == "linear" else RotaryScalingType.dynamic + self.rotary_embedding_scale = rotary_embedding_scaling["factor"] + assert self.rotary_embedding_scale > 1.0 + self.rotary_embedding_dim = 0 + self.neox_rotary_style = neox_rotary_style + if self.position_embedding_type == PositionEmbeddingType.rope_gpt_neox: + self.rotary_embedding_dim = int(self.attention_head_size * + rotary_embedding_percentage) + + self.dtype = dtype + self.quant_mode = quant_mode + if use_int8_kv_cache: + # TODO: remove use_int8_kv_cache as can be replaced by quant_mode.has_kv_cache_quant() + # Merge int8 setting into quant_mode + self.quant_mode = self.quant_mode.set_int8_kv_cache() + + self.use_int8_kv_cache = use_int8_kv_cache + if self.use_int8_kv_cache: + self.kv_orig_quant_scale = Parameter(shape=(1, ), dtype='float32') + self.kv_quant_orig_scale = Parameter(shape=(1, ), dtype='float32') + else: + self.register_parameter('kv_orig_quant_scale', None) + self.register_parameter('kv_quant_orig_scale', None) + + self.use_fp8_qdq = self.quant_mode.has_fp8_qdq() + if self.use_fp8_qdq: + self.qkv = FP8Linear(hidden_size, + hidden_size + + (2 * tp_size * self.num_attention_kv_heads * + self.attention_head_size), + bias=True, + dtype=dtype, + tp_group=tp_group, + tp_size=tp_size, + gather_output=False) + self.dense = FP8RowLinear(hidden_size, + hidden_size, + bias=bias, + dtype=dtype, + tp_group=tp_group, + tp_size=tp_size, + instance_id=instance_id) + else: + self.qkv = ColumnLinear(hidden_size, + hidden_size + + (2 * tp_size * self.num_attention_kv_heads * + self.attention_head_size), + bias=True, + dtype=dtype, + tp_group=tp_group, + tp_size=tp_size, + gather_output=False) + self.dense = RowLinear(hidden_size, + hidden_size, + bias=bias, + dtype=dtype, + tp_group=tp_group, + tp_size=tp_size, + instance_id=instance_id) + + if relative_attention: + self.rel_attn_table = Parameter(shape=(num_attention_heads // + tp_size, num_buckets), + dtype=dtype) + + self.use_dynamic_ntk = use_dynamic_ntk + self.use_logn_attn = use_logn_attn + + def forward( + self, + hidden_states: Tensor, + use_cache=False, + kv_cache_params=None, + attention_params=None, + workspace=None, + ): + if not default_net().plugin_config.gpt_attention_plugin: + raise ValueError('QWen is only supported with GPTAttention plugin') + + assert isinstance(hidden_states, Tensor) + qkv = self.qkv(hidden_states) + + kv_orig_quant_scale = self.kv_orig_quant_scale.value if self.use_int8_kv_cache else None + kv_quant_orig_scale = self.kv_quant_orig_scale.value if self.use_int8_kv_cache else None + + # return outputs + context, past_key_value = gpt_attention( + tensor=qkv, + past_key_value=kv_cache_params.get_first_past_key_value(), + sequence_length=attention_params.sequence_length, + host_past_key_value_lengths=kv_cache_params. + host_past_key_value_lengths, + host_max_kv_cache_lengths=kv_cache_params.host_max_kv_cache_lengths, + context_lengths=attention_params.context_lengths, + cache_indirection=kv_cache_params.cache_indirection, + host_request_types=attention_params.host_request_types, + num_heads=self.num_attention_heads, + num_kv_heads=self.num_attention_kv_heads, + hidden_size_per_head=self.attention_head_size, + q_scaling=self.q_scaling, + rotary_embedding_dim=self. + rotary_embedding_dim, # when we use it 0, we will not use rotary embedding in plugin + rotary_embedding_scale_type=self.neox_rotary_style, + rotary_embedding_max_positions=self.max_position_embeddings, + position_embedding_type=PositionEmbeddingType.rope_gpt_neox, + kv_orig_quant_scale=kv_orig_quant_scale, + kv_quant_orig_scale=kv_quant_orig_scale, + kv_cache_quant_mode=QuantMode.from_description( + use_int8_kv_cache=self.use_int8_kv_cache), + kv_cache_block_pointers=kv_cache_params. + get_first_kv_cache_block_pointers(), + max_context_length=attention_params.max_context_length, + mask_type=self.attention_mask_type.value, + host_context_lengths=attention_params.host_context_lengths) + + context = self.dense(context, workspace=workspace) + + if use_cache: + return (context, past_key_value) + else: + return context + + +class QWenBlock(Module): + + def __init__(self, + layer_id, + hidden_size, + seq_length, + num_attention_heads, + max_position_embeddings, + num_layers, + dtype=None, + attention_mask_type=AttentionMaskType.causal, + apply_query_key_layer_scaling=False, + hidden_act='silu', + position_embedding_type=PositionEmbeddingType.rope_gpt_neox, + rotary_base=10000.0, + rotary_scaling=None, + quant_mode=QuantMode(0), + mlp_hidden_size=None, + neox_rotary_style=True, + bias=False, + tp_group=None, + tp_size=1, + rms_norm_eps=1e-06): + super().__init__() + self._layer_id = layer_id # useful for debugging + self.hidden_size = hidden_size + self.seq_length = seq_length + self.mlp_hidden_size = mlp_hidden_size + self.neox_rotary_style = neox_rotary_style + self.bias = bias + self.hidden_act = hidden_act + self.dtype = dtype + self.attention_mask_type = attention_mask_type + self.apply_query_key_layer_scaling = apply_query_key_layer_scaling + self.tp_group = tp_group + self.tp_size = tp_size + self.num_attention_heads = num_attention_heads + self.max_position_embeddings = max_position_embeddings + self.num_layers = num_layers + self.position_embedding_type = position_embedding_type + + self.ln_1 = RmsNorm(normalized_shape=hidden_size, + eps=rms_norm_eps, + dtype=dtype) + + self.attention = QWenAttention( + hidden_size=self.hidden_size, + num_attention_heads=self.num_attention_heads, + max_position_embeddings=self.max_position_embeddings, + num_layers=self.num_layers, + seq_length=self.seq_length, + dtype=self.dtype, + attention_mask_type=self.attention_mask_type, + bias=bias, + position_embedding_type=self.position_embedding_type, + rotary_embedding_base=rotary_base, + rotary_embedding_scaling=rotary_scaling, + neox_rotary_style=neox_rotary_style, + tp_group=self.tp_group, + tp_size=self.tp_size, + use_int8_kv_cache=quant_mode.has_int8_kv_cache(), + ) + if not mlp_hidden_size: + mlp_hidden_size = hidden_size * 4 + + self.mlp = GatedMLP(hidden_size=hidden_size, + ffn_hidden_size=mlp_hidden_size // 2, + hidden_act=hidden_act, + dtype=dtype, + bias=False, + tp_group=tp_group, + tp_size=tp_size, + quant_mode=quant_mode, + instance_id=2 * layer_id + 1) + self.ln_2 = RmsNorm(normalized_shape=hidden_size, + eps=rms_norm_eps, + dtype=dtype) + + def forward( + self, + hidden_states: Tensor, + use_cache=False, + kv_cache_params=None, + attention_params=None, + all_reduce_workspace=None, + ): + residual = hidden_states + hidden_states = self.ln_1(hidden_states) + attention_output = self.attention( + hidden_states, + use_cache=use_cache, + kv_cache_params=kv_cache_params, + attention_params=attention_params, + workspace=all_reduce_workspace, + ) + if use_cache: + attention_output, presents = attention_output + + hidden_states = residual + attention_output + + residual = hidden_states + + hidden_states = self.ln_2(hidden_states) + + hidden_states = self.mlp(hidden_states) + + hidden_states = residual + hidden_states + if use_cache: + return (hidden_states, presents) + return hidden_states + + +class QWenModel(Module): + + def __init__( + self, + num_layers, + num_heads, + hidden_size, + seq_length, + vocab_size, + hidden_act, + max_position_embeddings, + dtype, + mlp_hidden_size=None, + position_embedding_type=PositionEmbeddingType.rope_gpt_neox, + neox_rotary_style=True, + bias=False, + rotary_base=10000.0, + rotary_scaling=None, + mapping=Mapping(), + quant_mode=QuantMode(0), + use_parallel_embedding=False, + embedding_sharding_dim=0, + rms_norm_eps=1e-06, + ): + super().__init__() + self.mapping = mapping + if self.mapping.is_first_pp_rank(): + self.vocab_embedding = Embedding( + num_embeddings=vocab_size, + embedding_dim=hidden_size, + dtype=dtype, + tp_size=mapping.tp_size if use_parallel_embedding else 1, + tp_group=mapping.tp_group if use_parallel_embedding else None, + sharding_dim=embedding_sharding_dim, + tp_rank=mapping.tp_rank) + + self.layers = ModuleList([ + QWenBlock(layer_id=i, + hidden_size=hidden_size, + seq_length=seq_length, + num_attention_heads=num_heads, + num_layers=num_layers, + max_position_embeddings=max_position_embeddings, + dtype=dtype, + hidden_act=hidden_act, + quant_mode=quant_mode, + mlp_hidden_size=mlp_hidden_size, + position_embedding_type=position_embedding_type, + rotary_base=rotary_base, + rotary_scaling=rotary_scaling, + neox_rotary_style=neox_rotary_style, + bias=bias, + tp_group=mapping.tp_group, + tp_size=mapping.tp_size, + rms_norm_eps=rms_norm_eps) + for i in self.get_transformer_layers(self.mapping, num_layers) + ]) + + self.ln_f = RmsNorm(normalized_shape=hidden_size, + eps=rms_norm_eps, + dtype=dtype) + + def forward(self, + input_ids, + position_ids=None, + use_cache=False, + kv_cache_params=None, + attention_params=None, + hidden_states=None, + all_reduce_workspace=None): + + if kv_cache_params.past_key_value is None: + tuple([None] * len(self.layers)) + + kv_cache_params.fill_none_tensor_list(len(self.layers)) + + if use_cache: + presents = [] + + if self.mapping.is_first_pp_rank(): + hidden_states = self.vocab_embedding(input_ids) + else: + hidden_states = recv(hidden_states, self.mapping.prev_pp_rank()) + self.register_network_output(f"embd", hidden_states) + + for layer, past, pointer, max_kv_cache_length in zip( + self.layers, kv_cache_params.past_key_value, + kv_cache_params.kv_cache_block_pointers, + kv_cache_params.host_max_kv_cache_lengths): + hidden_states = layer( + hidden_states, + use_cache=use_cache, + kv_cache_params=KeyValueCacheParams( + past_key_value=[past], + host_past_key_value_lengths=kv_cache_params. + host_past_key_value_lengths, + host_max_kv_cache_lengths=max_kv_cache_length, + kv_cache_block_pointers=[pointer], + cache_indirection=kv_cache_params.cache_indirection), + attention_params=attention_params, + all_reduce_workspace=all_reduce_workspace) + + if use_cache: + presents.append(hidden_states[1]) + hidden_states = hidden_states[0] + + if self.mapping.is_last_pp_rank(): + hidden_states = self.ln_f(hidden_states) + else: + hidden_states = send(hidden_states, self.mapping.next_pp_rank()) + + if use_cache: + return (hidden_states, tuple(presents)) + return hidden_states + + +class QWenForCausalLM(QWenModel, GenerationMixin): + + def __init__( + self, + num_layers, + num_heads, + num_kv_heads, + hidden_size, + seq_length, + vocab_size, + hidden_act, + max_position_embeddings, + dtype, + logits_dtype="float32", + mlp_hidden_size=None, + neox_rotary_style=True, + rotary_base=10000.0, + rotary_scaling=None, + mapping=Mapping(), + quant_mode=QuantMode(0), + use_parallel_embedding=False, + embedding_sharding_dim=0, + rms_norm_eps=1e-06, + ): + self.mapping = mapping + if isinstance(dtype, str): + self.dtype = str_dtype_to_trt(dtype) + else: + assert isinstance(dtype, trt.DataType) + self.dtype = dtype + if isinstance(logits_dtype, str): + self.logits_dtype = str_dtype_to_trt(logits_dtype) + else: + assert isinstance(logits_dtype, trt.DataType) + self.logits_dtype = logits_dtype + self.num_layers = num_layers + self.num_heads = num_heads + if num_kv_heads is None or num_kv_heads <= 0: + num_kv_heads = num_heads + self.num_kv_heads = num_kv_heads + self.hidden_size = hidden_size + self.vocab_size = vocab_size + self.tp_size = mapping.tp_size + + self.kv_dtype = self.dtype + if quant_mode.has_int8_kv_cache(): + self.kv_dtype = str_dtype_to_trt('int8') + elif quant_mode.has_fp8_kv_cache(): + self.kv_dtype = str_dtype_to_trt('fp8') + self.quant_mode = quant_mode + self.use_parallel_embedding = use_parallel_embedding + self.embedding_sharding_dim = embedding_sharding_dim + + super().__init__(num_layers=num_layers, + num_heads=num_heads, + hidden_size=hidden_size, + seq_length=seq_length, + vocab_size=vocab_size, + hidden_act=hidden_act, + max_position_embeddings=max_position_embeddings, + dtype=dtype, + mlp_hidden_size=mlp_hidden_size, + neox_rotary_style=neox_rotary_style, + rotary_base=rotary_base, + rotary_scaling=rotary_scaling, + mapping=mapping, + quant_mode=quant_mode, + use_parallel_embedding=use_parallel_embedding, + embedding_sharding_dim=embedding_sharding_dim, + rms_norm_eps=rms_norm_eps) + vocab_size_padded = pad_vocab_size(vocab_size, mapping.tp_size) + if self.mapping.is_last_pp_rank(): + self.lm_head = ColumnLinear(hidden_size, + vocab_size_padded, + bias=False, + dtype=dtype, + tp_group=mapping.tp_group, + tp_size=mapping.tp_size, + gather_output=True) + + def forward(self, + input_ids, + position_ids=None, + use_cache=False, + last_token_ids=None, + kv_cache_params=None, + attention_params=None, + hidden_states=None, + all_reduce_workspace=None): + hidden_states = super().forward(input_ids, position_ids, use_cache, + kv_cache_params, attention_params, + hidden_states, all_reduce_workspace) + if use_cache: + hidden_states, presents = hidden_states + + if self.mapping.is_last_pp_rank(): + hidden_states = gather_last_token_logits( + hidden_states, last_token_ids, + default_net().plugin_config.remove_input_padding) + + # [batch_size, hidden_size] -> [batch_size, vocab_size] + lm_logits = self.lm_head(hidden_states) + lm_logits.mark_output('logits', self.logits_dtype) + else: + hidden_states.mark_output('hidden_states_output', self.dtype) + + if use_cache and default_net().plugin_config.paged_kv_cache == False: + for i, present in zip( + self.get_transformer_layers(self.mapping, self.num_layers), + presents): + present.mark_output(f'present_key_value_{i}', self.kv_dtype) + if self.mapping.is_last_pp_rank(): + return (lm_logits, presents) + return (hidden_states, presents) + else: + if self.mapping.is_last_pp_rank(): + return lm_logits + return hidden_states + + def prepare_inputs( + self, + max_batch_size, + max_input_len, + max_new_tokens, + use_cache, + max_beam_width: int = 1, + max_num_tokens: int = None, + ): + '''@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() + ''' + + # Prepare inputs + head_size = self.hidden_size // self.num_heads + remove_input_padding = default_net().plugin_config.remove_input_padding + use_gpt_attention_plugin = default_net( + ).plugin_config.gpt_attention_plugin + use_gemm_plugin = default_net().plugin_config.gemm_plugin + paged_kv_cache = default_net().plugin_config.paged_kv_cache + tokens_per_block = default_net().plugin_config.tokens_per_block + use_custom_all_reduce = default_net( + ).plugin_config.use_custom_all_reduce + + model_inputs = self.prepare_basic_inputs( + max_batch_size, + max_beam_width, + max_input_len, + max_new_tokens, + self.num_kv_heads, + head_size, + self.num_layers, + self.kv_dtype, + remove_input_padding=remove_input_padding, + use_gpt_attention_plugin=use_gpt_attention_plugin, + use_gemm_plugin=use_gemm_plugin, + use_custom_all_reduce=use_custom_all_reduce, + paged_kv_cache=paged_kv_cache, + tokens_per_block=tokens_per_block, + dtype=self.dtype, + num_heads=self.num_heads, + mapping=self.mapping, + max_num_tokens=max_num_tokens, + ) + + return (model_inputs['input_ids'], model_inputs['position_ids'], True, + model_inputs['last_token_ids'], + KeyValueCacheParams( + past_key_value=model_inputs['past_key_value'], + host_past_key_value_lengths=model_inputs[ + 'host_past_key_value_lengths'], + host_max_kv_cache_lengths=model_inputs[ + 'host_max_kv_cache_lengths'], + kv_cache_block_pointers=model_inputs[ + 'kv_cache_block_pointers_list'], + cache_indirection=model_inputs['cache_indirection'], + ), + AttentionParams( + sequence_length=model_inputs['sequence_length'], + context_lengths=model_inputs['context_lengths'], + host_context_lengths=model_inputs['host_context_lengths'], + max_context_length=max_input_len, + host_request_types=model_inputs['host_request_types']), + model_inputs['hidden_states_input'], + model_inputs['all_reduce_workspace']) diff --git a/tensorrt_llm/module.py b/tensorrt_llm/module.py index 662161afa7b6..810b464a547c 100644 --- a/tensorrt_llm/module.py +++ b/tensorrt_llm/module.py @@ -52,13 +52,33 @@ def __getattr__(self, name): type(self).__name__, name)) def __setattr__(self, name, value) -> None: - if isinstance(value, Parameter): - parameters = self.__dict__.get('_parameters') - parameters[name] = value + # Improved module setattr to handle one edge case: + # attribute could be first set to None and later reset to Parameter / Module class + + try: + super().__getattribute__(name) + + except AttributeError: + # if base class doesn't have the attribute, no matter we init or reset: + # - keep Parameter and Module attrs in this Module class + # - leave all other attrs in base class + if isinstance(value, Parameter): + self.__dict__.get('_parameters')[name] = value + elif isinstance(value, Module): + self.__dict__.get('_modules')[name] = value + else: + super().__setattr__(name, value) + else: - modules = self.__dict__.get('_modules') - if isinstance(value, Module): - modules[name] = value + # if base class has the attribute, reset as follows: + # - when reset as Parameter or Module attr, remove from base & add to this Module class + # - other types reset and remain in base class + if isinstance(value, Parameter): + super().__delattr__(name) + self.__dict__.get('_parameters')[name] = value + elif isinstance(value, Module): + super().__delattr__(name) + self.__dict__.get('_modules')[name] = value else: super().__setattr__(name, value) diff --git a/tensorrt_llm/parameter.py b/tensorrt_llm/parameter.py index f78bfa57a054..8b79ffe3de9e 100644 --- a/tensorrt_llm/parameter.py +++ b/tensorrt_llm/parameter.py @@ -47,8 +47,9 @@ def __init__(self, v_range = 0.1 if dtype == trt.DataType.INT8: - value = torch.randint(int(-128 * v_range), - int(128 * v_range), (shape), + upper = math.ceil(128 * v_range) + value = torch.randint(-upper, + upper, (shape), dtype=trt_dtype_to_torch(dtype), device='cuda') # value ~ U[int(-128 * v_range), int(128 * v_range)] diff --git a/tensorrt_llm/plugin/plugin.py b/tensorrt_llm/plugin/plugin.py index 91a1e9eab10d..5244bb6a56a0 100644 --- a/tensorrt_llm/plugin/plugin.py +++ b/tensorrt_llm/plugin/plugin.py @@ -79,6 +79,7 @@ def init(self): self.paged_kv_cache = False self.tokens_per_block = 0 self.lookup_plugin = False + self.lora_plugin = False def enable_qk_half_accum(self): self.attention_qk_half_accumulation = True @@ -175,3 +176,7 @@ def set_quantize_tensor_plugin(self): def set_lookup_plugin(self, dtype='float16'): self.lookup_plugin = dtype return self + + def set_lora_plugin(self, dtype='float16'): + self.lora_plugin = dtype + return self diff --git a/tensorrt_llm/profiler.py b/tensorrt_llm/profiler.py index b9ff6176c50a..1c545c2b445c 100644 --- a/tensorrt_llm/profiler.py +++ b/tensorrt_llm/profiler.py @@ -170,6 +170,8 @@ def device_memory_info( self, device: Optional[Union[torch.device, int]] = None, ) -> int: + if device is None: + device = torch.cuda.current_device() index = device.index if isinstance(device, torch.device) else device if index not in self.device_handles: handle = pynvml.nvmlDeviceGetHandleByIndex(index) diff --git a/tensorrt_llm/quantization/layers.py b/tensorrt_llm/quantization/layers.py index 1a34790be8e7..daaa1c060407 100644 --- a/tensorrt_llm/quantization/layers.py +++ b/tensorrt_llm/quantization/layers.py @@ -998,6 +998,7 @@ def __init__(self, apply_query_key_layer_scaling=False, attention_mask_type=AttentionMaskType.padding, bias=True, + qkv_bias_only=False, dtype=None, position_embedding_type=PositionEmbeddingType.learned_absolute, tp_group=None, @@ -1064,7 +1065,7 @@ def __init__(self, hidden_size, hidden_size + 2 * self.num_kv_heads * tp_size * self.attention_head_size, - bias=bias, + bias=(bias or qkv_bias_only), dtype=dtype, tp_group=tp_group, tp_size=tp_size, @@ -1079,13 +1080,16 @@ def __init__(self, tp_size=tp_size, quant_mode=quant_mode) + self.use_lora = False + def forward(self, hidden_states: Tensor, attention_mask=None, use_cache=False, kv_cache_params=None, attention_params=None, - workspace=None): + workspace=None, + lora_params=None): # TODO add in-flight batching to SmoothQuant if default_net().plugin_config.smooth_quant_gemm_plugin: qkv = self.qkv(hidden_states) diff --git a/tensorrt_llm/runtime/__init__.py b/tensorrt_llm/runtime/__init__.py index d1075424d7d7..bbfa62f54b64 100644 --- a/tensorrt_llm/runtime/__init__.py +++ b/tensorrt_llm/runtime/__init__.py @@ -12,9 +12,12 @@ # 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 .generation import SamplingConfig # autoflake: skip from .generation import (ChatGLMGenerationSession, GenerationSession, - ModelConfig, SamplingConfig, to_word_list_format) + ModelConfig, to_word_list_format) from .kv_cache_manager import GenerationSequence, KVCacheManager +from .lora_manager import LoraManager # autoflake: skip +from .model_runner import ModelRunner from .session import Session, TensorInfo __all__ = [ @@ -22,9 +25,11 @@ 'GenerationSession', 'GenerationSequence', 'KVCacheManager', + 'LoraManager' 'SamplingConfig', 'Session', 'TensorInfo', 'ChatGLMGenerationSession', 'to_word_list_format', + 'ModelRunner', ] diff --git a/tensorrt_llm/runtime/generation.py b/tensorrt_llm/runtime/generation.py index 244ed7227cb3..f1c20a3ea9ae 100755 --- a/tensorrt_llm/runtime/generation.py +++ b/tensorrt_llm/runtime/generation.py @@ -30,6 +30,7 @@ from ..mapping import Mapping from ..quantization import QuantMode from .kv_cache_manager import GenerationSequence, KVCacheManager +from .lora_manager import LoraManager from .session import _scoped_stream @@ -246,6 +247,7 @@ class ModelConfig: model_name: str = "" paged_kv_cache: bool = False cross_attention: bool = False + head_size: int = None has_position_embedding: bool = True has_token_type_embedding: bool = False tokens_per_block: int = 64 @@ -254,6 +256,7 @@ class ModelConfig: gather_all_token_logits: bool = False dtype: str = "" use_custom_all_reduce: bool = False + lora_plugin: bool = False @dataclass @@ -261,7 +264,12 @@ class SamplingConfig: end_id: int pad_id: int + max_new_tokens: int = field(default=20) num_beams: int = field(default=1) + max_kv_cache_length: Optional[int] = field(default=None) + output_sequence_lengths: bool = field(default=False) + return_dict: bool = field(default=False) + temperature: Union[float, torch.Tensor] = field(default=1.0) top_k: Union[int, torch.Tensor] = field(default=1) top_p: Union[float, torch.Tensor] = field(default=0.0) @@ -279,6 +287,15 @@ class SamplingConfig: output_cum_log_probs: bool = field(init=False, default=False) output_log_probs: bool = field(init=False, default=False) + def update(self, **kwargs): + unused_kwargs = dict() + for key, value in kwargs.items(): + if hasattr(self, key): + setattr(self, key, value) + else: + unused_kwargs[key] = value + return unused_kwargs + class GenerationSession(object): @@ -310,7 +327,8 @@ def __init__(self, f'cuda:{self.runtime.runtime_rank % mapping.gpus_per_node}') torch.cuda.set_device(self.device) # dynamic_decoder currently use torch's current stream, so must let TRT enqueue use same stream here - if stream is None: + self.stream = stream + if self.stream is None: self.stream = torch.cuda.Stream(self.device) torch.cuda.set_stream(self.stream) self.debug_mode = debug_mode @@ -368,7 +386,8 @@ def __init__(self, if model_config.has_position_embedding and self.mapping.is_first_pp_rank( ): expected_tensor_names += ['position_ids'] - if model_config.has_token_type_embedding: + if model_config.has_token_type_embedding and self.mapping.is_first_pp_rank( + ): expected_tensor_names += ['token_type_ids'] expected_tensor_names += ['cache_indirection'] @@ -394,7 +413,8 @@ def __init__(self, 'host_past_key_value_lengths' ] expected_tensor_names += [ - f'host_max_kv_cache_length_{i}' for i in range(self.num_layers) + f'host_max_kv_cache_length_{i}' + for i in range(self.first_layer, self.last_layer) ] if model_config.remove_input_padding: expected_tensor_names.append('host_context_lengths') @@ -410,10 +430,12 @@ def __init__(self, if model_config.cross_attention: expected_tensor_names += [ - f'cross_present_key_value_{i}' for i in range(self.num_layers) + f'cross_present_key_value_{i}' + for i in range(self.first_layer, self.last_layer) ] expected_tensor_names += [ - f'cross_past_key_value_{i}' for i in range(self.num_layers) + f'cross_past_key_value_{i}' + for i in range(self.first_layer, self.last_layer) ] expected_tensor_names += [ 'encoder_output', 'encoder_input_lengths', @@ -423,6 +445,13 @@ def __init__(self, if self.mapping.tp_size > 1 and model_config.use_custom_all_reduce: expected_tensor_names += ['all_reduce_workspace'] + if model_config.lora_plugin: + expected_tensor_names += ['lora_ranks'] + expected_tensor_names += [ + f'lora_weights_pointers_{i}' + for i in range(self.first_layer, self.last_layer) + ] + found_tensor_names = [ self.runtime.engine.get_tensor_name(i) for i in range(self.runtime.engine.num_io_tensors) @@ -493,7 +522,7 @@ def num_heads_kv(self): @property def head_size(self): - return self.hidden_size // self.num_heads + return self.hidden_size // self.num_heads if self._model_config.head_size is None else self._model_config.head_size @property def quant_mode(self): @@ -541,6 +570,10 @@ def has_position_embedding(self): def has_token_type_embedding(self): return self._model_config.has_token_type_embedding + @property + def use_lora_plugin(self): + return self._model_config.lora_plugin + def __setup_decoder(self, input_ids: torch.Tensor, sampling_config: SamplingConfig, host_context_lengths: torch.Tensor): @@ -606,8 +639,8 @@ def __setup_decoder(self, input_ids: torch.Tensor, dtype=torch.float32) assert ( - scfg.presence_penalty == 0.0 or scfg.repetition_penalty == 0.0 - ), f"presence_penalty({scfg.presence_penalty}) and repetition_penalty({scfg.repetition_penalty}) cannot be larger than 0.0 at the same time." + scfg.presence_penalty == 0.0 or scfg.repetition_penalty == 1.0 + ), f"presence_penalty({scfg.presence_penalty}) and repetition_penalty({scfg.repetition_penalty}) cannot be non-default values at the same time." if isinstance(scfg.min_length, torch.Tensor): assert scfg.min_length.dtype == torch.int32, f"scfg.min_length.dtype ({scfg.min_length.dtype}) must be torch.int32" @@ -662,6 +695,7 @@ def __setup_decoder(self, input_ids: torch.Tensor, device=self.device) max_context_length = host_context_lengths.max() + # setup output ids buffer if input_ids.shape[0] != host_context_lengths.shape[0]: # dim 0 of input_ids is not batch size, which means remove_padding is enabled split_ids_list = list( @@ -680,7 +714,7 @@ def __setup_decoder(self, input_ids: torch.Tensor, tiled_input_ids = tiled_input_ids.reshape(batch_size, scfg.num_beams, max_context_length) - tiled_input_ids.permute(2, 0, 1) + tiled_input_ids.permute(2, 0, 1) # TODO: delete? self.output_ids = torch.cat( (tiled_input_ids, torch.full((batch_size, scfg.num_beams, @@ -787,7 +821,9 @@ def setup(self, max_new_tokens: int, beam_width: int = 1, max_kv_cache_length: Optional[int] = None, - encoder_max_input_length: Optional[int] = None): + encoder_max_input_length: Optional[int] = None, + lora_manager: LoraManager = None, + lora_uids: List[str] = None): # Store these params related to buffer size to check against # the input shape with the params given in decode() self.batch_size = batch_size @@ -798,7 +834,7 @@ def setup(self, self.encoder_max_input_length = encoder_max_input_length if max_kv_cache_length is None: self.max_kv_cache_length = self.max_seq_length - logger.info( + logger.debug( "The max_kv_cache_length is not set, we will use max_seq_length by default." ) self.host_max_kv_cache_lengths = [ @@ -829,7 +865,8 @@ def setup(self, self.max_seq_length) if max_kv_cache_length.shape[0] != self.num_layers: logger.error( - "max_kv_cache_length tensor's size is not equal to num_layers!" + "max_kv_cache_length tensor's size is not equal to num_layers! " + "Note that num_layers = num_total_layers // pipeline_parallelism_size." ) assert False self.host_max_kv_cache_lengths = [ @@ -840,6 +877,7 @@ def setup(self, ] else: assert False, "invalid max_kv_cache_length!" + self.lora_manager = lora_manager self.buffer = {} if self.mapping.is_last_pp_rank(): @@ -850,6 +888,7 @@ def setup(self, dtype=self._tensor_dtype('logits'), device=self.device) if self.cross_attention: + # use shape info to pass max length info in remove padding mode self.buffer['encoder_max_input_length'] = torch.empty( (encoder_max_input_length, ), dtype=self._tensor_dtype('encoder_max_input_length'), @@ -927,6 +966,44 @@ def setup(self, dtype=torch.int64, device="cpu") + if self.use_lora_plugin and self.lora_manager is not None: + assert lora_uids is not None + lora_weights_pointers_list = [ + torch.zeros(size=(batch_size, 2), + dtype=torch.int64).contiguous().cpu() + for _ in range(self.num_layers) + ] + self.buffer.update({ + 'lora_ranks': + torch.zeros(size=(batch_size, ), + dtype=torch.int32).contiguous().cpu() + }) + + for idx in range(self.num_layers): + layer_idx = idx + self.first_layer + self.buffer.update({ + f'lora_weights_pointers_{layer_idx}': + torch.zeros(size=(batch_size, 2), + dtype=torch.int64).contiguous().cpu() + }) + for batch_idx in range(batch_size): + lora_uid = lora_uids[batch_idx] + if lora_uid is not None: + self.buffer['lora_ranks'][ + batch_idx] = self.lora_manager.uid_to_low_ranks( + lora_uid) + + self.buffer[f'lora_weights_pointers_{layer_idx}'][ + batch_idx][ + 0] = self.lora_manager.lora_weights_pointers_list[ + layer_idx][lora_uid][0] + self.buffer[f'lora_weights_pointers_{layer_idx}'][ + batch_idx][ + 1] = self.lora_manager.lora_weights_pointers_list[ + layer_idx][lora_uid][1] + else: + self.buffer['lora_ranks'][batch_idx] = 0 + self.buffer_allocated = True def _get_context_shape_buffer(self, @@ -1087,6 +1164,7 @@ def _get_context_shape_buffer(self, }) if self.use_gpt_attention_plugin: + # context request host_request_types = torch.zeros_like(context_lengths, device='cpu').int() ctx_shape.update({ @@ -1100,15 +1178,15 @@ def _get_context_shape_buffer(self, }) ctx_buffer.update({ f'host_max_kv_cache_length_{idx}': - self.host_max_kv_cache_lengths[idx], + self.host_max_kv_cache_lengths[idx - self.first_layer], }) ctx_buffer.update({ 'sequence_length': self.sequence_length_buffer, 'host_past_key_value_lengths': torch.tensor( - [0, 1] * batch_size, dtype=torch.int32 - ), # field 0: past_key_value_length, field 1: is_context + [0] * batch_size, dtype=torch.int32 + ), # field 0: past_key_value_length, field 1: is_context (deprecated). changed to [0], otherwise affects batch padded input mode 'host_request_types': host_request_types.contiguous(), }) @@ -1124,6 +1202,16 @@ def _get_context_shape_buffer(self, ctx_shape['all_reduce_workspace'] = self.all_reduce_workspace.shape ctx_buffer['all_reduce_workspace'] = self.all_reduce_workspace + if self.use_lora_plugin: + ctx_shape['lora_ranks'] = self.buffer['lora_ranks'].shape + ctx_buffer['lora_ranks'] = self.buffer['lora_ranks'] + for idx in range(self.num_layers): + layer_idx = idx + self.first_layer + ctx_shape[f'lora_weights_pointers_{layer_idx}'] = self.buffer[ + f'lora_weights_pointers_{layer_idx}'].shape + ctx_buffer[f'lora_weights_pointers_{layer_idx}'] = self.buffer[ + f'lora_weights_pointers_{layer_idx}'] + return ctx_shape, ctx_buffer def _get_next_step_shape_buffer(self, @@ -1192,7 +1280,11 @@ def _get_next_step_shape_buffer(self, next_step_buffer['position_ids'] = position_ids.contiguous() if self.cross_attention: - next_step_shape['encoder_output'] = encoder_output.shape + # hack: disable (or minimize) cross qkv computation at generation phase + # TODO: enable [0,0,.] true zero tensor input; or use IfConditionalLayer + next_step_shape['encoder_output'] = [ + 1, 1, encoder_output.shape[-1] + ] # encoder_output.shape next_step_shape[ 'encoder_input_lengths'] = encoder_input_lengths.shape next_step_shape['encoder_max_input_length'] = self.buffer[ @@ -1276,11 +1368,19 @@ def _get_next_step_shape_buffer(self, f'cross_past_key_value_{idx}'] = cross_cache_shape if self.use_gpt_attention_plugin: + # generation requests host_request_types = torch.ones_like(context_lengths, device='cpu').int() + # previous [past_kv_length, is_context] has been deprecated. only past_kv_length should be given here + # Note we should use max_context_length here to align to max -- but isn't this done in attn plugin's max_element() already? + host_past_key_value_lengths = torch.tensor( + [max_context_length + step] * (batch_size * beam_width), + dtype=torch.int32, + device='cpu') next_step_shape.update({ 'sequence_length': (batch_size * beam_width, ), - 'host_past_key_value_lengths': (batch_size * beam_width, ), + 'host_past_key_value_lengths': + host_past_key_value_lengths.shape, 'host_request_types': host_request_types.shape }) @@ -1290,19 +1390,13 @@ def _get_next_step_shape_buffer(self, }) next_step_buffer.update({ f'host_max_kv_cache_length_{idx}': - self.host_max_kv_cache_lengths[idx], + self.host_max_kv_cache_lengths[idx - self.first_layer], }) next_step_buffer.update({ # Sequence lengths are not used in the context phase actually. - 'sequence_length': - self.sequence_length_buffer, - 'host_past_key_value_lengths': - torch.tensor( - [max_context_length + step, 0] * (batch_size * beam_width), - dtype=torch.int32 - ), # field 0: past_key_value_length, field 1: is_context - 'host_request_types': - host_request_types, + 'sequence_length': self.sequence_length_buffer, + 'host_past_key_value_lengths': host_past_key_value_lengths, + 'host_request_types': host_request_types, }) if self.remove_input_padding: next_step_buffer[ @@ -1321,6 +1415,18 @@ def _get_next_step_shape_buffer(self, 'all_reduce_workspace'] = self.all_reduce_workspace.shape next_step_buffer['all_reduce_workspace'] = self.all_reduce_workspace + if self.use_lora_plugin: + next_step_shape['lora_ranks'] = self.buffer['lora_ranks'].shape + next_step_buffer['lora_ranks'] = self.buffer['lora_ranks'] + for idx in range(self.num_layers): + layer_idx = idx + self.first_layer + next_step_shape[ + f'lora_weights_pointers_{layer_idx}'] = self.buffer[ + f'lora_weights_pointers_{layer_idx}'].shape + next_step_buffer[ + f'lora_weights_pointers_{layer_idx}'] = self.buffer[ + f'lora_weights_pointers_{layer_idx}'] + return next_step_shape, next_step_buffer def _prepare_context_inputs(self, batch_size, context_lengths, @@ -1555,18 +1661,24 @@ def handle_per_step( batch_size, self.vocab_size_padded) if step == 0 and beam_width > 1: - + # these tiled tensors are returned by handle_per_step(), so they can relay to the next generation calls if not self.use_gpt_attention_plugin: attention_mask = _tile_beam_width(attention_mask, beam_width) context_lengths = _tile_beam_width(context_lengths, beam_width) host_context_lengths = _tile_beam_width(host_context_lengths, beam_width) + if encoder_input_lengths is not None: + encoder_input_lengths = _tile_beam_width( + encoder_input_lengths, beam_width) + if tasks is not None: tasks = _tile_beam_width(tasks, beam_width) # Move tiling before logit computing of context if not self.paged_kv_cache: for key in self.buffer.keys(): + # Note: this tiles both self attn cache and cross attn cache! + # both names contain "present_key_value" if "present_key_value" in key: self.buffer[key] = _tile_beam_width( self.buffer[key], beam_width) @@ -1607,6 +1719,7 @@ def handle_per_step( tasks, prompt_vocab_size, encoder_output, encoder_input_lengths) self.runtime._set_shape(next_context, next_step_shape) self.runtime._set_buffer(next_context, next_step_buffer) + if self.debug_mode: self.debug_buffer = next_step_buffer if self.cuda_graph_mode: @@ -1657,6 +1770,7 @@ def handle_per_step( next_token_logits = logits.reshape( (batch_size, beam_width, -1)).to(self.decoder_logits_dtype) decode_step = step + max_context_length + should_stop = self.dynamic_decoder.forward( next_token_logits, decode_step, max_context_length, self.max_kv_cache_length, ite, batch_size, self.end_ids, @@ -1691,7 +1805,7 @@ def handle_per_step( # We set this to False for all sequences, since we use only length criterion to stop now self.kv_cache_manager.step([False] * batch_size) - return should_stop, next_step_buffer, tasks, context_lengths, host_context_lengths, attention_mask, context_logits + return should_stop, next_step_buffer, tasks, context_lengths, host_context_lengths, attention_mask, context_logits, encoder_input_lengths def decode_regular(self, batch_size: int, @@ -1735,7 +1849,7 @@ def get_outputs_dict(output_ids): return outputs for step in range(0, self.max_new_tokens): - should_stop, next_step_buffer, tasks, context_lengths, host_context_lengths, attention_mask, logits = self.handle_per_step( + should_stop, next_step_buffer, tasks, context_lengths, host_context_lengths, attention_mask, logits, encoder_input_lengths = self.handle_per_step( cache_indirections, step, batch_size, max_context_length, beam_width, input_ids, hidden_states, scfg, kv_cache_block_pointers, prompt_embedding_table, tasks, @@ -1744,11 +1858,13 @@ def get_outputs_dict(output_ids): sequence_lengths, next_step_buffer, stop_words_list, bad_words_list, no_repeat_ngram_size, encoder_output, encoder_input_lengths) - if step == 0: - context_logits = logits if self.gather_all_token_logits: - generation_logits.append( - next_step_buffer['logits'].clone().detach()) + if self.mapping.is_last_pp_rank(): + if step == 0: + context_logits = logits + else: + generation_logits.append( + next_step_buffer['logits'].clone().detach()) if should_stop is not None and should_stop.item(): final_output_ids = self.finalize_decoder( @@ -1759,6 +1875,12 @@ def get_outputs_dict(output_ids): return get_outputs_dict(final_output_ids) else: return final_output_ids + elif self.mapping.is_last_pp_rank( + ) and self.gather_all_token_logits: + outputs = {} + outputs['context_logits'] = context_logits + outputs['generation_logits'] = generation_logits + return outputs else: return None @@ -1770,6 +1892,11 @@ def get_outputs_dict(output_ids): return get_outputs_dict(final_output_ids) else: return final_output_ids + elif self.mapping.is_last_pp_rank() and self.gather_all_token_logits: + outputs = {} + outputs['context_logits'] = context_logits + outputs['generation_logits'] = generation_logits + return outputs else: return None @@ -1813,7 +1940,7 @@ def get_outputs_dict(output_ids): return outputs for step in range(0, self.max_new_tokens): - should_stop, next_step_buffer, tasks, context_lengths, host_context_lengths, attention_mask, logits = self.handle_per_step( + should_stop, next_step_buffer, tasks, context_lengths, host_context_lengths, attention_mask, logits, encoder_input_lengths = self.handle_per_step( cache_indirections, step, batch_size, max_context_length, beam_width, input_ids, hidden_states, scfg, kv_cache_block_pointers, prompt_embedding_table, tasks, diff --git a/tensorrt_llm/runtime/lora_manager.py b/tensorrt_llm/runtime/lora_manager.py new file mode 100644 index 000000000000..3af6e36e8ead --- /dev/null +++ b/tensorrt_llm/runtime/lora_manager.py @@ -0,0 +1,74 @@ +import json + +import numpy as np + +from .._utils import _str_to_np_dict, fromfile, numpy_to_torch + + +class LoraManager(object): + + def __init__(self, model_dir, model_config): + ''' + Load lora modules, could be move to client side + ''' + self._model_config = model_config + self._lora_uid_to_key = {} + self._lora_uid_to_low_ranks = {} + self._lora_weights = [] + self._lora_weights_pointers_list = [ + ] # shape: [layer, lora_module_numbers, 2] + + with open(model_dir / "lora_weights.json", 'r') as f: + config = json.load(f) + lora_config = config['lora_config'] + for key in lora_config['lora_kqv_adapter']: + self._lora_uid_to_key[lora_config['lora_kqv_adapter'][key] + ['key']] = key + + for layer_idx in range(model_config.num_layers): + self._lora_weights_pointers_list.append({}) + + for uid, key in self._lora_uid_to_key.items(): + low_rank = int(lora_config['lora_kqv_adapter'][key]['low_rank']) + self._lora_uid_to_low_ranks[lora_config['lora_kqv_adapter'][key] + ['key']] = low_rank + prefix = f"model.model.language_model.encoder.layers.{layer_idx}.self_attention.adapter_layer.lora_kqv_adapter.{key}" + t_in = numpy_to_torch( + np.ascontiguousarray( + fromfile(model_dir, f'{prefix}.linear_in.weight.bin', + [model_config.hidden_size, low_rank], + _str_to_np_dict['bfloat16']).transpose( + 1, 0))).cuda() + + t_out = numpy_to_torch( + np.ascontiguousarray( + fromfile(model_dir, f'{prefix}.linear_out.weight.bin', + [low_rank, model_config.hidden_size * 3], + _str_to_np_dict['bfloat16']).transpose( + 1, 0))).cuda() + + self._lora_weights_pointers_list[layer_idx].update({ + uid: [ + t_in.contiguous().data_ptr(), + t_out.contiguous().data_ptr() + ] + }) + + self._lora_weights.append(t_in) + self._lora_weights.append(t_out) + + def uid_to_key(self, uid: str): + assert isinstance(uid, str) + return self._lora_uid_to_key[uid] + + def uid_to_low_ranks(self, uid: str): + assert isinstance(uid, str) + return self._lora_uid_to_low_ranks[uid] + + @property + def lora_weights(self): + return self._lora_weights + + @property + def lora_weights_pointers_list(self): + return self._lora_weights_pointers_list diff --git a/tensorrt_llm/runtime/model_runner.py b/tensorrt_llm/runtime/model_runner.py new file mode 100644 index 000000000000..05ebd10c0031 --- /dev/null +++ b/tensorrt_llm/runtime/model_runner.py @@ -0,0 +1,335 @@ +# 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 copy +import json +from pathlib import Path +from typing import List, Optional, Tuple, Union + +import torch + +import tensorrt_llm +import tensorrt_llm.profiler as profiler +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization import QuantMode +from tensorrt_llm.runtime import GenerationSession, ModelConfig, SamplingConfig + + +def get_engine_name(model: str, dtype: str, tp_size: int, pp_size: int, + rank: int) -> str: + """ + Get the serialized engine file name. + + Args: + model (str): + Model name, e.g., bloom, gpt. + dtype (str): + Data type, e.g., float32, float16, bfloat16, + tp_size (int): + The size of tensor parallel. + pp_size (int): + The size of pipeline parallel. + rank (int): + The rank id. + + Returns: + str: The serialized engine file name. + """ + if pp_size == 1: + return '{}_{}_tp{}_rank{}.engine'.format(model, dtype, tp_size, rank) + return '{}_{}_tp{}_pp{}_rank{}.engine'.format(model, dtype, tp_size, + pp_size, rank) + + +def read_config(config_path: Path) -> Tuple[ModelConfig, dict]: + """ + Read the engine config file and create a ModelConfig instance, return the ModelConfig instance + and other config fields in a dict. + + Args: + config_path (Path): + The path of engine config file. + + Returns: + Tuple[ModelConfig, dict]: A ModelConfig instance and other config fields. + """ + with open(config_path, 'r') as f: + config = json.load(f) + + builder_config = config['builder_config'] + model_name = builder_config['name'] + dtype = builder_config['precision'] + tp_size = builder_config['tensor_parallel'] + pp_size = builder_config.get('pipeline_parallel', 1) + world_size = tp_size * pp_size + assert world_size == tensorrt_llm.mpi_world_size(), \ + f'Engine world size ({tp_size} * {pp_size}) != Runtime world size ({tensorrt_llm.mpi_world_size()})' + + num_heads = builder_config['num_heads'] + assert num_heads % tp_size == 0, \ + f"The number of heads ({num_heads}) is not a multiple of tp_size ({tp_size})" + num_kv_heads = builder_config.get('num_kv_heads', num_heads) + # TODO: multi_query_mode should be removed + multi_query_mode = builder_config.get('multi_query_mode', False) + if multi_query_mode: + logger.warning( + "`multi_query_mode` config is deprecated. Please rebuild the engine." + ) + num_kv_heads = 1 + num_heads = num_heads // tp_size + num_kv_heads = (num_kv_heads + tp_size - 1) // tp_size + + hidden_size = builder_config['hidden_size'] // tp_size + vocab_size = builder_config['vocab_size'] + num_layers = builder_config['num_layers'] + + cross_attention = builder_config.get('cross_attention', False) + has_position_embedding = builder_config.get('has_position_embedding', True) + has_token_type_embedding = builder_config.get('has_token_type_embedding', + False) + gather_all_token_logits = builder_config.get('gather_all_token_logits', + False) + max_prompt_embedding_table_size = builder_config.get( + 'max_prompt_embedding_table_size', 0) + quant_mode = QuantMode(builder_config.get('quant_mode', 0)) + + plugin_config = config['plugin_config'] + use_gpt_attention_plugin = bool(plugin_config['gpt_attention_plugin']) + remove_input_padding = plugin_config['remove_input_padding'] + paged_kv_cache = plugin_config['paged_kv_cache'] + tokens_per_block = plugin_config['tokens_per_block'] + use_custom_all_reduce = plugin_config.get('use_custom_all_reduce', False) + + model_config = ModelConfig( + vocab_size=vocab_size, + num_layers=num_layers, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + hidden_size=hidden_size, + gpt_attention_plugin=use_gpt_attention_plugin, + remove_input_padding=remove_input_padding, + model_name=model_name, + paged_kv_cache=paged_kv_cache, + cross_attention=cross_attention, + has_position_embedding=has_position_embedding, + has_token_type_embedding=has_token_type_embedding, + tokens_per_block=tokens_per_block, + max_prompt_embedding_table_size=max_prompt_embedding_table_size, + quant_mode=quant_mode, + gather_all_token_logits=gather_all_token_logits, + dtype=dtype, + use_custom_all_reduce=use_custom_all_reduce) + + other_config = { + 'world_size': world_size, + 'tp_size': tp_size, + 'pp_size': pp_size, + 'max_batch_size': builder_config['max_batch_size'], + 'max_input_len': builder_config['max_input_len'] + } + return model_config, other_config + + +class ModelRunner: + """ + An interface class that wraps GenerationSession and provides generation methods. + """ + + def __init__(self, session: GenerationSession, max_batch_size: int, + max_input_len: int) -> None: + """ + Create a ModelRunner instance. + You are recommended to use the from_dir method to load the engine and create a ModelRunner instance. + + Args: + session (GenerationSession): + The TensorRT session created from an engine. + max_batch_size (int): + The maximum batch size allowed for the input. + max_input_len (int): + The maximum input length allowed for the input. + """ + self.session = session + self.max_batch_size = max_batch_size + self.max_input_len = max_input_len + + @classmethod + def from_dir(cls, + engine_dir: str, + rank: int = 0, + debug_mode: bool = False) -> 'ModelRunner': + """ + Create a ModelRunner instance from an engine directory. + + Args: + engine_dir (str): + The directory that contains the serialized engine files and config files. + rank (int): + The runtime rank id. + debug_mode (int): + Whether or not to turn on the debug mode. + Returns: + ModelRunner: An instance of ModelRunner. + """ + # session setup + engine_dir = Path(engine_dir) + config_path = engine_dir / "config.json" + model_config, other_config = read_config(config_path) + world_size = other_config.pop('world_size') + tp_size = other_config.pop('tp_size') + pp_size = other_config.pop('pp_size') + runtime_mapping = tensorrt_llm.Mapping(world_size=world_size, + rank=rank, + tp_size=tp_size, + pp_size=pp_size) + torch.cuda.set_device(rank % runtime_mapping.gpus_per_node) + + engine_name = get_engine_name(model_config.model_name, + model_config.dtype, tp_size, pp_size, + rank) + serialize_path = engine_dir / engine_name + + profiler.start('load tensorrt_llm engine') + with open(serialize_path, 'rb') as f: + engine_buffer = f.read() + + if model_config.model_name in ('chatglm_6b', 'glm_10b'): + session_cls = tensorrt_llm.runtime.ChatGLMGenerationSession + else: + session_cls = tensorrt_llm.runtime.GenerationSession + session = session_cls(model_config, + engine_buffer, + runtime_mapping, + debug_mode=debug_mode) + profiler.stop('load tensorrt_llm engine') + loading_time = profiler.elapsed_time_in_sec("load tensorrt_llm engine") + logger.info(f'Load engine takes: {loading_time} sec') + + return cls(session, **other_config) + + @property + def remove_input_padding(self) -> bool: + return self.session.remove_input_padding + + def _prepare_inputs(self, batch_input_ids: List[torch.Tensor], + pad_id: int) -> Tuple[torch.Tensor]: + # Remove potential additional dim, cast to int32 + batch_input_ids = [ + x.flatten().type(torch.int32) for x in batch_input_ids + ] + input_lengths = [x.size(0) for x in batch_input_ids] + max_length = max(input_lengths) + if max_length > self.max_input_len: + raise RuntimeError( + f"Maximum input length ({max_length}) exceeds the engine limit ({self.max_input_len})" + ) + batch_size = len(batch_input_ids) + if batch_size > self.max_batch_size: + raise RuntimeError( + f"Input batch size ({batch_size}) exceeds the engine limit ({self.max_batch_size})" + ) + + if self.remove_input_padding: + batch_input_ids = torch.concat(batch_input_ids).unsqueeze(0) + else: + # Right padding for trt-llm + paddings = [ + torch.ones(max_length - l, dtype=torch.int32) * pad_id + for l in input_lengths + ] + batch_input_ids = [ + torch.cat([x, pad]) for x, pad in zip(batch_input_ids, paddings) + ] + batch_input_ids = torch.stack(batch_input_ids) + input_lengths = torch.tensor(input_lengths, dtype=torch.int32) + return batch_input_ids, input_lengths + + def _prepare_outputs(self, outputs: dict, + input_lengths: torch.Tensor) -> dict: + if 'context_logits' in outputs: + batch_size = input_lengths.size(0) + context_logits = outputs['context_logits'] + if self.remove_input_padding: + context_logits = context_logits.flatten(end_dim=1) + + seg_points = [0] + input_lengths.cumsum(dim=0).tolist() + context_logits = [ + context_logits[s:e] + for s, e in zip(seg_points[:-1], seg_points[1:]) + ] + else: + context_logits = [ + context_logits[bidx, :input_lengths[bidx]] + for bidx in range(batch_size) + ] + outputs['context_logits'] = context_logits + + return outputs + + def generate(self, + batch_input_ids: List[torch.Tensor], + sampling_config: Optional[SamplingConfig] = None, + **kwargs) -> Union[torch.Tensor, dict]: + """ + Generates sequences of token ids. + The generation-controlling parameters are set in the sampling_config; it will be set to a default one if not passed. + You can override any sampling_config's attributes by passing corresponding parameters. + + Args: + batch_input_ids (List[torch.Tensor]): + A list of input id tensors. Each tensor is of shape (sequence_length, ). + sampling_config (Optional[SamplingConfig]): + The sampling configuration to be used as base parametrization for the generation call. + The passed **kwargs matching the sampling_config's attributes will override them. + If the sampling_config is not provided, a default will be used. + kwargs (Dict[str, Any]: + Ad hoc parametrization of sampling_config. + The passed **kwargs matching the sampling_config's attributes will override them. + Returns: + torch.Tensor or dict: + If return_dict=False, the method returns generated output_ids. + If return_dict=True, the method returns a dict of output_ids, + sequence_lengths (if sampling_config.output_sequence_lengths=True), + context_logits and generation_logits (if self.session.gather_all_token_logits=True). + """ + # Use sampling_config like HF's generation_config + if sampling_config is None: + sampling_config = SamplingConfig(end_id=None, pad_id=None) + else: + sampling_config = copy.deepcopy(sampling_config) + sampling_config.update(**kwargs) + + batch_size = len(batch_input_ids) + batch_input_ids, input_lengths = self._prepare_inputs( + batch_input_ids, sampling_config.pad_id) + + self.session.setup( + batch_size=batch_size, + max_context_length=input_lengths.max().item(), + max_new_tokens=sampling_config.max_new_tokens, + beam_width=sampling_config.num_beams, + max_kv_cache_length=sampling_config.max_kv_cache_length) + + batch_input_ids = batch_input_ids.cuda() + input_lengths = input_lengths.cuda() + outputs = self.session.decode( + batch_input_ids, + input_lengths, + sampling_config, + output_sequence_lengths=sampling_config.output_sequence_lengths, + return_dict=sampling_config.return_dict) + if sampling_config.return_dict: + outputs = self._prepare_outputs(outputs, input_lengths) + return outputs diff --git a/tensorrt_llm/runtime/session.py b/tensorrt_llm/runtime/session.py index 3f9b2de53d55..41d9f5d898d5 100644 --- a/tensorrt_llm/runtime/session.py +++ b/tensorrt_llm/runtime/session.py @@ -19,6 +19,7 @@ from typing import Any, Dict, List, Optional import tensorrt as trt +import torch from .._utils import trt_dtype_to_torch from ..logger import logger @@ -144,6 +145,25 @@ def _print_io_info(self): logger.info( f"Tensor:{name=:}, {mode=:}, {shape=:}, {dtype=:}, {tformat=:}") + def set_shapes(self, + tensor_dict: Dict[str, torch.Tensor], + context: Optional[trt.IExecutionContext] = None): + if context is None: + context = self.context + + for i in range(self.engine.num_io_tensors): + name = self.engine.get_tensor_name(i) + if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: + ok = context.set_input_shape(name, tensor_dict[name].shape) + logger.debug( + f"setting input tensor {name} with shape {tensor_dict[name].shape}" + ) + if not ok: + raise ValueError( + f"Couldn't assign {name} with shape {tensor_dict[name].shape}, " + f"engine supports [min, opt, max] = {self.engine.get_profile_shape(context.active_optimization_profile, name)}" + ) + def infer_shapes( self, inputs: List[TensorInfo], diff --git a/tests/bindings/test_bindings.py b/tests/bindings/test_bindings.py index 1d0ffcaba2d1..8ee7d25deb7d 100644 --- a/tests/bindings/test_bindings.py +++ b/tests/bindings/test_bindings.py @@ -1,3 +1,4 @@ +import inspect import json import tempfile from pathlib import Path @@ -208,6 +209,10 @@ def test_gpt_model_config(): gpt_model_config.compute_context_logits = True assert gpt_model_config.compute_context_logits + assert not gpt_model_config.compute_generation_logits + gpt_model_config.compute_generation_logits = True + assert gpt_model_config.compute_generation_logits + assert gpt_model_config.model_variant == _tb.GptModelVariant.GPT model_variant = _tb.GptModelVariant.GLM gpt_model_config.model_variant = model_variant @@ -341,3 +346,11 @@ def check_properties(the_object, properties, model_config): world_config) == json_config["name"] + "_float32_tp1_rank3.engine" assert gpt_json_config.engine_filename( world_config, "llama") == "llama_float32_tp1_rank3.engine" + + +def test_gpt_session(): + members = {name: tpe for (name, tpe) in inspect.getmembers(_tb.GptSession)} + assert isinstance(members["model_config"], property) + assert isinstance(members["world_config"], property) + assert isinstance(members["device"], property) + assert "generate" in members diff --git a/tests/bindings/test_gpt_session.py b/tests/bindings/test_gpt_session.py new file mode 100644 index 000000000000..8a0762f5bfd8 --- /dev/null +++ b/tests/bindings/test_gpt_session.py @@ -0,0 +1,216 @@ +import logging as _log +import os as _os +import pathlib as _pl +import subprocess as _sp +import sys as _sys +import typing as _tp + +import numpy as _np +import pytest +import torch as _tor + +import tensorrt_llm.bindings as _tb + + +@pytest.fixture(scope="module") +def llm_root() -> _pl.Path: + environ_root = _os.environ.get("LLM_ROOT", None) + return _pl.Path(environ_root) if environ_root is not None else _pl.Path( + __file__).parent.parent.parent + + +@pytest.fixture(scope="module") +def llm_model_root() -> _pl.Path | None: + return _os.environ.get("LLM_MODEL_ROOT", None) + + +@pytest.fixture(scope="module") +def resource_path(llm_root: _pl.Path) -> _pl.Path: + return llm_root / "cpp" / "tests" / "resources" + + +@pytest.fixture(scope="module") +def engine_path(resource_path: _pl.Path) -> _pl.Path: + return resource_path / "models" / "rt_engine" + + +@pytest.fixture(scope="module") +def data_path(resource_path: _pl.Path) -> _pl.Path: + return resource_path / "data" + + +def run_command(command: _tp.Sequence[str], + cwd: _pl.Path, + *, + shell=False, + env=None) -> None: + _log.info("Running: cd %s && %s", str(cwd), " ".join(command)) + _sp.check_call(command, cwd=cwd, shell=shell, env=env) + + +def prepare_model_tests( + llm_root: _pl.Path, + resource_path: _pl.Path, + model_name: str, + model_cache_arg=[], +): + scripts_dir = resource_path / "scripts" + python_exe = _sys.executable + model_env = {**_os.environ, "PYTHONPATH": f"examples/{model_name}"} + build_engines = [ + python_exe, + str(scripts_dir / f"build_{model_name}_engines.py") + ] + model_cache_arg + run_command(build_engines, cwd=llm_root, env=model_env) + + generate_expected_output = [ + python_exe, + str(scripts_dir / f"generate_expected_{model_name}_output.py") + ] + run_command(generate_expected_output, cwd=llm_root, env=model_env) + + +def sequence_lengths(sequences: _np.ndarray, pad_id: int) -> _np.ndarray: + return _np.apply_along_axis(lambda x: _np.searchsorted(x, True), 1, + sequences == pad_id).astype("int32") + + +@pytest.mark.parametrize( + "variant, results_file", + [ + ("fp32-default", "output_tokens_fp32_tp1_pp1.npy"), + ("fp32-plugin", "output_tokens_fp32_plugin_tp1_pp1.npy"), + ("fp16-default", "output_tokens_fp16_tp1_pp1.npy"), + ("fp16-plugin", "output_tokens_fp16_plugin_tp1_pp1.npy"), + # ("fp16-plugin-packed", "output_tokens_fp16_plugin_packed_tp1_pp1.npy"), + # ("fp16-plugin-packed-paged", "output_tokens_fp16_plugin_packed_paged_tp1_pp1.npy"), + ]) +def test_gpt_session(variant, results_file, llm_root: _pl.Path, + resource_path: _pl.Path, engine_path: _pl.Path, + data_path: _pl.Path, llm_model_root): + model_dir = "gpt2" + tp_size = 1 + pp_size = 1 + beam_width = 1 + max_batch_size = 8 + end_id = 50256 + pad_id = 50256 + repetitions = 2 + + # load input data + input_path = data_path / "input_tokens.npy" + assert input_path.is_file() + given_input = _np.load(input_path).astype("int32") + input_shape = given_input.shape + assert len(input_shape) == 2 + num_given_inputs = input_shape[0] + assert max_batch_size <= num_given_inputs + max_input_length = input_shape[1] + given_input_lengths = sequence_lengths(given_input, pad_id) + assert _np.all(given_input_lengths <= max_input_length) + + # load expected output data + results_path = data_path / model_dir / ( + "sampling" + if beam_width == 1 else f"beam_search_{beam_width}") / results_file + + if not results_path.exists(): + model_cache_arg = ["--model_cache", + str(llm_model_root) + ] if llm_model_root is not None else [] + prepare_model_tests(llm_root, resource_path, "gpt", model_cache_arg) + + assert results_path.is_file() + expected_output = _np.load(results_path) + output_shape = expected_output.shape + assert len(output_shape) == 2 + assert num_given_inputs * beam_width == output_shape[0] + max_seq_length = output_shape[1] + assert max_input_length <= max_seq_length + expected_output_lengths = sequence_lengths(expected_output, end_id) + assert _np.all(expected_output_lengths <= max_seq_length) + + gpu_size_path = f"tp{tp_size}-pp{pp_size}-gpu" + model_path = engine_path / model_dir / variant / gpu_size_path + assert model_path.is_dir() + config_path = model_path / "config.json" + config_json = _tb.GptJsonConfig.parse_file(str(config_path)) + assert config_json.tensor_parallelism == tp_size + assert config_json.pipeline_parallelism == pp_size + world_config = _tb.WorldConfig.mpi(tensor_parallelism=tp_size, + pipeline_parallelism=pp_size) + engine_filename = config_json.engine_filename(world_config) + assert (model_path / engine_filename).is_file() + session_config = _tb.GptSessionConfig(max_batch_size, beam_width, + max_seq_length) + + model_config = config_json.model_config + session = _tb.GptSession(session_config, model_config, world_config, + str(model_path / engine_filename)) + assert isinstance(session, _tb.GptSession) + assert isinstance(session.model_config, _tb.GptModelConfig) + assert isinstance(session.world_config, _tb.WorldConfig) + assert session.device == world_config.device + cuda_device = _tor.device("cuda", world_config.device) + + max_new_tokens = max_seq_length - max_input_length + sampling_config = _tb.SamplingConfig(beam_width) + sampling_config.temperature = [1.0] + sampling_config.min_length = [1] + sampling_config.random_seed = [42] + sampling_config.top_k = [0] + sampling_config.top_p = [0.0] + + packed_input = model_config.use_packed_input + assert not packed_input + input_ids = _tor.from_numpy( + given_input[:max_batch_size, :max_input_length]).to(cuda_device) + assert input_ids.dtype == _tor.int32 + input_lengths = _tor.from_numpy( + given_input_lengths[:max_batch_size]).to(cuda_device) + assert input_lengths.dtype == _tor.int32 + generation_input = _tb.GenerationInput(end_id, pad_id, input_ids, + input_lengths, packed_input) + generation_input.max_new_tokens = max_new_tokens + + for r in range(repetitions): + output_ids = _tor.empty((max_batch_size, max_seq_length), + dtype=_tor.int32, + device=cuda_device) + output_lengths = _tor.empty((max_batch_size, ), + dtype=_tor.int32, + device=cuda_device) + generation_output = _tb.GenerationOutput(output_ids, output_lengths) + num_steps = 0 + + def on_token_generated(ids, step, finished): + assert ids.shape == (max_batch_size, 1, max_seq_length) + nonlocal num_steps + assert step == num_steps + num_steps += 1 + # check that we only finish after producing `maxNewTokens` tokens + assert not finished or num_steps == max_new_tokens + # check that `finished` is set to true after producing `maxNewTokens` tokens + assert num_steps != max_new_tokens or finished + + generation_output.on_token_generated = on_token_generated + + session.generate(generation_output, generation_input, sampling_config) + observed_output = output_ids.squeeze().cpu().numpy() + assert observed_output.shape == (max_batch_size, max_seq_length) + observed_output_lengths = output_lengths.squeeze().cpu().numpy() + assert _np.all(observed_output_lengths <= max_seq_length) + + for batch_idx in range(max_batch_size): + expected_length = expected_output_lengths[batch_idx] + observed_length = observed_output_lengths[batch_idx] + assert expected_length == observed_length, (batch_idx, + expected_length, + observed_length) + expected = expected_output[batch_idx, :expected_length] + observed = observed_output[batch_idx, :expected_length] + unmatched = expected != observed + if _np.any(unmatched): + assert False, (batch_idx, _np.where(unmatched), + _np.column_stack( + (expected, observed))[unmatched]) diff --git a/tests/model/test_gpt.py b/tests/model/test_gpt.py index cdefaef936c8..a553845d22b0 100644 --- a/tests/model/test_gpt.py +++ b/tests/model/test_gpt.py @@ -30,6 +30,8 @@ import tensorrt_llm from tensorrt_llm import Builder from tensorrt_llm._utils import str_dtype_to_torch +from tensorrt_llm.functional import RotaryScalingType +from tensorrt_llm.layers import PositionEmbeddingType from tensorrt_llm.network import net_guard from tensorrt_llm.plugin.plugin import ContextFMHAType from tensorrt_llm.runtime import ModelConfig, SamplingConfig @@ -899,6 +901,36 @@ def test_greedy_search_float32(self, use_refit, streaming): np.testing.assert_allclose(ref.cpu().numpy(), res.cpu().numpy()) + def test_rope_scaling_is_set_in_attention(self): + num_layers = 2 + position_embedding_type = PositionEmbeddingType.rope_gpt_neox + rotary_embedding_percentage = 0.3 + rotary_base = 99999.1 + rotary_scaling = {"type": "linear", "factor": 2.72} + tensorrt_llm_gpt = tensorrt_llm.models.GPTLMHeadModel( + num_layers=num_layers, + num_heads=4, + hidden_size=128, + vocab_size=256, + hidden_act='gelu', + max_position_embeddings=1024, + dtype=trt.float16, + position_embedding_type=position_embedding_type, + rotary_embedding_percentage=rotary_embedding_percentage, + rotary_base=rotary_base, + rotary_scaling=rotary_scaling, + ) + for layer_i in range(num_layers): + assert tensorrt_llm_gpt.layers[ + layer_i].attention.rotary_embedding_base == rotary_base + assert tensorrt_llm_gpt.layers[ + layer_i].attention.rotary_embedding_scale == rotary_scaling[ + "factor"] + assert tensorrt_llm_gpt.layers[ + layer_i].attention.rotary_embedding_scale_type == RotaryScalingType.linear + assert tensorrt_llm_gpt.layers[ + layer_i].attention.position_embedding_type == position_embedding_type + if __name__ == '__main__': unittest.main() diff --git a/tests/model/test_gpt_e2e.py b/tests/model/test_gpt_e2e.py index 2e2bda4236c1..2b29cfc6dd81 100644 --- a/tests/model/test_gpt_e2e.py +++ b/tests/model/test_gpt_e2e.py @@ -143,8 +143,7 @@ def check_accuracy(engine_dir, input_tokens, max_output_len): hidden_size = config['builder_config']['hidden_size'] // world_size vocab_size = config['builder_config']['vocab_size'] num_layers = config['builder_config']['num_layers'] - multi_query_mode = config['builder_config']['multi_query_mode'] - num_kv_heads = 1 if multi_query_mode else num_heads + num_kv_heads = config['builder_config']['num_kv_heads'] runtime_rank = tensorrt_llm.mpi_rank() runtime_mapping = tensorrt_llm.Mapping(world_size, diff --git a/tests/tools/plugin_gen/test_plugin_gen.py b/tests/tools/plugin_gen/test_plugin_gen.py index 7880d2ffcf67..54188bc1f13d 100644 --- a/tests/tools/plugin_gen/test_plugin_gen.py +++ b/tests/tools/plugin_gen/test_plugin_gen.py @@ -1,4 +1,5 @@ import os +from importlib.metadata import version import pytest @@ -20,6 +21,10 @@ def gen_trt_plugins(*args, **kwargs): def is_triton_installed() -> bool: + # the triton detection does not work in PyTorch NGC 23.10 container + if version('triton') != "2.1.0+440fd1b": + return False + return os.path.exists(TRITON_COMPILE_BIN) From b2da17f7648286d42d01c4e513564d2c59c894a9 Mon Sep 17 00:00:00 2001 From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Date: Fri, 17 Nov 2023 06:24:50 -0800 Subject: [PATCH 2/3] Update submodule --- 3rdparty/cutlass | 2 +- 3rdparty/json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/3rdparty/cutlass b/3rdparty/cutlass index fc9ebc645b63..39c6a83f231d 160000 --- a/3rdparty/cutlass +++ b/3rdparty/cutlass @@ -1 +1 @@ -Subproject commit fc9ebc645b63f3a6bc80aaefde5c063fb72110d6 +Subproject commit 39c6a83f231d6db2bc6b9c251e7add77d68cbfb4 diff --git a/3rdparty/json b/3rdparty/json index 5fec8034933e..bc889afb4c5b 160000 --- a/3rdparty/json +++ b/3rdparty/json @@ -1 +1 @@ -Subproject commit 5fec8034933ef434a98dfbd2551b052c56345869 +Subproject commit bc889afb4c5bf1c0d8ee29ef35eaaf4c8bef8a5d From 2cd379ed3b649dfdadbc54e9dd1de524bc207e4e Mon Sep 17 00:00:00 2001 From: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com> Date: Fri, 17 Nov 2023 14:44:39 +0000 Subject: [PATCH 3/3] update --- .../aarch64-linux-gnu/libtensorrt_llm_batch_manager_static.a | 3 +++ .../libtensorrt_llm_batch_manager_static.pre_cxx11.a | 3 +++ cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/version.txt | 3 +++ 3 files changed, 9 insertions(+) create mode 100644 cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/libtensorrt_llm_batch_manager_static.a create mode 100644 cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/libtensorrt_llm_batch_manager_static.pre_cxx11.a create mode 100644 cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/version.txt diff --git a/cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/libtensorrt_llm_batch_manager_static.a b/cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/libtensorrt_llm_batch_manager_static.a new file mode 100644 index 000000000000..223613c73e6a --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/libtensorrt_llm_batch_manager_static.a @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b867c2e048671eecc421244d436436782093baf02f0fd5d49232b3d3042e55ea +size 1688216 diff --git a/cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/libtensorrt_llm_batch_manager_static.pre_cxx11.a b/cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/libtensorrt_llm_batch_manager_static.pre_cxx11.a new file mode 100644 index 000000000000..bac8af4a8964 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/libtensorrt_llm_batch_manager_static.pre_cxx11.a @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:db433a13ec6a017638bbb97b53a98624ad675b395787c99054d48ab370f5e3a0 +size 1697778 diff --git a/cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/version.txt b/cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/version.txt new file mode 100644 index 000000000000..c357c2d84af3 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/aarch64-linux-gnu/version.txt @@ -0,0 +1,3 @@ +81f472ac2b68edd03a0265299744347f libtensorrt_llm_batch_manager_static.a +4e5e3bbdfffa6deb6a50c541a946ac7a libtensorrt_llm_batch_manager_static.pre_cxx11.a +7edd8a21 commit