diff --git a/docs/source/blogs/Best_perf_practice_on_DeepSeek-R1_in_TensorRT-LLM.md b/docs/source/blogs/Best_perf_practice_on_DeepSeek-R1_in_TensorRT-LLM.md index 22d4688c503d..18c5ce06c0f2 100644 --- a/docs/source/blogs/Best_perf_practice_on_DeepSeek-R1_in_TensorRT-LLM.md +++ b/docs/source/blogs/Best_perf_practice_on_DeepSeek-R1_in_TensorRT-LLM.md @@ -414,13 +414,13 @@ Note that, `max_batch_size` and `max_num_tokens` can easily affect the performan Generally, you should make sure that `max_batch_size` is not too low to bottleneck the throughput, and `max_num_tokens` needs to be large enough so that it covers the max input sequence length of the samples in dataset, as mentioned in below section "WIP: Chunked context support on DeepSeek models". -For more details on `max_batch_size` and `max_num_tokens`, refer to [Tuning Max Batch Size and Max Num Tokens](../legacy/performance/performance-tuning-guide/tuning-max-batch-size-and-max-num-tokens.md). +For more details on `max_batch_size` and `max_num_tokens`, refer to the performance tuning guidance in the developer documentation. ### MLA chunked context MLA currently supports the chunked context feature on both Hopper and Blackwell GPUs. You can use `--enable_chunked_context` to enable it. This feature is primarily designed to reduce TPOT (Time Per Output Token). The default chunk size is set to `max_num_tokens`. If you want to achieve a lower TPOT, you can appropriately reduce the chunk size. However, please note that this will also decrease overall throughput. Therefore, a trade-off needs to be considered. -For more details on `max_num_tokens`, refer to [Tuning Max Batch Size and Max Num Tokens](../legacy/performance/performance-tuning-guide/tuning-max-batch-size-and-max-num-tokens.md). +For more details on `max_num_tokens`, refer to the performance tuning guidance in the developer documentation. ### Out of memory issues diff --git a/docs/source/blogs/XQA-kernel.md b/docs/source/blogs/XQA-kernel.md index d4018966cd2d..456024a10ed5 100644 --- a/docs/source/blogs/XQA-kernel.md +++ b/docs/source/blogs/XQA-kernel.md @@ -2,7 +2,7 @@ XQA kernel provides optimization for [MQA](https://arxiv.org/abs/1911.02150) and [GQA](https://arxiv.org/abs/2305.13245v3) during generation phase. It also provides optimization for beam search. Using tensor cores for acceleration, reducing data loading and conversion, it delivers increased throughput within the same latency budget. Increased throughput allows serving greater number of user requests while providing the same experience. -Support matrix and usage flags are described in [docs/source/legacy/advanced/gpt-attention](../legacy/advanced/gpt-attention.md#xqa-optimization). +Support matrix and usage flags are described in the [Attention](../features/attention.md) documentation. **Increased Throughput:** Looking at the Throughput-Latency curves below, we see that the enabling of XQA optimization increases throughput. Higher throughput equates to serving more users, and we can see that TPOT on the Y-axis flattens out when XQA gets enabled. diff --git a/docs/source/commands/trtllm-build.rst b/docs/source/commands/trtllm-build.rst deleted file mode 100644 index d209d3eac9ca..000000000000 --- a/docs/source/commands/trtllm-build.rst +++ /dev/null @@ -1,7 +0,0 @@ -trtllm-build -=========================== - -.. argparse:: - :module: tensorrt_llm.commands.build - :func: parse_arguments - :prog: trtllm-build diff --git a/docs/source/legacy/advanced/images/8x_l20_L40S_node_architecture.png b/docs/source/features/images/8x_l20_L40S_node_architecture.png similarity index 100% rename from docs/source/legacy/advanced/images/8x_l20_L40S_node_architecture.png rename to docs/source/features/images/8x_l20_L40S_node_architecture.png diff --git a/docs/source/legacy/advanced/lowprecision-pcie-allreduce.md b/docs/source/features/lowprecision-allreduce.md similarity index 69% rename from docs/source/legacy/advanced/lowprecision-pcie-allreduce.md rename to docs/source/features/lowprecision-allreduce.md index b7ab50703701..54eddafffb7a 100644 --- a/docs/source/legacy/advanced/lowprecision-pcie-allreduce.md +++ b/docs/source/features/lowprecision-allreduce.md @@ -5,8 +5,7 @@ Note: This feature is optimized for PCIe-based GPU topologies and may affect model accuracy. Please evaluate precision impact for your specific workload. ``` - -TRT-LLM supports `low-precision-allreduce`, a communication optimization that accelerates AllReduce operations in PCIe-based GPU environments. This feature quantizes FP16/BF16 data to FP8 during network transmission, reducing communication volume and improving performance. +TensorRT LLM supports `low-precision-allreduce`, a communication optimization that accelerates AllReduce operations in PCIe-based GPU environments. This feature quantizes FP16/BF16 data to FP8 during network transmission, reducing communication volume and improving performance. It is available on the PyTorch backend through the `allreduce_strategy` field of `LlmArgs`. ## Algorithm @@ -35,20 +34,16 @@ Low-Precision-AllReduce is specifically designed for the topology shown above, w ## Usage -The Low-Precision-AllReduce algorithm can be enabled in two ways: +Enable the Low-Precision-AllReduce algorithm by setting the `allreduce_strategy` field in `LlmArgs`: -1. **Direct specification** in your code: -``` -AllReduce allreduce(mapping=mapping, strategy=AllReduceStrategy.LOWPRECISION); -``` +```python +from tensorrt_llm import LLM -2. Enable by LlmArgs -``` -Set allreduce_strategy field in LlmArgs. -Candidates of strategies are "AUTO", "NCCL", "UB", "MINLATENCY", "ONESHOT", "TWOSHOT", "LOWPRECISION" and "MNNVL". -If no strategy is set, AUTO will be set. +llm = LLM(model="", allreduce_strategy="LOWPRECISION") ``` +The candidate strategies are `"AUTO"`, `"NCCL"`, `"UB"`, `"MINLATENCY"`, `"ONESHOT"`, `"TWOSHOT"`, `"LOWPRECISION"` and `"MNNVL"`. If no strategy is set, `AUTO` is used. + ## Performance and Accuracy Considerations Low-Precision-AllReduce reduces communication volume by using FP8 data format for transmission. This optimization: @@ -58,4 +53,4 @@ Low-Precision-AllReduce reduces communication volume by using FP8 data format fo Users should evaluate the precision impact on their specific models and workloads. -**Note**: When compiling TensorRT-LLM without enabling the `ENABLE_FP8` option, setting Low Precision allreduce will not take effect. +**Note**: When TensorRT LLM is built without the `ENABLE_FP8` option, setting the low-precision AllReduce strategy will not take effect. diff --git a/docs/source/index.rst b/docs/source/index.rst index 9d92a8770148..b65f8a62b3d8 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -72,6 +72,7 @@ Welcome to TensorRT LLM's Documentation! features/overlap-scheduler.md features/paged-attention-ifb-scheduler.md features/parallel-strategy.md + features/lowprecision-allreduce.md features/quantization.md features/sampling.md features/additional-outputs.md @@ -127,10 +128,9 @@ Welcome to TensorRT LLM's Documentation! .. toctree:: :maxdepth: 2 - :caption: Use TensorRT Engine - :hidden: + :caption: Migration - legacy/tensorrt_quickstart.md + legacy/tensorrt-backend-removal.md Indices and tables ================== diff --git a/docs/source/legacy/advanced/disaggregated-service.md b/docs/source/legacy/advanced/disaggregated-service.md deleted file mode 100644 index ac6cbad472e9..000000000000 --- a/docs/source/legacy/advanced/disaggregated-service.md +++ /dev/null @@ -1,86 +0,0 @@ -(disaggregated-service)= - -# Disaggregated-Service (Prototype) - -```{note} -Note: -This feature is currently in prototype, and the related API is subject to change in future versions. -``` -Currently TRT-LLM supports `disaggregated-service`, where the context and generation phases of a request can run on different executors. TRT-LLM's disaggregated service relies on the executor API, please make sure to read the [executor page](executor.md) before reading the document. - -For more information on disaggregated service in LLM inference, one can refer to papers such as [DistServe](https://arxiv.org/abs/2401.09670), [SplitWise](https://arxiv.org/abs/2311.18677). - -An [architectural and performance overview](../../../docs/source/blogs/tech_blog/blog05_Disaggregated_Serving_in_TensorRT-LLM.md), as well as [usage examples](../../../examples/disaggregated/README.md), are provided. - -## Environment Variables - -TRT-LLM uses some environment variables to control the behavior of disaggregated service. - - -* `TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP`: If set to `1`, generationExecutor will not overlap KV cache transfer with model inference. The default value is `0`. - -* `TRTLLM_ENABLE_KVCACHE_RECEIVE_PARALLEL`: When the generation rank receives KV cache from multiple context ranks within a single context instance, it will receive KV cache from each rank sequentially. If set to `1`, the generation rank will receive KV cache from each rank within one context instance in parallel. The default value is `0`. - -* `TRTLLM_REQUEST_KV_CACHE_CONCURRENT`: If set to `1`, generationExecutor prepares independent resources for each context executor to receive KV cache, requests whose KV cache are received from different context executors will be processed concurrently. If set to `0`, the generation executor will reuse the same resource to process KV cache transfer for each request sequentially, reducing the resources used by KV cache transmission and thereby lowering the risk of running out of memory. The default value is `0`. - -* `TRTLLM_TRY_ZCOPY_FOR_KVCACHE_TRANSFER`: TRT-LLM typically copies non-contiguous data into a temporary buffer before sending KV cache. If set to `1`, TRT-LLM will attempt to directly transmit each KV cache block, eliminating extra copies. The default value is `0`. - -* `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE`: By default, TRT-LLM uses a `stream-ordered memory allocator` to allocate temporary buffers. If this environment variable is set to #Size, TRT-LLM will use `cudaMalloc` to allocate buffer of size #Size for KV cache transmission. The default value is `512MB`. Users can set `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=1GB` to allocate a 1 GB buffer with `cudaMalloc` for KV cache transmission. - -* `TRTLLM_KVCACHE_TRANSFER_USE_ASYNC_BUFFER`: If set to `1`, TRT-LLM will use `cudaMallocAsync` to allocate buffers for KV cache transmission. The default value is `0`. This environment variable only takes effect when `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE` is greater than 0. - -* `TRTLLM_KVCACHE_SEND_MAX_CONCURRENCY_NUM`: The maximum number of concurrent KV cache sends. The default value is `1`. This environment variable only takes effect when `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE` is greater than 0. - -There are some other useful environment variables that may help when encountering failures or performance issues. - -* `NCCL_GRAPH_MIXING_SUPPORT`: With the default value `1`, the CUDA driver may create too many CUDA streams while working with one CUDA graph, leading to performance drop. Setting it to `0` will reduce the number of CUDA streams, but please make sure there are no other NCCL ops outside the one CUDA graph, otherwise it's unsafe. - -* `UCX_MAX_RNDV_RAILS`: With the default value `2`, UCX attempts to use two InfiniBand (IB) NIC devices per GPU for Rendezvous (RNDV) transfers. When both the context and generation instances enable tensor- and expert-parallel (TEP), multiple TP ranks may transfer KV cache concurrently. Because each TP rank can use up to two NIC devices, some NIC devices can be shared across GPUs, causing contention and reduced throughput. Setting `UCX_MAX_RNDV_RAILS=1` can reduce contention in this case. - -## Troubleshooting and FAQ - -### General FAQs - -*Q. What are the limitations of disaggregated-service in TRT-LLM?* - -A. Currently, only `decoder-only engine` and `beamWidth=1` are supported, and the KV cache at each layer of the model is required to be homogeneous, with the same data type and the same number of attention heads. - -*Q. Is the engine used by disaggregated-service different from other engines?* - -A. No. There are no special requirements for the arguments to build engine. - -*Q. Do the engines used by the context executor and generation executor need to be the same?* - -A. No. The engines used by context executor and generation executor can be different, and their parallelism can be heterogeneous, i.e., TP,PP can be different, and TRT-LLM will handle the heterogeneity of KV cache. - -*Q. Does TRT-LLM support running multiple context executor instances and generation executor instances?* - -A. Yes. TRT-LLM supports running multiple context executors and generation executors at the same time, and each executor can use different engine, but it is the user's responsibility to route requests to different executors and manage `requestId`. - -*Q. Can an executor handle both context-only requests and generation-only requests?* - -A. Yes, but it's not recommended, TRT-LLM does not implement proper scheduling for the case where the executor handles mixed context-only requests and generation-only requests, it's better to run context-only requests and generation-only requests on different executors. - -*Q. Does disaggregated-service in TRT-LLM support multi-gpu and multi-node?* - -A. Yes, it's recommended that different executor use different GPUs. We support context-only executor and generation-only executor run on same node or different nodes. The `participantIds` and `deviceIds` used by each executor need to be explicitly set by the user, and the `participantIds` of each executor must not be intersecting. - -### Debugging FAQs - -*Q. Does TRT-LLM support using GPU direct RDMA for inter-node KV Cache transfer?* - -A. Yes, TRT-LLM supports using GPU direct RDMA for inter-node KV cache transfer. - -*Q. What causes the substantial bandwidth fluctuations in kvCache transfers, especially during the first few requests following service initialization?* - -A. The communication for kvCache transfer between executors are established dynamically. The connection establishment process incurs significant overhead, which explains the apparently lower kvCache transfer bandwidth observed during the initial requests after service startup. This lower bandwidth reflects the inclusion of connection establishment overhead. When conducting benchmarks, it is recommended to perform a warm-up phase to ensure accurate performance measurements. - -*Q. When my servers are running on different NVLink domains, some servers hang or have a lower performance. How to fix that?* - -A. NVLink domain can be found with `nvidia-smi -q` in the `Fabric.ClusterUUID` field. A few UCX environment variables can be adjusted when your servers have different NVLink domains: - -* `UCX_CUDA_IPC_ENABLE_MNNVL`: Set to `n`. This also can reduce UCX timeout error messages like `UCX ERROR cuMemImportFromShareableHandle failed: invalid resource handle`, although these errors don't necessarily cause your trtllm-serve to fail. - -* `UCX_NET_DEVICES`: Check if this is set correctly, or unset this variable to allow UCX to use all possible devices. - -* `UCX_RNDV_SCHEME`: Set to `get_zcopy` or `put_zcopy` on GB200 for better performance. The default value is `auto`. diff --git a/docs/source/legacy/advanced/executor.md b/docs/source/legacy/advanced/executor.md deleted file mode 100644 index 1f59d0802a16..000000000000 --- a/docs/source/legacy/advanced/executor.md +++ /dev/null @@ -1,153 +0,0 @@ -(executor)= - -# Executor API - -TensorRT-LLM includes a high-level C++ API called the Executor API which allows you to execute requests -asynchronously, with in-flight batching, and without the need to define callbacks. - -A software component (referred to as "the client" in the text that follows) can interact -with the executor using the API defined in the [`executor.h`](source:cpp/include/tensorrt_llm/executor/executor.h) file. -For details about the API, refer to the {ref}`_cpp_gen/executor.rst`. - -The following sections provide an overview of the main classes defined in the Executor API. - -## API - -### The Executor Class - -The `Executor` class is responsible for receiving requests from the client, and providing responses for those requests. The executor is constructed by providing a path to a directory containing the TensorRT-LLM engine or buffers containing the engine and the model JSON configuration. The client can create requests and enqueue those requests for execution using the `enqueueRequest` or `enqueueRequests` methods of the `Executor` class. Enqueued requests will be scheduled for execution by the executor, and multiple independent requests can be batched together at every iteration of the main execution loop (a process often referred to as continuous batching or iteration-level batching). Responses for a particular request can be awaited for by calling the `awaitResponses` method, and by providing the request id. Alternatively, responses for any requests can be awaited for by omitting to provide the request id when calling `awaitResponses`. The `Executor` class also allows to cancel requests using the `cancelRequest` method and to obtain per-iteration and per-request statistics using the `getLatestIterationStats`. - -### The Request Class - -The `Request` class is used to define properties of the request, such as the input token ids and the maximum number of tokens to generate. The `streaming` parameter can be used to indicate if the request should generate a response for each new generated tokens (`streaming = true`) or only after all tokens have been generated (`streaming = false`). Other mandatory parameters of the request include the sampling configuration (defined by the `SamplingConfig` class) which contains parameters controlling the decoding process and the output configuration (defined by the `OutputConfig` class) which controls what information should be included in the `Result` for a particular response. - -Optional parameters can also be provided when constructing a request such as a list of bad words, a list of stop words, a client id, or configurations objects for prompt tuning, LoRA, or speculative decoding, or a number of sequences to generate for example. - -### The Response Class - -The `awaitResponses` method of the `Executor` class returns a vector of responses. Each response contains the request id associated with this response, and also contains either an error or a `Result`. Check if the response has an error by using the `hasError` method before trying to obtain the `Result` associated with this response using the `getResult` method. - -### The Result Class - -The `Result` class holds the result for a given request. It contains a Boolean parameter called `isFinal` that indicates if this is the last `Result` that will be returned for the given request id. It also contains the generated tokens. If the request is configured with `streaming = false` and `numReturnSequences = 1`, a single response will be returned, the `isFinal` Boolean will be set to `true` and all generated tokens will be included in the `outputTokenIds`. If `streaming = true` and `numReturnSequences = 1` is used, a `Result` will include one or more tokens (depending on the request `returnAllGeneratedTokens` parameter) except the last result and the `isFinal` flag will be set to `true` for the last result associated with this request. - -The request `numReturnSequences` parameter controls the number of output sequences to generate for each prompt. When this option is used, the Executor will return at least `numReturnSequences` responses for each request, each containing one Result. In beam search (`beamWidth > 1`), the number of beams to be returned will be limited by `numReturnSequences` and the `sequenceIndex` attribute of the `Result` class will always be zero. Otherwise, in sampling (`beamWidth = 1`), the `sequenceIndex` attribute indicates the index of the generated sequence in the result (`0 <= sequenceIndex < numReturnSequences`). It contains a Boolean parameter called `isSequenceFinal` that indicates if this is the last result for the sequence and also contains a Boolean parameter `isFinal` that indicates when all sequences for the request have been generated. When `numReturnSequences = 1`, `isFinal` is identical to `isSequenceFinal`. - -Here is an example that shows how a subset of 3 responses might look like for `numReturnSequences = 3`: - -``` -Response 1: requestId = 1, Result with sequenceIndex = 0, isSequenceFinal = false, isFinal = false -Response 2: requestId = 1, Result with sequenceIndex = 1, isSequenceFinal = true, isFinal = false -Response 3: requestId = 1, Result with sequenceIndex = 2, isSequenceFinal = false, isFinal = false -``` - -In this example, each response contains one result for different sequences. The `isSequenceFinal` flag of the second Result is set to true, indicating that it is the last result for `sequenceIndex = 1`, however, the isFinal flag of each Response is set to false because sequences 0 and 2 are not completed. - -### Sending Requests with Different Beam Widths - -The executor can process requests with different beam widths if the following conditions are met: - -- The model was built with a `max_beam_width > 1`. -- The executor is configured with a `maxBeamWidth > 1` (the configured `maxBeamWidth` must be less than or equal to the model's `max_beam_width`). -- The requested beam widths are less than or equal to the configured `maxBeamWidth`. - -The executor may schedule successive requests with the same beam width at the same time. For successive requests with two different beam widths, `x` and `y`, requests with beam width `y` are not scheduled until all requests with beam width `x` have been processed. -This allows the runtime to reconfigure itself for a new beam width when no requests are in flight. The reconfiguration happens automatically each time requests with a different beam width than currently configured are detected. Waiting for previous requests to finish and reconfiguring the runtime may cause significant overhead and reduce overall throughput. - -### Controlling output with Logits Post-Processor - -Optionally, you can alter the logits produced by the network by providing an instance of `Executor::LogitsPostProcessorConfig`. For instance, this feature can be used to generate JSON formatted output. {cpp:class}`Executor::LogitsPostProcessorConfig ` specifies a map of named callbacks in the following form - -```cpp -std::unordered_map)>> -``` - -The map key is the name associated with that logits post-processing callback. Each request can then specify the name of the logits post-processor to use for that particular request, if any. - -The first argument to the callback is the request id, second is the logits tensor, third are the tokens produced by the request so far, fourth is the operation stream used by the logits tensor, and last one is an optional client id. The callback returns a modified tensor of logits. Multiple requests can share same client id and callback can use different logic based on client id. - -You must use the stream to access the logits tensor. For example, to perform an addition with a bias tensor, the addition operation is enqueued on that stream. Alternatively, you can call `stream->synchronize()`, however, that will slow down the entire execution pipeline. - -The executor also includes a {cpp:class}`LogitsPostProcessorBatched ` method that enables altering logits of multiple requests in a batch. The batched method allows further optimizations and reduces callback overheads. - -```cpp -std::function const&, std::vector&, std::vector> const&, StreamPtr const&, std::vector> const&)> -``` - -A single batched callback can be specified in `LogitsPostProcessorConfig`. Each request can opt to apply this callback by specifying the name of the logits post-processor as `Request::kBatchedPostProcessorName`. - -Note: Neither callback variant is supported with the `STATIC` batching type for the moment. - -In a multi-GPU run, the callback is invoked on all ranks in the first tensor-parallel group, by default. To ensure correct execution, replicate the client-side state that is accessed by the callback on these ranks. If replication is expensive or infeasible, use `LogitsPostProcessorConfig::setReplicate(false)` to invoke the callback only on rank 0. The executor broadcasts the sampled tokens internally to ensure correct execution. - -### Structured output with guided decoding -Guided decoding controls the generation outputs to be amenable to pre-defined structured formats, e.g., JSON or XML. Currently, guided decoding is supported with the [XGrammar](https://github.com/mlc-ai/xgrammar) backend. - -To enable guided decoding, a valid instance of `GuidedDecodingConfig` must be provided when constructing `Executor`. `GuidedDecodingConfig` should be constructed with some tokenizer information, including `encodedVocab`, `tokenizerStr` (optional) and `stopTokenIds` (optional). Given a Hugging Face tokenizer, these can be extracted by: - -```python -encoded_vocab = tokenizer.get_vocab() -encoded_vocab = [token for token, _ in sorted(encoded_vocab.items(), key=lambda x: x[1])] -tokenizer_str = tokenizer.backend_tokenizer.to_str() -stop_token_ids = [tokenizer.eos_token_id] -``` - -Refer to [`tensorrt_llm/llmapi/tokenizer.py`](source:tensorrt_llm/llmapi/tokenizer.py) for more details. You may dump these materials to disk, and reload them to C++ runtime for use. - -Each request can be optionally specified with a `GuidedDecodingParams`, which defines the desired structured format. Currently, it supports four types: -* `GuidedDecodingParams::GuideType::kJSON`: The generated text is amenable to JSON format; -* `GuidedDecodingParams::GuideType::kJSON_SCHEMA`: The generated text is amenable to JSON format with additional restrictions; -* `GuidedDecodingParams::GuideType::kREGEX`: The generated text is amenable to regular expression; -* `GuidedDecodingParams::GuideType::kEBNF_GRAMMAR`: The generated text is amenable to the extended Backus-Naur form (EBNF) grammar. - -The latter three types should be used with the schema/regex/grammar provided to `GuidedDecodingParams`. - -### Obtaining Arbitrary Output Tensors -The executor API gives the user the possibility to read the arbitrary outputs from the model. For example, it is possible to obtain hidden states or logits. - -#### Mark Tensors As Output -For a tensor to be obtainable using this feature, it needs to be marked as an output in the model definition (e.g. add `topk_logits.mark_output("TopKLogits")`) before building the TRT engine. - -#### Configure The Executor -Assuming the TensorRT engine you are planning to use has a tensor named `TopKLogits` marked as output, you should then configure the `Executor` to read from this output tensor by passing its name to the `ExecutorConfig` configuration object: -```cpp -auto const executorConfig = ExecutorConfig{}; - -std::vector additionalOutputs{ - executor::AdditionalModelOutput{"TopKLogits", /*whether or not to get the output for the context too */ true}}; -executorConfig.setAdditionalModelOutputs(additionalOutputs); - -// ... set more configuration options if needed -// ... create the `Executor` instance -``` - -### Request Additional Output -Construct a request to enqueue in the executor to query this tensor output: -```cpp -executor::Request request{requestTokens, parameters.maxOutputLength, true, executor::SamplingConfig{}, - executor::OutputConfig{false, false, false, true, false, false, additionalOutputs}}; -executor.enqueueRequest(request); -``` - -The output can be found at the `additionalOutputs` property of each response. - -#### Note on context outputs - -If KV cache reuse is enabled, context outputs will not contain outputs for the part of the context that has been reused. This part of the outputs can only be obtained from the prior request with the same prefix that generated this part of the KV cache. - -## C++ Executor API Example - -Two C++ examples are provided that show how to use the Executor API and can be found in the [`examples/cpp/executor`](source:examples/cpp/executor/) folder. - -## Python Bindings for the Executor API - -Python bindings for the Executor API are also available to use the Executor API from Python. The Python bindings are defined in [bindings.cpp](source:cpp/tensorrt_llm/nanobind/executor/bindings.cpp) and once built, are available in package `tensorrt_llm.bindings.executor`. Running `help('tensorrt_llm.bindings.executor')` in a Python interpreter will provide an overview of the classes available. - -In addition, three Python examples are provided to demonstrate how to use the Python bindings to the Executor API for single and multi-GPU models. They can be found in [`examples/bindings`](source:examples/bindings). - -## In-flight Batching with the Triton Inference Server - -A Triton Inference Server C++ [backend](https://github.com/triton-inference-server/tensorrtllm_backend) is provided with TensorRT-LLM that -includes the mechanisms needed to serve models using in-flight batching. That -backend is also a good starting example of how to implement in-flight batching using -the TensorRT-LLM C++ Executor API. diff --git a/docs/source/legacy/advanced/expert-parallelism.md b/docs/source/legacy/advanced/expert-parallelism.md deleted file mode 100644 index 55d948a80cae..000000000000 --- a/docs/source/legacy/advanced/expert-parallelism.md +++ /dev/null @@ -1,30 +0,0 @@ -(expert-parallelism)= - -# Expert Parallelism in TensorRT-LLM - -## Mixture of Experts (MoE) - -Mixture of Experts (MoE) architectures have become widespread, with models such as [Mistral Mixtral 8×7B](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1). Specifically, MoE’s structure supports multiple parallel feed-forward neural-network (FFN) layers (called experts) in place of the single FFN layer in a dense model. When tokens arrive, the router layer selects the top-k experts for each token, and the corresponding hidden state of each token is dispatched to those experts. As a result, there are multiple tokens’ hidden states that are dispatched to each expert. - -moe_structure - -the MOE structure in Switch Transformer: [https://arxiv.org/pdf/2101.03961.pdf](https://arxiv.org/pdf/2101.03961.pdf) - -## Tensor Parallel vs Expert Parallel - -Parallelism on multi-GPUs is necessary if the MoE model cannot be accommodated by a single GPU’s memory. We support three kinds of parallel patterns for MoE structure: Tensor Parallel (default pattern), Expert Parallel, and a hybrid of the two. - -tensor parallel vs expert parallel - -Tensor Parallel evenly splits each expert’s weight and distributes them to different GPUs, which means each GPU holds partial weight of all experts, While Expert Parallel evenly distributes some of the experts’ full weight to different GPUs, which means each GPU holds part of the experts’ full weight. As a result, each GPU rank in the Tensor Parallel group receives all tokens’ hidden states for all experts, then computes using the partial weights, while for Expert Parallel, each GPU rank only receives part of tokens’ hidden states for experts on this rank, then computes using the full weights. - -When both Tensor Parallel and Expert Parallel are enabled, each GPU handles a portion of the expert weights matrices (as in EP mode) and these weights are further sliced across multiple GPUs (as in TP mode). This hybrid approach aims to balance the workload more evenly across GPUs, enhancing efficiency and reducing the likelihood of bottlenecks associated with EP mode alone. - - -## How to Enable - -The default parallel pattern is Tensor Parallel. You can enable Expert Parallel or hybrid parallel by setting `--moe_tp_size` and `--moe_ep_size` when calling `convert_checkpoint.py`. If only `--moe_tp_size` is provided, TRT-LLM will use Tensor Parallel for the MoE model; if only `--moe_ep_size` is provided, TRT-LLM will use Expert Parallel; if both are provided, the hybrid parallel will be used. - -Ensure the product of `moe_tp_size` and `moe_ep_size` is equal to `tp_size`, since the total number of MoE parallelism across all GPUs must match the total number of parallelism in other parts of the model. - -The other parameters related to the MoE structure, such as `num_experts_per_tok` (TopK in previous context) and `num_local_experts`, can be found in the model’s configuration file, such as the one for [Mixtral 8x7B model](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1/blob/main/config.json). diff --git a/docs/source/legacy/advanced/gpt-attention.md b/docs/source/legacy/advanced/gpt-attention.md deleted file mode 100644 index d5b4d29718bf..000000000000 --- a/docs/source/legacy/advanced/gpt-attention.md +++ /dev/null @@ -1,343 +0,0 @@ -(gpt-attention)= - -# Multi-Head, Multi-Query, and Group-Query Attention - -This document details the implementation of multi-head attention (MHA), -multi-query attention (MQA) and group-query attention (GQA) for auto-regressive -GPT-like models in TensorRT-LLM. As a quick reminder, the multi-head attention -is the sequence of a batched matmul, a softmax and another batched matmul -described in the -[Attention Is All You Need](https://arxiv.org/abs/1706.03762) article. [Multi-query Attention (MQA)](https://arxiv.org/abs/1911.02150) and [Group-query Attention (GQA)](https://arxiv.org/abs/2307.09288) are variants of MHA that use fewer, so-called, K/V head than the number of query heads. TensorRT-LLM, MHA, MQA and GQA are implemented by the operator [`tensorrt_llm.functional.gpt_attention`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/functional.py). - -## Important Note - -As discussed below, the current implementation supports two input modes: Padded -and packed (non-padded). As the packed mode is always more memory-efficient and -faster than the padded mode, ***support for padded mode may be removed in the -future***. - -## Padded and Packed Tensors - -In TensorRT-LLM, the GPT attention operator supports two different types -of QKV inputs: Padded and packed (i.e. non padded) inputs. The mode is -determined by the global configuration parameter `remove_input_padding` defined -in [`tensorrt_llm.plugin`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/plugin/plugin.py). - -When padding is enabled (that is, `remove_input_padding` is `False`), the sequences -that are shorter than the `max_sequence_length` are padded to that maximum -length. It may result in excessive memory consumption as well as unneeded -computations on padding tokens (in the various matrix multiplications that -surround the MHA block). - -To overcome that problem, TensorRT-LLM supports a mode without padding where -the different tokens are packed together and the user provides the operator -with a 1D tensor containing the lengths of the different sequences. It is -recommended that users to always use packed mode (and support for the padded -mode may be removed in the future). - -## Context and Generation Phases - -The GPT attention operator encapsulates different implementations for both -context and generation phases in auto-regressive models like GPT. - -### Context Phase - -If the `context_fmha_type` is set to `disabled` (refer to -[`tensorrt_llm.plugin`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/plugin/plugin.py)), -the implementation maps to a sequence of GPU kernels that will store the -intermediate `Q*K^T` tensor in memory before calling the softmax operator. It -is the slowest method and the memory footprint is significant (quadratically -depends on the sequence length). - -Otherwise, if `context_fmha_type` is set to a `enabled` or -`enabled_with_fp32_acc` (accumulation in the first batched matmul is forced to -FP32), that function will trigger a kernel that performs the MHA/MQA block -using a single kernel. For short sequences, that kernel uses a vanilla -implementation of MHA/MQA. For larger sequences, this kernel uses the Flash -Attention algorithm as described in -[FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://arxiv.org/abs/2205.14135) -and -[FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://arxiv.org/abs/2307.08691). - -Currently, the implementation triggers extra kernels that apply pre-processing -to the elements (like RoPE) and populate the KV cache (see below). In a future -release, the number of such kernels is planned on being reduced in order to -improve the overall performance. - -#### FP8 Context FMHA - -When FP8 quantization is activated, the attention can be further accelerated by -enabling FP8 Context FMHA (`use_fp8_context_fmha = enable`). - -FP8 Paged Context FMHA is also supported with the fp8 quantization workflow. -You need to specify `use_fp8_context_fmha = enable` and -`use_paged_context_fmha = enable` at the same time. - -Please be aware that this feature is only supported on Ada and Hopper. - -### Generation Phase - -The generation phase is implemented using a single kernel called the masked -multi-head attention in TensorRT-LLM. That kernel is able to apply -pre-processing on the Q, K, and V elements on-the-fly: adds the QKV bias, applies -RoPE, and performs dequantization and quantization. TensorRT-LLM will continue to add (or -enable) additional features in future releases. For example, enable the support -for IA3. - -_The masked MHA kernel has a special version that distributes the work across -multiple CUDA thread-blocks on the GPU for cases where the GPU occupancy is -low. That mode called multi-block is turned on by default starting from TRT-LLM 0.13, -and can be disabled using `--multi_block_mode=False` during runtime. -Users are recommended to test that mode in scenarios where both the batch -size and the number of heads in the model are relatively small. The exact -definition of small in that context will depend on the model of the GPU and is -hard to predict but to provide with a rule of thumb, it is worth testing that -mode when `batch_size * num_heads` is less than the number of multi-processors -on the GPU (that suggestion may evolve in the future as more research is -conducted and the software improves)_. - -_Note that even if the multi-block mode is enabled, the attention operator will -not immediately trigger the multi-block version of the GPU kernel. There is a -minimum number of tokens (input + generated) that are required for the -multi-block version to become more efficient than the "vanilla" implementation -that uses a single CUDA thread-block per head. It is controlled by an internal -heuristic._ - -Another note is that as the masked MHA kernels use shared memory size -proportional to sequence length, so there can be some cases that GPU's shared -memory is not enough when multi-block mode is not enabled. To get masked MHA -kernel work in these cases, multi-block mode is forced on and a warning log is -printed. - -#### XQA Optimization - -Another optimization for MQA/GQA in the generation phase is called XQA optimization. - -Support matrix of the XQA optimization: - - FP16 / BF16 compute data type. - - FP16 / BF16 / FP8 / INT8 KV cache data type. - - Paged KV cache (8 / 16 / 32 / 64 / 128 tokens per block). - -By default, this is enabled. Note that a heuristic algorithm -is also used to decide whether to use XQA kernel or masked MHA kernel to get -better performance. -If you want to use that kernel whenever possible, set `TRTLLM_FORCE_XQA=1` to force use of the XQA kernel when the model config is supported. -Supported configurations can be found using the `shouldUse` function of the `DecoderXQARunner` class in -`cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.h`. - - -(inflight-batching)= - -## In-flight Batching - -TensorRT-LLM supports in-flight batching of requests (also known as continuous -batching or iteration-level batching) for higher serving throughput. With this feature, -sequences in context phase can be processed together with sequences in -generation phase. The purpose of that technique is to better interleave -requests to reduce latency as well as make better use of the GPUs. -For efficiency reasons (1), the support for inflight batching ***requires the -input tensors to be packed (no padding)***. - -***In the current implementation, the sequences that are going through the -context phase must be before the sequences in the generation phase in the input -tensor. For example, for sequences `S0`, `S1` and `S2`, if `S0` and `S2` are in -context phase (and `S1` in generation), tokens from `S0` and `S2` must appear -before the tokens of `S1` in the input tensor***. The constraint may or may not -be relaxed in a future version. - -_(1) Padding sequences in the generation phase, that contain a single token, to -the length of the maximum input sequence is inefficient use of resources_. - - - -## Chunked Context - -In the original state, the common behavior was to process all context tokens at -once. This feature splits the context into several chunks. In this way, the -context chunks can be batched with more tokens during the generation phase, -which is expected to increase the total throughput. Chunking contexts also removes -constraints on input length. To enable this feature, the FMHA paged kv-cache also -needs to be enabled. Except for the last one, the size of the context chunk needs -to be an integer multiple of the kv-cache block size. Refer to -[the chunked context feature](../../features/long-sequence.md) for usage. - -## KV Cache - -In the generation phase, a common optimization is to provide the MHA kernel -with a cache containing the values of the past K and V elements that have -already been computed. That cache is known as the KV cache. TensorRT-LLM uses -that technique to accelerate its generation phase. In TensorRT-LLM, there is -one KV cache per Transformer layer, which means that there are as many KV -caches as layers in a model. The current version of TensorRT-LLM supports two -different types of KV caches: **contiguous** and **paged** KV caches. - -### Contiguous KV Cache - -The contiguous KV cache is a monolithic tensor. Its shape is: -``` -[max_batch_size * max_beam_width, 2, num_heads, max_seqlen, hidden_dim_per_head]. -``` - -That implementation uses a lot more memory than needed when the sequences are -shorter than the maximum sequence length (even if they end up close to the -limit after the generation of many output tokens, it may take a lot of steps to -reach that point). - -### Paged KV Cache - -The paged KV cache decomposes the KV cache into blocks that are distributed to -the different requests by a cache manager during processing. That cache manager -keeps track of the sequences, allocate new blocks from a pool and recycle those -blocks when required. See the simplified implementation of -[`tensorrt_llm.runtime.KVCacheManager`](source:tensorrt_llm/runtime/kv_cache_manager.py). -A more efficient C++ implementation is included in the -[Batch Manager](source:cpp/include/tensorrt_llm/batch_manager). - -## INT8/FP8 KV Caches - -In its current implementation, even if the rest of the network runs in INT8 or -FP8, the GPT attention operator works with FP32, FP16, and BFloat16 inputs and -outputs. However, TensorRT-LLM supports INT8 and FP8 -(`kv_cache_quant_mode=QuantMode.INT8_KV_CACHE` and -`kv_cache_quant_mode=QuantMode.FP8_KV_CACHE`) KV caches. - -The GPT attention operator populates the KV cache. When INT8 or FP8 KV caches -are enabled, the input values have to be quantized to 8 bits using a scaling -factor. For quantization, the scaling factor is stored in the -`kv_cache_scaling_factor` tensor. Its shape is `[1]` and only per-tensor -quantization is supported in the current version. Quantization uses inversed scale -since it does multiply as `fp_value * (1.0 / kv_cache_scaling_factor)` in plugin. - -During generation, the values read from the cache are dequantized on-the-fly in -the MHA/MQA kernel, dequantization can be described as -`quantized_value * kv_cache_scaling_factor`. - - -## Sliding Window Attention, Cyclic (Rolling Buffer) KV Cache - -TensorRT-LLM has a feature called `Cyclic KV Cache`, which treats the kv cache -as a circular buffer. This means that it only stores the kv cache for the last N -tokens, where N is determined by the `max_attention_window_size` parameter in -`GenerationSession.setup`. You can see examples of this in the `run.py` or -`summarize.py` files. When the cache is full, new tokens’ kv cache will -overwrite the "least recently used" caches. - -In the context phase, if the input length surpasses the `max_attention_window_size`, -`Sliding Window Attention` will be activated. This serves the same function as -the `sliding window_size`. - -This feature helps to reduce the memory footprint of the kv cache when -dealing with very long sequences. - -The feature, which allows different `max_attention_window_size` values -for each layer, is also supported. To utilize this feature, simply provide an -`int32 torch.Tensor` or `list` to the `GenerationSession.setup` when using python -runtime session, or provide a vector to the `KvCacheConfig` when using cpp runtime. -If the number of the provided elements is less than the number of layers, the provided -tensor/list/vector will be repeated multiple times to the number of layers and then be -saved as a new tensor. This tensor will serve as the buffer for `max_attention_window_size`, -setting unique values for each layer. However, it’s important to note that the -memory allocation for the kv cache still relies on the buffer’s maximum value. - -_Note that the cyclic kv cache feature doesn't work with beam searching currently as -the context kv cache are shared across beams. - -## StreamingLLM - -The StreamingLLM feature uses a window attention to perform efficient and stable LLM -on long texts, which means that only `N` tokens need to be stored in the KV cache. -Similar to the cyclic KV cache feature in TensorRT-LLM, `max_attention_window_size` -parameter is used to determine `N`. Different from the cyclic KV cache feature, -the first `S` tokens, called sink tokens, are always kept in the attention window, -where `S` is determined by `sink_token_length` parameter in `GenerationSession.setup`. -But in context phase, the self-attention is dense in the official implementation of -StreamingLLM, and it uses all of the tokens for computation and only saves `N` tokens -to the KV cache. - -In addition, the relative position embedding is also changed in StreamingLLM. -When determining the relative distance and adding positional information to tokens, -StreamingLLM uses the positions within the cache rather than those in the original text. - -`streamingllm` flag is used to enable this feature. - -## Beam-Search - -The GPT attention operator supports beam-search. In the context phase, a single -beam is computed per input sequence. In the generation phase, the MHA/MQA/GQA -kernel uses an additional tensor to reconstruct the correct path for each beam. -That tensor is called the `cache_indirection`. Its shape is `[batch_size, -beam_width, max_seqlen]`. - -For a sequence `si`, a beam `bi` and a token `ti`, the element -`cache_indirection[si][bi][ti]` is an integer between `0` and `beam_width-1` -that indicates which path in the beam to read the K and V elements from in the -KV cache. This tensor is populated in the sampling stage. - -## Input QKV tensor - -The input QKV tensor packs the Q, K and V tensors (concatenated along the last -dimension) after the projection of the hidden states. It is a 3D tensor. RoPE -and quantization to INT8 or FP8 (when needed) are performed by the GPT -attention operator. - -In padded mode, its shape is `[batch_beam_size, max_seqlen, 3 * hidden_dim]` -where `batch_beam_size` is the batch size (number of sequences) for the context -phase and the batch size multiplied by the beam width for the generation phase. -Having different beam widths per sequence in padded mode is not supported. - -In packed mode, its shape is `[num_tokens, 3 * hidden_dim]` where -`num_tokens` is the total number of tokens in the batch. For the sequences in -context phase, the number of tokens of a sequence corresponds to its input -length (even if the beam width is greater than `1` for beam search). For the -sequences in generation phase, there are `beam_width` tokens per sequence. The -beam width can be different for each sequence. - -In other words, the pseudo-code to compute the number of tokens is: - -```python -num_tokens = 0 - -# Add the length of each sequence in context phase. -for seq in context_phase: - num_tokens += seq.length - -# Add the width of the beam for each sequence in generation phase. -for seq in generation_phase: - num_tokens += seq.beam_width -``` - -### Rotary Positional Embedding (RoPE) - -The GPT attention operation can perform the computation of the Rotary -Positional Embedding (RoPE). When that operation is enabled, -`rotary_embedding_dim` is set to a value greater than 0, it is fused with other -operations. The GPT operator supports GPT-NeoX and GPT-J forms of RoPE by -setting `position_embedding_type` to `PositionEmbeddingType.rope_gpt_neox` -or `PositionEmbeddingType.rope_gptj`. - -### ALiBi - -The GPT attention operator can apply ALiBi to the result of the `Q*K^T` -product. The bias is computed on-the-fly from the ALiBi slopes in the optimized -kernel. - -### Scaling factor(s) - -In MHA, the output of the `Q*K^T` product is scaled by a constant value that -is computed as: - -``` -norm_factor = 1.f / (q_scaling * sqrt(head_size)). -``` - -### Cross Attention - -On top of the MHA as self attention needed by GPT-style decoder-only models, `gpt_attention` also supports cross attention. - -This enables using `gpt_attention` in a broader aspect as a generic decoder component. For example, the Encoder-Decoder model uses `gpt_attention` to issue both the self attention and cross attention modules in its Decoder. - -### Relative Attention Bias (RAB) - -Relative attention bias (RAB) is a kind of relative position modeling, adding an attention bias (`Q*K^T+bias`) according to relative positions. RAB is a lightweight method to include the information of relative positions, and is used in the popular Encoder-Decoder model [T5](https://huggingface.co/docs/transformers/model_doc/t5) and also other models in the T5 family. - -RAB is supported in two modes: i) regular mode which user passes in relative attention bias computed ahead of MHA. ii) implicit mode which computes the relative attention bias on the fly in MHA. The implicit mode suits the case when the relative attention bias is too large to fit in memory and can be turned on by passing in `max_distance`. diff --git a/docs/source/legacy/advanced/gpt-runtime.md b/docs/source/legacy/advanced/gpt-runtime.md deleted file mode 100644 index e384c45175e4..000000000000 --- a/docs/source/legacy/advanced/gpt-runtime.md +++ /dev/null @@ -1,202 +0,0 @@ -(gpt-runtime)= - -# C++ GPT Runtime - -TensorRT-LLM includes a C++ component to execute TensorRT engines built with -the Python API as described in the {ref}`architecture-overview` section. -That component is called the C++ runtime. - -The API of the C++ runtime is composed of the classes declared in -[`cpp/include/tensorrt_llm/runtime`](https://github.com/NVIDIA/TensorRT-LLM/tree/main/cpp/include/tensorrt_llm/runtime) and -implemented in [`cpp/tensorrt_llm/runtime`](https://github.com/NVIDIA/TensorRT-LLM/tree/main/cpp/tensorrt_llm/runtime). - -Even if the different components described in that document mention GPT in -their name, they are not restricted to this specific model. Those classes can -be used to implement auto-regressive models like BLOOM, GPT-J, GPT-NeoX or -LLaMA, for example. - -Complete support of encoder-decoder models, like T5, will be added to -TensorRT-LLM in a future release. An experimental version, only in Python for -now, can be found in the [`examples/models/core/enc_dec`](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/enc_dec) folder. - -## Overview - -Runtime models are described by an instance of the -[`ModelConfig`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/include/tensorrt_llm/runtime//modelConfig.h) -class and a pointer to the TensorRT engine that must be -executed to perform the inference. -The environment is configured through the -[`WorldConfig`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/include/tensorrt_llm/runtime/worldConfig.h) -(that name comes from -[MPI](https://en.wikipedia.org/wiki/Message_Passing_Interface) and its "famous" -`MPI_COMM_WORLD` default communicator). -The [`SamplingConfig`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/include/tensorrt_llm/runtime/samplingConfig.h) -class encapsulates parameters that control the -[generation](https://huggingface.co/blog/how-to-generate) of new tokens. - -### Model Configuration - -The model configuration is an instance of the -[`ModelConfig`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/include/tensorrt_llm/runtime//modelConfig.h) class. -That class encapsulates the following parameters (they are declared as private -member variables and exposed through getters and setters): - - * `vocabSize`, the size of the vocabulary, - * `numLayers`, the number of layers in the model, - * `numHeads`, the number of heads in the attention block, - * `numKvHeads`, the number of heads for K and V in the attention component. - When the number of K/V heads is the same as the number of (Q) heads, the - model uses multi-head attention. When the number of K/V heads is 1, it uses - multi-query attention. Otherwise, it uses group-query attention. Refer to {ref}`gpt-attention` for more information, - * `hiddenSize`, the size of the hidden dimension, - * `dataType`, the datatype that was used to build the TensorRT engine and that - must be used to run the model during inference, - * `useGptAttentionPlugin`, indicates if the {ref}`gpt-attention` operator was compiled using the - [GPT Attention plugin](https://github.com/NVIDIA/TensorRT-LLM/tree/main/cpp/tensorrt_llm/plugins/gptAttentionPlugin), - * `inputPacked`, indicates that the input must be packed (or padded when set - to `false`). For performance reasons, it is recommended to always use packed, - even if its default is set to `false` (will be changed in a future release). - Refer to {ref}`gpt-attention` for more information, - * `pagedKvCache`, indicates if the K/V cache uses paging. - Refer to {ref}`gpt-attention` for more information, - * `tokensPerBlock`, is the number of tokens in each block of the K/V cache. - It's relevant when the paged K/V cache is enabled. By default, the value is - 64. Refer to {ref}`gpt-attention` for more information, - * `quantMode`, controls the quantization method. Refer to {ref}`precision` for more information. - * `maxBatchSize`, indicates the maximum batch size that the TensorRT engine - was built for, - * `maxInputLen`, the maximum size of the input sequences, - * `maxSequenceLen`, the maximum total size (input+output) of the sequences. - -### World Configuration - -Familiarity with -[MPI](https://en.wikipedia.org/wiki/Message_Passing_Interface), is not required -to utilize the TensorRT-LMM C++ runtime. There are two main things -you need to know: -* The C++ Runtime in TensorRT-LLM uses -[processes](https://en.wikipedia.org/wiki/Process_(computing)) to execute -TensorRT engines on the different GPUs. Those GPUs can be located on a single -node as well as on different nodes in a cluster. Each process is called a -*rank* in MPI. -* The ranks are grouped in communication groups. The -TensorRT-LLM C++ Runtime calls that group the *world*. - -The world configuration is an instance of the -[`WorldConfig`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/include/tensorrt_llm/runtime/worldConfig.h) -class, which encapsulates the following parameters: - -* `tensorParallelism`, the number of ranks that collaborate together to - implement Tensor Parallelism (TP). With TP, each GPU performs computations for - all the layers of the model. Some of those computations are distributed - across the GPU. TP is more balanced than Pipeline Parallelism (PP), in most cases, but - requires higher bandwidth between the GPUs. It is the recommended setting in - the presence of NVLINK between GPUs, -* `pipelineParallelism`, the number of ranks that collaborate together to - implement Pipeline Parallelism (PP). With PP, each GPU works on a subset of - consecutive layers. Communications between the GPUs happen only at the - boundaries of the subsets of layers. It is harder to guarantee the full - utilization of the GPUs with PP but it requires less memory bandwidth. It - is the recommended setting in the absence of NVLINK between GPUs, -* `rank`, the unique identifier of the rank, -* `gpusPerNode`, indicates the number of GPUs on each node. Having that - information allows the C++ runtime to optimize communications between GPUs in - a node (like taking advantage of the - [NVLINK](https://www.nvidia.com/en-us/data-center/nvlink/) - interconnect between GPUs of an A100 - [DGX](https://www.nvidia.com/en-us/data-center/dgx-platform/) - node). - -### Sampling Parameters - -The [`SamplingConfig`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/include/tensorrt_llm/runtime/samplingConfig.h) -class encapsulates parameters that control the -[generation](https://huggingface.co/blog/how-to-generate) of new tokens. -A comparison of selecting decoding method is listed as the table below (`X` means it is not supported yet). -Except for the `beamWidth` parameter, all the fields are optional and the -runtime will use a default value if no values are provided by the user. For -vector fields, the TensorRT-LLM runtime supports one value per sequence (that is, -the vector contains `batchSize` values). If all the sequences use the same -value for a given parameter, the vector can be limited to a single element -(that is, `size() == 1`). - -| Method name in HF | Condition in HF | Method name in TRT-LLM | Condition in TRT-LLM | -| :------------------------------: | :---------------------------------------------------: | :--------------------: | :--------------------------------------------: | -| assisted decoding | `assistant_model` or `prompt_lookup_num_tokens!=None` | X | | -| beam-search decoding | `num_beams>1` and `do_sample=False` | beam search | `beamWidth > 1` | -| beam-search multinomial sampling | `num_beams>1` and `do_sample=True` | X | | -| constrained beam-search decoding | `constraints!=None` or `force_words_ids!=None` | X | | -| contrastive search | `penalty_alpha>0` and `top_k>1` | X | | -| diverse beam-search decoding | `num_beams>1` and `num_beam_groups>1` | X | | -| greedy decoding | `num_beams=1` and `do_sample=False` | sampling | `beamWidth == 1` and `topK=0` and `topP=0.0f` | -| multinomial sampling | `num_beams=1` and `do_sample=True` | sampling | `beamWidth == 1` and (`topK>0` or `topP>0.0f`) | - -***General*** - -| Name in TRT-LLM | Description | Data type | Range of value | Default value | Name in HF | -| :-----------------: | :-------------------------------------------------------------------------------: | :-----------: | :---------------------------------------------------------------------------------------: | :---------------------------------------------------: | :--------------------: | -| `temperature` | modulation of logits in sampling workflow | List\[Float\] | \[0.0f, $+\infty$\) | `1.0f` (no modulation) | `temperature` | -| `minLength` | lower-bound on the number of tokens generated | List\[Int\] | \[0, $+\infty$\) | `0` (no effect (the first generated token can be EOS) | `min_length` | -| `repetitionPenalty` | penalize repetitive tokens
multiplicative, irrespective of appearances count | List\[Float\] | \[0.0f, $+\infty$\)
`< 1.0f` encourages repetition
`> 1.0f` discourages it | `1.0f` (no effect) | `repetition_penalty` | -| `presencePenalty` | penalize existed tokens
additive, irrespective of appearances count | List\[Float\] | \($-\infty$, $+\infty$\)
`< 0.0f` encourages repetition
`> 0.0f` discourages it | `0.0f` (no effect) | no | -| `frequencyPenalty` | penalize existed tokens
additive, dependent on appearances count | List\[Float\] | \($-\infty$, $+\infty$\)
`< 0.0f` encourages repetition
`> 0.0f` discourages it | `0.0f` (no effect) | no | -| `noRepeatNgramSize` | | List\[Int\] | \[0, $+\infty$\)
`> 0` all ngrams of that size can only occur once | `0` (no effect) | `no_repeat_ngram_size` | - -* The tokens of input prompt are included during adopting `repetitionPenalty`, `presencePenalty`, and `frequencyPenalty` onto logits. - -* The parameters `repetitionPenalty`, `presencePenalty`, and `frequencyPenalty` are not mutually exclusive. - -***Sampling*** - -| Name in TRT-LLM | Description | Data type | Range of value | Default value | Name in HF | -| :-------------: | :---------------------------------------------------------------------: | :-----------: | :---------------: | :--------------: | :--------: | -| `randomSeed` | random seed for random number generator | Int64 | \[0, 2^64-1\] | `0` | no | -| `topK` | the number of logits to sample from | List\[Int\] | \[0, 1024\] | `0` | `top_k` | -| `topP` | the top-P probability to sample from | List\[Float\] | \[0.0f, 1.0f\] | `0.0f` | `top_p` | -| `topPDecay` | the decay in the `topP` algorithm | List\[Float\] | \(0.0f, 1.0f\] | `1.0f` | no | -| `topPMin` | the decay in the `topP` algorithm | List\[Float\] | \(0.0f, 1.0f\] | `1.0e-6,f` | no | -| `topPResetIds` | the decay in the `topP` algorithm | List\[Int\] | \[-1, $+\infty$\) | `-1` (no effect) | no | -| `minP` | scale the most likely token to determine the minimum token probability. | List\[Float\] | \[0.0f, 1.0f\] | `0.0` (no effect) | `min_p` | - - * If setting `topK = 0` and `topP = 0.0f`, greedy search is performed. - * If setting `topK > 0` and `topP = 0.0f`, `topK` tokens of highest probabilities will become the candidates of sampling (named `TopK sampling` in TRT-LLM). - * If setting `topK = 0` and `topP > 0.0f`, tokens will be sorted with probability descendly, then the tokens with highest probabilities which the accumulated probability larger than `topP` will become the candidates of sampling (named `TopP sampling` in TRT-LLM). - * If setting `topK > 0` and `topP > 0.0f`, `topK` tokens of highest probabilities will be selected, then those selected tokens will be sorted with probability descendly and their probability will be normalized, then the tokens with highest normalized probabilities which the accumulated probability larger than `topP` will become the candidates of sampling (named `TopKTopP sampling` in TRT-LLM) - - * If different `topK` values are provided for the different sequences in the batch, the performance of the implementation will depend on the largest value. For efficiency reasons, we recommend to batch requests with similar `topK` values together. - - * `topPDecay`, `topPMin` and `topPResetIds` are explained in - [_Factuality Enhanced Language Models for Open-Ended Text Generation_](https://arxiv.org/abs/2206.04624). - `topPDecay` is the decay, `topPMin` is the lower-bound and `topPResetIds` indicates where to reset the decay. - - * `minP` is explained in [_Turning Up the Heat: Min-p Sampling for Creative and Coherent LLM Outputs_](https://arxiv.org/abs/2407.01082). - - * TensorRT-LLM does not generate all possible tokenizations of a word. Therefore, stop words may appear in the output if there are multiple ways to tokenize a stop word and the token sequence in the output differs from the one in `stopWords`. - -***Beam-search*** - -| Name in TRT-LLM | Description | Data type | Range of value | Default value | Name in HF | -| :-----------------------: | :-----------------------------: | :-----------------: | :----------------------: | :-----------------------: | :-----------------: | -| `beamWidth` | width for beam-search algorithm | Int | \[0, 1024\] | `0` (disable beam search) | `beam_width` | -| `beamSearchDiversityRate` | diversity of generated tokens | List\[Float\] | \[0, $+\infty$\) | `0.0f` | `diversity_penalty` | -| `lengthPenalty` | penalize longer sequences | List\[Float\] | \[0, $+\infty$\) | `0.0f` | `length_penalty` | -| `earlyStopping` | see description below | List\[Int\] | \($-\infty$, $+\infty$\) | `0` | `early_stopping` | -| `beamWidthArray` | see description below | List\[List\[Int\]\] | \[0, 1024\] | `` | no | - - * Beam-search algorithm: [beam search](https://en.wikipedia.org/wiki/Beam_search). - * Parameter `diversity_penalty` in HF is only used for `diverse beam-search decoding` (or named `Group-Beam-Search`), which is not supported by TRT-LLM yet. - * If setting `earlyStopping = 1`, decoding will stop once `beamWidth` finished sentences are generated. - * If setting `earlyStopping = 0`, decoding will keep going until no better sentences (with better score) can be generated. - * If setting `earlyStopping` to other values, decoding will stop only depending on `lengthlengthPenalty`. - * `beamWidthArray` is a list of beam width for each step. Using `beamWidthArray = [20,40,80]` as an example, -beam width will be 20 for the first step, 40 for second step, 80 for the later all steps. - * The `beamWidth` parameter is a scalar value. It means that in this release of -TensorRT-LLM, it is not possible to specify a different width for each input -sequence. This limitation is likely to be removed in a future release. - -### Internal Components - -The [`TllmRuntime`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/runtime/tllmRuntime.h) is in charge of the execution of the TensorRT engine. -The `TllmRuntime` class is an internal component and you are not expected to use that class directly. -The [`GptDecoder`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/include/tensorrt_llm/runtime/gptDecoder.h) generates tokens from the logits. -The `GptDecoder` can be used directly to implement a custom generation loop and for use cases that cannot be satisfied by the TRT-LLM implementation. diff --git a/docs/source/legacy/advanced/graph-rewriting.md b/docs/source/legacy/advanced/graph-rewriting.md deleted file mode 100644 index 40552a0497d9..000000000000 --- a/docs/source/legacy/advanced/graph-rewriting.md +++ /dev/null @@ -1,197 +0,0 @@ -(graph-rewriting)= - - -# Graph Rewriting Module - -TensorRT-LLM uses a declarative approach to define neural networks and contains -techniques to optimize the underlying graph. It provides a wrapper similar to PyTorch's Module. When a user invokes the `forward` method, the layers are lowered to TensorRT's `ILayer`s and become part of an `INetworkDefinition`. The Graph Rewriting (GW) module can be used to manipulate the network at the `ILayer`/`INetworkDefinition` level. - -## When to Use Graph Rewriting? - -For network manipulation, there are two options in TensorRT-LLM: - -1. **Module Rewriting:** This method modifies the members of `Module` instances before triggering the `forward` method (that is, creating the TensorRT graph). It works on the highest level of the network representation and facilitates the modification of sequences of operations (like modifying the GEMM + activation for SmoothQuant), - -2. **Graph Rewriting:** Graph Rewriting manipulates TensorRT's `INetworkDefinition` after the `forward` method is triggered. It operates at a finer-grained `ILayer` level and can alter the structure across multiple Module instances. It is typically used for layer fusion. - -Graph Rewriting (GW) is ideally used in the following conditions: - -1. When only `ILayer`/`INetworkDefinition` is available, -2. When Module Rewriting would lead to nested control flow or scattered functionality. - -## Graph Rewriting APIs - -Several core APIs are provided for Graph Rewriting: - -### Tensor-Related Methods - -- `Tensor.get_parent`: Get the `ILayer` that produces this tensor, -- `Tensor.get_users`: Get the consumer `ILayer`s of this tensor, -- `replace_all_uses_with`: Replace this tensor with another tensor in all consumer `ILayer`s. - -### FLayerInfo for Retrieving High-Level Information for a Functional - -For all the layers located in `functional.py`, the original input information is missing once lowered to `INetworkDefinition`, especially for TensorRT plugins, which are opaque in the Python world. `FLayerInfo` holds their original information as a high-level signature containing inputs like `Tensor`, Python attributes, and more. There is a Network-wise singleton called `FLayerInfoMemo` to map each `ILayer` to its corresponding `FLayerInfo`. - -For `FLayerInfo`: - -- `FLayerInfo.replace_input_with`: Replace some input tensor with another tensor, -- `FLayerInfo.replace_output_uses_with`: Redirect the usage of the original output tensors to a set of new tensors. - -For `FLayerInfoMemo`: - -- `FLayerInfoMemo.instance()`: Get the singleton instance, -- `FLayerInfoMemo.get`: Get the corresponding `FLayerInfo` for an `ILayer`. - -`FLayerInfo` remains consistent with the actual `ILayer` during GW, making it safe to use. - -### Pattern and Pattern Manager - -There are two kinds of patterns: - -- `PatternRewriter`: Used for defining a rewriting pattern, which actually alters the network. - - `match`: Match the pattern; returns true if a layer is matched, - - `rewrite`: Manipulate a layer, - - `match_and_rewrite`: Combines both `match` and `rewrite`, used for complex states that need to pass from `match` to `rewrite`. - -- `PatternAnalyzer`: Used for defining an analysis pattern, which collects information from the network. - - `match`: Match the pattern, - - `analyze`: Perform analysis on a list of layers. - -There are two managers for managing multiple `PatternRewriter` or `PatternAnalyzer`: - -- `RewritePatternManager`: - - `add`: Add a pattern with its label and benefit; the benefit specifies its privilege, - - `get`: Get a pattern by label, - - `rewrite`: Apply the rewriting patterns contained to a network. - -- `AnalysisPatternManager`: - - `add`: Add a pattern with its label and benefit; the benefit specifies its privilege, - - `get`: Get a pattern by label, - - `analyze`: Apply the analysis patterns contained to a network. - -### @record_signature to Decorate Functionals Requiring FLayerInfo - -The `@record_signature` decorator is used to record the `FLayerInfo` for a functional. While FLayerInfo is vital for GW when analyzing or rewriting certain functionals, it is used in an "add as needed" manner. If you are adding GW patterns, ensure that the functional requires the `@record_signature` decorator. - -## Classical Workflow - -There are specific routines for defining a GW pattern. Let's start with a simple example: replacing a sum layer with a subtract layer, which can also be found in the `test_graph_rewriting.py` file. - -```python -class NaivePatternRewriter_ReplaceAddWithSub(PatternRewriter): - - def __init__(self): - super().__init__('replace_add_with_sub', - root_layer={trt.LayerType.ELEMENTWISE}, - separate_match_rewrite=True) - - def match(self, layer: Layer): - # The rewriter will stop at the first matched layer, and then the Rewriter will enter the rewrite() to do the rewriting. - return layer.as_layer().op == trt.ElementWiseOperation.SUM - - def rewrite(self, layer: Layer) -> None: - # The layer here should be an Elementwise_SUM layer. - with net_guard(layer.network): - # There are several stages to replace some subgraph with another subgraph: - - # Stage 1: Get the input tensors and output tensors of the subgraph to replace. - # - For Elementwise_SUM, there are two inputs and one output. - a, b = layer.get_inputs(0, 1) - o = layer.get_outputs(0)[0] - - # Stage 2: Create a new subgraph that takes the old one's inputs. - # - Here we insert an Elementwise_SUB layer, and 'c' is the output. - c = a - b - - # Stage 3: Redirect all the layers depending on the outputs of the old subgraph to the new subgraph's. - # - After this, the SUM becomes dangling and will be pruned by TensorRT when building the engine. - # - Note that there is no API in TensorRT python to remove a layer explicitly; `replace_all_uses_with` is the only way to "remove" a layer. - o.replace_all_uses_with(c) - - # Stage 4: Mark all the layers in the old subgraph as removed. - # - This helps the PatternRewriter to skip the removed layers. - layer.mark_as_removed() -``` - -In this example, we deal with `ILayer` rather than Plugins, so `FLayerInfo` is unnecessary. As illustrated in the `rewrite` method, there are four stages that are shared across nearly all rewrite patterns. - -Note that in GW, we **NEVER** rewrite a layer directly. Instead, we do it in two steps: first, create another layer with the same input and deprive all the users of the original outputs, redirecting them to the outputs of the new layers. In this way, the old layer will be dangling and pruned automatically by TensorRT during the engine building phase. This is a limitation of TensorRT since remove-layer-like APIs are not available in Python. - -In Stage 2, we rely on operators and layers commonly used during the network building phase. Ideally, you can replace them with any network structure during GW. - -For the usage of `FLayerInfo`, let's rewrite the `gpt_attention` to enable the `remove-padding` feature. `gpt_attention` is actually - - a TensorRT plugin, so we need `FLayerInfo` to hold the original Tensor-wise inputs to help create new `gpt_attention` layers. - -```python -class GPTAttentionPluginRemovePaddingRewritePass(PatternRewriter): - - def __init__(self): - super().__init__('gpt_attention_plugin_remove_padding', - root_layer={trt.LayerType.PLUGIN_V2}) - - def match_and_rewrite(self, layer: Layer) -> bool: - if layer.as_layer().type != trt.LayerType.PLUGIN_V2 or \ - layer.as_layer().plugin.plugin_namespace != 'tensorrt_llm' or \ - layer.as_layer().plugin.plugin_type != 'GPTAttention': - return False - - # Retrieve the FLayerInfo - flayer = FLayerInfoMemo.instance().get(layer.name) - assert flayer - # Although the layer is a plugin, which is a black box, we get some high-level input information from the FLayerInfo. - tensor_input: Tensor = flayer.get_input('qkv') - if tensor_input.shape[0] == 1: # Already in remove-padding mode - return False - - # Some information could be passed in from external - assert self.args is not None, "args should be passed in from RewritePatternManager.rewrite()" - batch_size, in_len, hidden_size = self.args['batch_size'], self.args['in_len'], self.args['hidden_size'] - - with net_guard(layer.network): - new_inputs = flayer.clone_inputs() - - # Step 1: Create new inputs and replace the original arglist. - input = Tensor( - name='qkv', - dtype=trt.float16, - shape=(1, batch_size * in_len, hidden_size), - ) - new_inputs['qkv'] = input - - # Step 2: Create a new plugin instance. - new_outs = gpt_attention(**new_inputs) - - # Step 3: Deprive all the users of the old plugin instance. - flayer.replace_outputs_uses_with(layer.network, new_outs) - - # Step 4: Remove the old plugin instance. - layer.mark_as_removed() - - return True -``` - -This is quite similar to the first example, with the focus on the `FLayerInfo` part. Through the code below, we can get the original inputs of this layer, enabling us to alter the inputs related to remove-padding and create a new layer to replace it. - -```python -flayer = FLayerInfoMemo.instance().get(layer.name) -assert flayer -``` - -```python -new_inputs = flayer.clone_inputs() - -# Step 1: Create new inputs and replace the original arglist. -input = Tensor( - name='tensor', - dtype=trt.float16, - shape=(1, batch_size * in_len, hidden_size), -) -new_inputs['tensor'] = input - -# Step 2: Create a new plugin instance. -new_outs = gpt_attention(**new_inputs) -``` - -For real examples, please refer to the `FuseAttentionWithBiasPass` in the `graph_rewriting.py`. diff --git a/docs/source/legacy/advanced/images/disaggregated-service_usage.png b/docs/source/legacy/advanced/images/disaggregated-service_usage.png deleted file mode 100644 index 6b98a2233227..000000000000 Binary files a/docs/source/legacy/advanced/images/disaggregated-service_usage.png and /dev/null differ diff --git a/docs/source/legacy/advanced/kv-cache-management.md b/docs/source/legacy/advanced/kv-cache-management.md deleted file mode 100644 index f4506d6ee996..000000000000 --- a/docs/source/legacy/advanced/kv-cache-management.md +++ /dev/null @@ -1,75 +0,0 @@ -(kv-cache-management)= - -# KV Cache Management: Pools, Blocks, and Events - -This document provides an overview of the internal hierarchy and event system for paged KV cache management, as implemented in the TensorRT-LLM codebase. - -For more information on KV cache reuse see [KV cache reuse](kv-cache-reuse.md). - ---- - -## Hierarchy: Pool, Block, and Page - -### **Block** -- **Definition:** The smallest unit of KV cache allocation. A `KVCacheBlock` holds metadata (not the actual data) for a chunk of KV cache. -- **Purpose:** Each block represents a fixed number of tokens' worth of KV data (can be specified by `tokens_per_block` parameter). -- **Usage:** Blocks are allocated, reused, or evicted as sequences are processed. - -### **Page** -- **Definition:** In this codebase, "page" is often used interchangeably with "block" (as in "paged KV cache"), but technically, a page could refer to a memory page (hardware-level), while a block is a logical unit for the cache. -- **In Practice:** The code uses "block" as the main unit; "page" is not a distinct class or struct. - -### **Pool** -- **Definition:** A pool is a contiguous memory buffer (or set of buffers) that holds the actual KV data for one or more layers. -- **Types:** There are primary pools (fast GPU memory) and secondary pools (slower, e.g., CPU or offload memory). -- **Organization:** Each pool can serve multiple layers that share the same KV head configuration. Pools are managed by `KVCacheBlockPool` and tracked in vectors in `WindowBlockManager`. -- **Block ↔ Pool:** Each block is an index into a pool; the pool provides the actual storage, while the block is the metadata handle. - -### **WindowBlockManager/BlockManager** - -TRT-LLM supports 2 complex features related to KV cache management: -1. **Variable Group-Query Attention (VGQA)** - i.e. a different `num_kv_heads` value for different layers. -2. **Variable Sliding Window Attention (VSWA)** - i.e. a different `attention_window_size` value for different layers. - -In order to support both of these features, the pool management works as described below. - -But in the simple, *most common case*, for most models, where -1. [MHA/MQA/Non-variable GQA](gpt-attention.md#multi-head-multi-query-and-group-query-attention), i.e., same `num_kv_heads` value for all layers, -2. Global attention/[SWA](gpt-attention.md#sliding-window-attention-cyclic-rolling-buffer-kv-cache), i.e., same `attention_window_size` value for all layers, - -only a *single* pool will be created within the structure described below. - -#### KV Cache Pool Management - -- **WindowBlockManager:** Manages blocks and pools for a specific attention window size. Within a `WindowBlockManager`, there can be multiple pools - each corresponding a unique number of KV heads - i.e., to support VGQA. -- **BlockManager:** Manages all `WindowBlockManager` instances, one per unique window size. - -**Hierarchy Summary:** -- **Pool** (memory buffer for KV data) - - Contains many blocks. -- **Blocks** (metadata for a chunk of the pool, each block = `tokens_per_block` tokens) - - (Optionally, blocks can be swapped between primary/secondary pools.) -- **BlockManager/WindowBlockManager**: Manage pools and blocks, handle allocation, reuse, and eviction. - ---- - -## Events in `KVCacheEventManager` - -The `KVCacheEventManager` is responsible for tracking and reporting significant changes in the state of the KV cache. Events are used for logging, debugging, or possibly for external monitoring. - -### **Types of Events** -- **Created Event:** When pools or blocks are created/allocated. -- **Updated Event:** When a block's state changes (e.g., moved between primary/secondary, priority updated). -- **Removed Event:** When a block is removed from the cache (evicted or released). -- **Stored Event:** When blocks are stored for potential reuse (e.g., after a sequence finishes and its blocks are reusable). - -### **What Triggers an Event?** -- **Allocation/Deallocation:** Creating or freeing memory pools or blocks. -- **Eviction/Reuse:** When a block is evicted, reused, or its priority changes. -- **Block Movement:** When a block is moved between memory levels (primary ↔ secondary). -- **Block Storage:** When blocks are stored for future reuse (e.g., after a sequence completes). - -**In summary:** -An "event" is any significant change in the lifecycle or state of a KV cache block or pool, tracked for monitoring, debugging, or optimization purposes. - ---- diff --git a/docs/source/legacy/advanced/kv-cache-reuse.md b/docs/source/legacy/advanced/kv-cache-reuse.md deleted file mode 100644 index 5f3a5d73cf34..000000000000 --- a/docs/source/legacy/advanced/kv-cache-reuse.md +++ /dev/null @@ -1,94 +0,0 @@ -(kv-cache-reuse)= - -# KV cache reuse - -This document describes how kv cache pages can be shared and reused by requests that start with the same prompt. This can greatly lower first token latency, the time it takes before the first output token is generated. Many use cases can benefit from this, including multi-turn requests and system prompts. - -## How to enable kv cache reuse - -There are two steps to enabling kv cache reuse. - -1. Model must support it - -KV cache reuse requires the model to be built for paged context attention. This is done with `trtllm-build`: - -```trtllm-build --use_paged_context_fmha enable``` - -2. KV cache reuse is enabled by default in KVCacheManager - -If you are running gptManagerBenchmark application, you can disable kv cache reuse with a command-line switch: - -```gptManagerBenchmark --enable_kv_cache_reuse enable=false``` - -If you are running a Triton server, you can enable kv cache reuse with a parameter: - -``` -parameters: { - key: "enable_kv_cache_reuse" - value: { - string_value: "true" - } -} -``` - -If you are writing your own application using Executor API, you can enable kv cache reuse by including `enableBlockReuse=true` when you create the `KvCacheConfig` object. Note that this is the default, if you wish to disable kv cache reuse, pass `enableBlockReuse=false` instead. - -### Enable kv cache reuse for p-tuning - -When using p-tuning, different requests may use same fake input ids (i.e. prompt ids whose values are larger than vocabulary size). That may lead to incorrect kv cache reuse, since TRT-LLM could not distinguish these requests only by input ids. To enable kv cache reuse for p-tuning correctly, users should provide an extra id (uint64) for each input id. Extra ids for normal input ids (i.e. text token ids) should always be 0, while fake input ids should have extra ids which are larger than 0. Requests using same prompt embeddings should use same extra ids, while requests using different prompt embeddings should use different extra ids. - -Example: -Assume vocabulary size is 100, which means normal text token ids are in range [0, 99] and prompt ids start from 100. - -```python -# Request 1 uses prompt embedding table 1 -input_ids = [100, 101, 102, 103, 1, 2, 3, 4] -extra_ids = [1, 1, 1, 1, 0, 0, 0, 0] - -# Request 2 uses prompt embedding table 2 -input_ids = [100, 101, 102, 103, 1, 2, 3, 4] -extra_ids = [2, 2, 2, 2, 0, 0, 0, 0] - -# Request 3 uses prompt embedding table 1 and different text tokens -input_ids = [100, 101, 102, 103, 5, 6, 7, 8] -extra_ids = [1, 1, 1, 1, 0, 0, 0, 0] -``` - -## Performance expectations - -KV cache state can be reused when two requests start with the same partial prompt. This reduces first token latency, the time it takes until the first output token is generated. Bigger savings are realized when the shared prompt is longer, relative to the overall prompt length. The biggest saving is realized when two identical requests are run back-to-back, in which case the latency for the first output token approaches latency for subsequent tokens. - -## Situations that can prevent kv cache reuse - -There are a few pitfalls that can prevent kv cache reuse when that seems possible. KV cache state only becomes reusable after the request that computed the state terminates. If you have a shared system prompt, the first request will compute kv cache state for the system prompt, the second request will reuse it, but only if the second request launches after the first request completed. If you run with a large batch-size, it is likely that many requests that share a common system prompt will be launched before the first request has terminated. No reuse will occur until one of the requests terminate, then subsequently scheduled requests can reuse. - -Kv cache state for system prompts will remain reusable until memory is needed for launching a new request or propagating an existing one. When this happens, reusable blocks are evicted based on LRU. System prompts that are frequently used have a better chance of remaining reusable, but there is no guarantee since launching new requests take priority over possible reuse. Running with a larger batch size, or larger output sequence lengths for example will reduce the probability of kv cache blocks being reused, since it increases memory needs. - -KV cache state is stored in blocks, each block holds multiple tokens. Only full blocks can be shared by multiple requests, thus the block size matters. Partially matched blocks can also be reused, but that creates a new copy of the block for each sequence. The block size is a trade-off, larger block size may improve efficiency of compute kernels, but it reduces the likelihood of kv cache state reuse. The block defaults to 128 tokens, this can be changed when the model is built with the trtllm-build command, for example - -```trtllm-build --tokens_per_block 32 ...``` - -will create a model where one KV cache block can hold 32 tokens. Note that tokens_per_block must be a power of 2. - -## Offloading to host memory - -Offloading to host memory increases likelihood of kv cache reuse. Reusable blocks that are needed for higher priority tasks, like propagating an already running request, are copied to a buffer in host memory instead of being evicted. This greatly extends the amount of memory available for reuse, allowing blocks to remain reusable much longer. On the other hand, offloading of blocks (and subsequent onboarding when a block is reused) has some cost since the blocks must be copied from CPU to GPU memory and vice versa. This cost is negligible on Grace-Hopper machines, and small enough to yield a net benefit for many use cases on x86 machines with Hopper GPUs. Offloading is unlikely to yield benefits on older architectures because of the (relatively) slow link between GPU and host memory. - -If you are running gptManagerBenchmark, you can enable offloading with a command-line switch. For example, - -```gptManagerBenchmark --kv_host_cache_bytes 45000000000``` - -will create a 45 GiB offloading buffer in host memory. Note that this buffer is pinned memory, allocating a lot of pinned memory on x86 machines can take a substantial amount of time (10s of seconds). This is a one-time cost. - -If you are running a Triton server, you can enable offloading to host memory with the kv_cache_host_memory_bytes parameter. For example, adding this to your model config file will create a 45 GiB offloading buffer in host memory. - -``` -parameters: { - key: "kv_cache_host_memory_bytes" - value: { - string_value: "45000000000" - } -} -``` - -If you are writing your own application using Executor API, you can enable offloading to host by including `hostCacheSize=45000000000` when you create the `KvCacheConfig` object. This will create a 45 GiB offloading buffer in host memory. diff --git a/docs/source/legacy/advanced/lora.md b/docs/source/legacy/advanced/lora.md deleted file mode 100644 index 5ac3aa7c5797..000000000000 --- a/docs/source/legacy/advanced/lora.md +++ /dev/null @@ -1,142 +0,0 @@ -(lora)= - -## Run gpt-2b + LoRA using Executor / cpp runtime - -First build a model with LoRA and inflight-batching enabled. - -```bash -git-lfs clone https://huggingface.co/qychen/luotuo-lora-7b-0.1 -git-lfs clone https://huggingface.co/kunishou/Japanese-Alpaca-LoRA-7b-v0 -BASE_MODEL=llama-7b-hf - -python examples/models/core/llama/convert_checkpoint.py --model_dir ${BASE_MODEL} \ - --output_dir /tmp/llama_7b/trt_ckpt/fp16/1-gpu/ \ - --dtype float16 - -trtllm-build --checkpoint_dir /tmp/llama_7b/trt_ckpt/fp16/1-gpu/ \ - --output_dir /tmp/llama_7b_with_lora_qkv/trt_engines/fp16/1-gpu/ \ - --remove_input_padding enable \ - --gpt_attention_plugin float16 \ - --context_fmha enable \ - --paged_kv_cache enable \ - --gemm_plugin float16 \ - --lora_plugin float16 \ - --max_batch_size 128 \ - --max_input_len 512 \ - --max_seq_len 562 \ - --lora_dir Japanese-Alpaca-LoRA-7b-v0 \ - --max_lora_rank 8 \ - --lora_target_modules "attn_q" "attn_k" "attn_v" -``` - -To pass LoRAs into the cpp runtime they must be converted to the format below. -The script below will convert a Hugging Face LoRA model to the correct NumPy tensor. - -```bash -python3 tensorrt_llm/examples/hf_lora_convert.py -i Japanese-Alpaca-LoRA-7b-v0 -o Japanese-Alpaca-LoRA-7b-v0-weights --storage-type float16 -python3 tensorrt_llm/examples/hf_lora_convert.py -i luotuo-lora-7b-0.1 -o luotuo-lora-7b-0.1-weights --storage-type float16 -``` - -Refer to the [tensorrtllm_backend documentation](https://github.com/triton-inference-server/tensorrtllm_backend/blob/main/docs/lora.md) for a Multi-LoRA example using Triton. - -### LoRA tensor format details - -To run inference using `Executor`, a `Request` must have a `LoraConfig` that contains a `task_id`, `weights` and `config` parameters. - -`task_id` the unique task ID for the given LoRA. - -To perform inference with a specific LoRA for the first time, `task_id`, `weights`, and `config` must all be given. The LoRA will be cached, so that subsequent requests for the same task only require `task_id`. -If the cache is full, the oldest LoRA will be evicted to make space for new ones. An error is returned if `task_id` is not cached. - -`weights` contains the weights for all the LoRAs. Currently, this should include weights for all TP and PP ranks. -The weights tensor has the shape `[num_lora_modules_layers, D x Hi + Ho x D ]`. The last dimension holds the in / out adapter weights for the associated module (for example, `attn_qkv`) and model layer. - -Each of the in / out tensors are first flattened and then concatenated together in the format above. -The first dimension (of size `num_lora_module_layers`) has an entry for each module-layer (that is, there is an entry for `attn_q layer1` and another for `attn_k layer1`). - -`D=adapter_size (i.e. R value), Hi=hidden_size_in, Ho=hidden_size_out.` - -`config` is a configuration tensor which identifies the moduleId, layerId, and adapter size of each element of `LoraWeights`. It has the shape `[num_lora_modules_layers, 3]`. The last dimension holds `[module_id, layer_idx, adapter_size D (i.e. R value)]`. - -This feature supports LoRAs as described in https://arxiv.org/pdf/2106.09685.pdf - -#### Example LoRA tensors - -Here is an example of `LoraWeights` and `LoraConfig` tensors for a model with tp=1, pp=1, 4 layers, and a hidden size of 4. -The following tensors are for a LoRA which has a `q` and `k` adapter. - -``` -# loraConfig -[ - [1, 0, 2] - [2, 0, 4] - [1, 1, 2] - [2, 1, 4] - [1, 2, 2] # Note that the final 2 layers only adapt `q` - [1, 3, 8] -] -# Note: The loraConfig tensor configures the loraWeights tensor. -# The contents of each row of loraWeights is specified be the corresponding row in loraConfig - -# loraWeights -# Note: that 'in weights' and 'out weights' are 'A' and 'B' in the LoRA paper. -[ - [ <2 x 4 in weights>, <4 x 2 out weights> ] # `q` adapter for layer 0 - [ <4 x 4 in weights>, <4 x 4 out weights> ] # `k` adapter for layer 0 - [ <2 x 4 in weights>, <4 x 2 out weights> ] # `q` adapter for layer 1 - [ <4 x 4 in weights>, <4 x 4 out weights> ] # `k` adapter for layer 1 - [ <2 x 4 in weights>, <4 x 2 out weights> ] # `q` adapter for layer 2 - [ <8 x 4 in weights>, <4 x 8 out weights> ] # `q` adapter for layer 3. Note the final layer has a adapter size of 8 -] - -``` - -#### LoRA Module id mapping - -| module name (as specified in `convert_checkpoint.py` scripts) | module id | description | -| --------------------------------------------- | --------- | ----------- | -| attn_qkv | 0 | compbined qkv adapter | -| attn_q | 1 | q adapter | -| attn_k | 2 | k adapter | -| attn_v | 3 | v adapter | -| attn_dense | 4 | adapter for the dense layer in attention | -| mlp_h_to_4h | 5 | for llama2 adapter for gated mlp layer after attention / RMSNorm: up projection | -| mlp_4h_to_h | 6 | for llama2 adapter for gated mlp layer after attention / RMSNorm: down projection | -| mlp_gate | 7 | for llama2 adapter for gated mlp later after attention / RMSNorm: gate | -| cross_attn_qkv | 8 | compbined qkv adapter for cross attention | -| cross_attn_q | 9 | q adapter for cross attention | -| cross_attn_k | 10 | k adapter for cross attention | -| cross_attn_v | 11 | v adapter for cross attention | -| cross_attn_dense | 12 | adapter for the dense layer in cross attention | -| moe_h_to_4h | 13 | for mixtral adapter for expert mlp layer: up projection | -| moe_4h_to_h | 14 | for mixtral adapter for expert mlp layer: down projection | -| moe_gate | 15 | for mixtral adapter for expert mlp layer: gate | -| moe_router | 16 | for mixtral adapter for expert router layer | -| mlp_router | 17 | for qwen2-moe adapter for shared expert gate layer | -| mlp_gate_up | 18 | adapter for gated mlp layer after attention / RMSNorm: gate + up projection | - -#### LoraCache configuration - -The core idea is that we will have a fixed size, 2-level LoRA cache in TRT-LLM. The higher level cache resides on the host and the lower level is on GPU (distinct from the existing KV cache). Sizes of both are user configurable. - -The CPU cache is configured to be a max size. The GPU cache is configured to a percentage of free GPU memory after engine load. As requests come in LoRAs are stored in the host cache. - -As requests are scheduled for execution LoRAs are loaded into the GPU cache. - -#### LoRA with tensor parallel - -The partition of tensor parallel for LoRA is special. There are two cases: `RowLinear` and `ColumnLinear`. Assume we have a linear layer and the input feature size is `K` and the output feature size is `N`. Then, the shape of the weight is `[K, N]`. - -First, consider this linear layer is a `ColumnLinear` layer. When we partition the weight, we split the weight by column with `tp_size`. Then, there are `tp_size` split weights and the shapes of these weights are `[K, N // tp_size]`. When we apply LoRA adapter on such `ColumnLinear` layer, the shapes of original two weights are `[K, lora_rank]` and `[lora_rank, N]`. So, we only partition the second weight and get `tp_size` split weights with shapes `[lora_rank, N // tp_size]`. For the first weight, each GPU maintains the same entire weight (with shape `[K, lora_rank]`). - -Next, consider this linear layer is a `RowLinear` layer. When we partition the weight, we split the weight by row with `tp_size`. Then, there are `tp_size` split weights and the shapes of these weights are `[K // tp_size, N]`. When we apply LoRA adapter on such `RowLinear` layer, the shapes of original two weights are `[K, lora_rank]` and `[lora_rank, N]`. So, we only partition the first weight and get `tp_size` split weights with shapes `[K // tp_size, lora_rank]`. For the second weight, each GPU maintains the same entire weight (with shape `[lora_rank, N]`). - -#### DoRA - -TensorRT LLM supports DoRA as described in https://arxiv.org/abs/2402.09353 . To enable DoRA, you must add the additional `--dora_plugin enable` flag to the `trtllm-build` command. - -The DoRA scales must be normalized before they are submitted to TensorRT LLM in an inference request. The normalization requires the base model weights. To normalize your adapter you may use the script provided in `tensorrt_llm/examples/dora/normalize_weights.py`. - -When using DoRA, the format of `LoraWeights` and `LoraConfig` changes slightly. -The shape of `LoraConfig` becomes `[module_id, layer_idx, adapter_size D (i.e. R value), is_dora]`, with `is_dora` a boolean flag that determines whether the supplied adapter contains DoRA scales or not. If the old config shape is used, it is assumed the adapter does not have DoRA scales. -The shape of `LoraWeights` becomes `[num_lora_modules_layers, D x Hi + Ho x D + Ho]`, and the last `Ho` values are the DoRA scale vector. diff --git a/docs/source/legacy/advanced/open-sourced-cutlass-kernels.md b/docs/source/legacy/advanced/open-sourced-cutlass-kernels.md deleted file mode 100644 index 80e9c066b642..000000000000 --- a/docs/source/legacy/advanced/open-sourced-cutlass-kernels.md +++ /dev/null @@ -1,26 +0,0 @@ -We have recently open-sourced a set of Cutlass kernels that were previously known as "internal_cutlass_kernels". Due to internal dependencies, these kernels were previously only available to users as static libraries. We have now decoupled these internal dependencies, making the kernels available as source code. - -The open-sourced Cutlass kernels are on the path `cpp/tensorrt_llm/kernels/cutlass_kernels`, including: -- `low_latency_gemm` -- `moe_gemm` -- `fp4_gemm` -- `allreduce_gemm` - -To ensure stability and provide an optimized performance experience, we have maintained the previous method of calling these kernels via static libraries as an alternative option. You can switch between open-sourced Cutlass kernels and static library Cutlass kernels through the `USING_OSS_CUTLASS_*` macro (where * represents the specific kernel name), enabling kernel-level control. By default, the open-source Cutlass kernels are used. -Note that support for these static libraries will be gradually deprioritized in the future and may eventually be deprecated. - -**Default Configuration (Using open-sourced Cutlass Kernels)** - -To build using the open-source Cutlass kernels (default setting), run: - -```bash -python3 ./scripts/build_wheel.py --cuda_architectures "90-real;100-real" -``` - -**Using Static Library Cutlass Kernels** - -If you prefer to use the Cutlass kernels from the static library, you can control this during compilation by setting the `USING_OSS_CUTLASS_*` macro to `OFF`. For example, to use the static library implementation specifically for `low_latency_gemm` and `moe_gemm` while keeping other kernels as OSS, use the following compilation command: - -```bash -python3 ./scripts/build_wheel.py --cuda_architectures "90-real;100-real" -D "USING_OSS_CUTLASS_MOE_GEMM=OFF;USING_OSS_CUTLASS_LOW_LATENCY_GEMM=OFF" -``` diff --git a/docs/source/legacy/advanced/speculative-decoding.md b/docs/source/legacy/advanced/speculative-decoding.md deleted file mode 100644 index 2faf885f8d10..000000000000 --- a/docs/source/legacy/advanced/speculative-decoding.md +++ /dev/null @@ -1,182 +0,0 @@ -# Speculative Sampling - -- [About Speculative Sampling](#about-speculative-sampling) -- [Performance Improvements](#Performance-improvements) -- [Draft-Target-Model](#Draft-Target-Model) -- [NGram](#ngram) -- [Medusa](#medusa) - - [Medusa Tree](#medusa-tree) - - [Using Medusa with TensorRT-LLM](#using-medusa-with-tensorrt-llm) - - [Limitations](#limitations) -- [ReDrafter](#redrafter) -- [EAGLE](#eagle) - - [Disaggregated Serving](#disaggregated-serving) -- [Lookahead decoding](#lookahead-decoding) - -## About Speculative Sampling - -Speculative Sampling (also referred to as Speculative Decoding) is a set of techniques designed to allow generation of more than one token per forward pass iteration. This can lead to a reduction in the average per-token latency **in situations where the GPU -is underutilized due to small batch sizes.** - -Speculative Sampling involves predicting a sequence of future tokens, referred to as draft tokens, using a method -that is substantially more efficient than repeatedly executing the target Large Language Model (LLM). -These draft tokens are then collectively validated by processing them through the target LLM in a single forward pass. -The underlying assumptions are twofold: - -1. processing multiple draft tokens concurrently will be as rapid as processing a single token -2. multiple draft tokens will be validated successfully over the course of the full generation - -If the first assumption holds true, the latency of speculative decoding will no worse than the standard approach. If the second holds, output token generation advances by statistically more than one token per forward pass. -The combination of both these allows speculative decoding to result in reduced latency. - -TensorRT-LLM supports several approaches for generating draft tokens, including: - -1. Utilizing a smaller, auxiliary model, known as the draft model approach. For more information, refer to the [Fast Inference from Transformers via Speculative Decoding paper](https://arxiv.org/pdf/2211.17192.pdf). -2. Implementing additional language model heads that predict tokens for future positions: - 1. [Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads paper](https://arxiv.org/abs/2401.10774). - 2. [Recurrent Drafter for Fast Speculative Decoding in Large Language Models](https://arxiv.org/html/2403.09919v1). - 3. [EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty](https://arxiv.org/pdf/2401.15077). -3. Utilizing prompt tokens as draft tokens. For more information, refer to [NGram](https://github.com/apoorvumang/prompt-lookup-decoding/). -4. Utilizing Jacobi-like decoding to predict and verify draft tokens using the same model which does not need additional fine-tuning. Refer to [Break the Sequential Dependency of LLM Inference Using Lookahead Decoding](https://arxiv.org/pdf/2402.02057). - - -## Performance Improvements - -It's important to note that the effectiveness of speculative decoding techniques is highly dependent -on the specific task at hand. For instance, forecasting subsequent tokens in a code-completion scenario -may prove simpler than generating a summary for an article. - -Furthermore, when integrating Medusa with a standard PyTorch model implementation which may not be as finely -tuned as TensorRT-LLM, the potential time savings are more pronounced. - -## Draft-Target-Model - -The Draft-Target-Model involves the use of two distinct models (a smaller Draft model and a larger Target model) trained independently but sharing the same vocabulary. For example, GPT 125M / 6.7B models serve as the Draft / Target model. - -The management of Draft and Target models is facilitated through two separate `Executor` instances. -It is essential that you to coordinate the interactions between the Draft and Target models effectively. -Initially, the Draft model is queried to generate up to `K` draft tokens. -These tokens are then forwarded to the Target model for verification. -Upon verification, the Target model may return up to `K+1` tokens. -Subsequently, the prompt, now updated with the accepted tokens, is sent back to the Draft model to initiate the generation of new draft tokens. -This iterative process continues until a predefined stop conditions are met. -An example orchestration script is available in the Triton backend repository’s -[draft-target-model client example](https://github.com/triton-inference-server/tensorrtllm_backend/blob/main/client/python/draft_target_model_client.py). - -We provide two styles of running Draft-Target-Model now: using TensorRT-LLM-BLS in Triton Inference Server, or using TensorRT-LLM directly. Detailed steps of running can be found in [examples/draft_target_model/README.md](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/draft_target_model/README.md) and the code can be found in [examples/ngram/run_dtm_ngram.py](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/ngram/run_dtm_ngram.py). - -## NGram - -The NGram speculative decoding directly copies from the input prompt and previous generated output as draft tokens while generating the later output. It works like Draft-Target-Model but involves only one Target LLM model without further fine-tuning. The NGram profit from the scenarios which have high n-gram overlap between input prompt and output, such as summarization, document QA, multi-turn chat, code editing, etc. - -See document in [examples/ngram/README.md](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/ngram/README.md) and the code can be found in [examples/ngram/run_dtm_ngram.py](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/ngram/run_dtm_ngram.py). - -## Medusa - -This approach leverages a single model to both generate and verify draft tokens. -It enhances the existing model by adding multiple extra language model heads, known as Medusa heads. -These additional heads are trained to predict future tokens while the base model remains unchanged. -Specifically, the first Medusa head is tasked with predicting the immediate next token, -the second head predicts the token after that, and so on. -With `K` Medusa heads, the model can forecast up to `K` tokens ahead. -The draft tokens generated by the Medusa heads during iteration `i` -are then verified and potentially accepted in the subsequent iteration, `i+1`. - -The true potential of the Medusa strategy is realized when more than one token per head is used, -employing a TopK approach to create multiple potential paths, essentially forming a tree, rather than -a single linear path as seen in the Draft model approach. To reduce redundant computations, many of these paths, -which often share common prefixes, are consolidated into a single path. -This is achieved by applying attention with a sparse mask that represents the various paths. Sparse mask formed by Medusa tree is described in detail later. - -By validating multiple paths simultaneously, there is an increased likelihood of accepting more than one token per iteration, -albeit at the expense of additional computational effort. - -It is crucial to recognize that as the number of potential paths grows exponentially with `K`, -it is not necessary to explore or validate all of them. A recommended strategy for managing this complexity is to prune the tree -by focusing only on the paths with higher-probability tokens. - -You must strike a balance between the breadth and depth of the tree you want to explore and the impact of a larger tree on the overall -performance for your specific application. - -In the TensorRT-LLM implementation of Medusa, the configuration of the tree is a runtime parameter. -This flexibility allows you to experiment and identify the optimal tree structure for your use case, -which can then be utilized in a production environment. - -### Medusa Tree - -Consider the following diagram, which illustrates how the hidden states from the last layer of the base model -are passed to the base model's language model (LM) head and to four Medusa heads (MHs). - -

- Example Medusa Tree -

- -In this example: - -1. The token l0 represents the actual token generated by the model. -All other tokens, denoted as phk, are predictions from the MHs, -where `h` indicates the Medusa head index (1-based) and `k` represents the TopK choice index (0-based). -1. Four MHs are used, which means the model is predicting four future tokens. -2. The first two MHs utilize Top-2 predictions, while the last two use Top-1. -For instance, p10 and p11 are the top and -second top predictions from the first Medusa Head (MH1). -1. A total of four paths are explored, which is fewer than the 16 that would be examined -if a complete binary tree were used (assuming Top-2 predictions for all MHs). -1. As some of these paths may be accepted, there are ten potential candidates, referred to as `medusa_choices`. -The number of tokens that can be accepted at each step, including the true token, -ranges from 1 (if all Medusa predictions are incorrect) to 5 (if all are correct). - -During the generation phase, the model receives an input of 10 tokens, -which corresponds to the last tokens of each candidate path, rather than just one. - -In TensorRT-LLM, you have the option to define such trees by providing all the Medusa choices -or by simply specifying the unique paths. - -- Since each candidate/path begins with the true token (l0), -there is no need to specify it separately. For the predicted tokens, only the TopK indices are required. -- For example, to specify the path l0p10p21p30, -one would use `[0,1,0]`. And -to specify the path l0p11p20, -one would use `[1,0]`. -- To specify all 4 paths in the example, use `medusa_choices=[[0,0,0,0], [0,1,0], [1,0], [1,1]]`. -- It's also possible to specify all candidates explicitly, similar to the Medusa repository. -For instance, `medusa_choices=[[0], [0,0], [0,0,0], [0,0,0,0], [0,1], -[0,1,0], [1], [1,0], [1,1]]`. Note that when specifying all the candidates explicitly, **we don't include -the empty `[]` candidate** for the case where only the true token is accepted, that is, all the predictions from MHs are wrong. -So, only `9` candidates are specified. - -**Specifying paths-only instead of all choices is currently supported only in the Python runtime.** - -### Using Medusa with TensorRT-LLM - -For guidance on constructing and executing Medusa with the Python runtime, consult the [Medusa README](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/medusa/README.md). When utilizing the Inflight Fused Batching (IFB) with the C++ API, it is necessary to define the `medusa_choices` explicitly within the model configuration. For detailed instructions, refer to the [model configuration in TensorRT-LLM backend](https://github.com/triton-inference-server/tensorrtllm_backend?tab=readme-ov-file#modify-the-model-configuration) for more details. - -#### Limitations - -- TensorRT-LLM supports Medusa only for Vicuna (fine tuned LLaMA). -However, similar to any new model, you can follow the same approach to define your own Medusa model and deploy with TensorRT-LLM. -- We match only tokens during the validation phase that is `medusa_temperature=0`. -- Beam search is **not** compatible with Medusa. - - -## ReDrafter - -The ReDrafter approach enhances the single-model Medusa method by predicting and verifying tokens using the same model. However, unlike Medusa, it predicts draft tokens using a recurrent predictor, where each draft token depends on the previous one. This method also allows the use of beam search to identify more prominent draft tokens. For more details, please read [the ReDrafter paper](https://arxiv.org/html/2403.09919v1). - -TensorRT-LLM implements the ReDrafter model such that logits prediction, beam search, and draft token acceptance are performed inside the TensorRT engine. This contrasts with standard model inference, which only predicts logits and performs decoding outside the engine. Since the engine predicts explicit draft tokens instead of implicit tokens decoded from logits, we categorize this speculative decoding method as `explicit_draft_tokens`. Please, visit the [ReDrafter README](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/redrafter/README.md) for information about building and running the model. ReDrafter supports both Inflight Fused Batching runtime and Python static batching runtime. - -## EAGLE - -The EAGLE approach enhances the single-model Medusa method by predicting and verifying tokens using the same model. Similarly to ReDrafter, it predicts draft tokens using a recurrent predictor where each draft token depends on the previous one. However, unlike ReDrafter, it uses a single-layer transformer model to predict draft tokens from previous hidden states and decoded tokens. In the EAGLE-1 decoding tree needs to be known during the decoding. In the EAGLE-2 this tree is asssembled during the execution by searching for the most probable hypothesis along the beam. - -Similarly to ReDrafter, TensorRT-LLM implements the EAGLE model such that logits prediction, draft tokens acceptance and draft token generation are performed inside of the TensorRT engine(EAGLE-1 and EAGLE-2 are both supported). Please, visit the [EAGLE README](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/eagle/README.md) for information about building and running the model. - -> **EAGLE3 note.** If the EAGLE3 draft head config omits `draft_vocab_size`, TensorRT-LLM assumes it matches `vocab_size` and emits a warning. Set `draft_vocab_size` explicitly if the draft head uses a different vocabulary. - -### Disaggregated Serving - -[Disaggregated Serving](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/features/disaggregated-service.md) with EAGLE3 using the two model approach is supported in the Pytorch backend. Please refer to the following [Dynamo example](https://github.com/ai-dynamo/dynamo/blob/main/examples/tensorrt_llm/llama4_plus_eagle.md) on how to run EAGLE3 with Disaggregated Serving for Llama 4 Maverick. - -## Lookahead Decoding - -Lookahead decoding algorithm operates through two parallel computation branches within the same model: a lookahead branch that generates n-grams using a fixed-sized 2D window, and a verification branch that validates promising n-gram candidates. This approach eliminates the necessity for additional model training or fine-tuning and can be enabled for any autoregressive model. Refer to the [Lookahead decoding README](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/lookahead/README.md) for information about building and running the model. diff --git a/docs/source/legacy/advanced/weight-streaming.md b/docs/source/legacy/advanced/weight-streaming.md deleted file mode 100644 index 0cbe4b031f7b..000000000000 --- a/docs/source/legacy/advanced/weight-streaming.md +++ /dev/null @@ -1,58 +0,0 @@ -(weight-streaming)= - -## Running With Weight Streaming to Reduce GPU Memory Consumption - -TensorRT Weight Streaming can offload some weights to the CPU memory and stream them to the GPU memory during runtime. -This can reduce the weights size in GPU memory, therefore, we can run larger models or larger batch sizes in the same GPU memory budget. - - -During build time, build the engine with `--weight-streaming --gemm_plugin disable` since Weight Streaming only supports non-plugin weights. During runtime, run with `--gpu_weights_percent x` to config the percent of weights that remained on the GPU. `x` can be a value from `0.0` to `1.0`. - -Here is an example to run llama-7b with Weight Streaming: -```bash - -# Convert model as normal. Assume hugging face model is in llama-7b-hf/ -python3 examples/models/core/llama/convert_checkpoint.py \ - --model_dir llama-7b-hf/ \ - --output_dir /tmp/llama_7b/trt_ckpt/fp16/1-gpu/ \ - --dtype float16 - -# Build engine that enabled Weight Streaming. -trtllm-build \ - --checkpoint_dir /tmp/llama_7b/trt_ckpt/fp16/1-gpu/ \ - --output_dir /tmp/llama_7b/trt_engines/fp16/1-gpu/ \ - --weight_streaming \ - --gemm_plugin disable \ - --max_batch_size 128 \ - --max_input_len 512 \ - --max_seq_len 562 - -# Run the engine with 20% weights in GPU memory. -python3 examples/summarize.py \ - --engine_dir /tmp/llama_7b/trt_engines/fp16/1-gpu/ \ - --batch_size 1 \ - --test_trt_llm \ - --hf_model_dir llama-7b-hf/ \ - --data_type fp16 \ - --gpu_weights_percent 0.2 - -``` - - -### API Changes - -To build engines with Weight Streaming enabled, some API changes are needed for the builder: -- Added a new bool member `weight_streaming` to class `BuildConfig`. -- Added a new bool parameter `weight_streaming` to method `create_builder_config` of class `Builder`. - -To run with Weight Streaming with `Executor`, there are some API change to its config `ExecutorConfig`: -- Added a new float parameter `gpuWeightsPercent` to the constructor of `ExecutorConfig`. -- Added two member functions `setGpuWeightsPercent` and `getGpuWeightsPercent` to set and get the GPU weights percentage. - -Here is an example to create an `Executor` with Weight Streaming: -```c++ -... -auto executorConfig = tle::ExecutorConfig(gpuWeightsPercent=0.5); -auto executor = tle::Executor("model_path", tensorrt_llm::executor::ModelType::kDECODER_ONLY, executorConfig); -... -``` diff --git a/docs/source/legacy/architecture/add-model.md b/docs/source/legacy/architecture/add-model.md deleted file mode 100644 index 7235ee07fbd5..000000000000 --- a/docs/source/legacy/architecture/add-model.md +++ /dev/null @@ -1,95 +0,0 @@ -(add-model)= - -# Adding a Model - -This document describes how to add a typical decoder-only model in TensorRT LLM. - -## Step 1. Write Modeling Part - -TensorRT LLM provides different levels of APIs: - -- Low-level functions, for example, `concat`, `add`, and `sum`. -- Basic layers, such as, `Linear` and `LayerNorm`. -- High-level layers, such as, `MLP` and `Attention`. -- Base class for typical decoder-only models, such as, `DecoderModelForCausalLM`. - -1. Create a model directory in `tensorrt_llm/models`, for example `my_model`. -2. Write a `model.py` with TensorRT LLM's APIs - -```python -class MyDecoderLayer(Module): - def __init__(self, config: PretrainedConfig, layer_idx: int): - self.layer_idx = layer_idx - self.config = config - self.input_layernorm = LayerNorm(...) - self.attention = Attention(...) - self.post_layernorm = LayerNorm(...) - self.mlp = MLP(...) - - def forward(self, hidden_states, ...): - # decoder layer forward - return hidden_states - -class MyModel(Module): - def __init__(self, config: PretrainedConfig): - self.config = config - self.vocab_embedding = Embedding(...) - self.layers = DecoderLayerList(MyDecoderLayer, config) - self.ln_f = LayerNorm(...) - - def forward(self, input_ids, ...): - # model forward - return hidden_states - - -class MyModelForCausalLM(DecoderModelForCausalLM): - def __init__(self, config: PretrainedConfig): - transformer = MyModel(config) - lm_head = ColumnLinear(...) - super().__init__(config, transformer, lm_head) -``` - - -## Step 2. Implement Weight Conversion - -The weights from source framework need to be converted and bound to the new added TensorRT LLM model. Here is an example of converting HuggingFace weights: - -```python -class MyModelForCausalLM(DecoderModelForCausalLM): - @classmethod - def from_hugging_face( - cls, - hf_model_dir, - dtype='float16', - mapping: Optional[Mapping] = None) -> MyModelForCausalLM - # create a TensorRT LLM MyModelForCausalLM model object - # convert HuggingFace checkpoint to TensorRT LLM expected weights dict - # load the weights to MyModelForCausalLM object -``` - -It's optional to develop a `convert_checkpoint.py` script in the `examples/my_model/` directory for the convenience of offline weights conversion. - -## Step 3. Register New Model - -Please register the new model class `MyModelForCausalLM` in `tensorrt_llm/models/__init__.py`. - -## Step 4. Verify New Model - -At last, let's verify the new model. The typical commands are as following: - -```bash -cd examples/my_model/ - -python convert_checkpoint.py --model_dir hf_model_dir --output_dir tllm_ckpt_dir - -trtllm-build --checkpoint_dir tllm_ckpt_dir --output_dir tllm_engine_dir - -# try the model with a single prompt -python ../run.py --engine_dir tllm_engine_dir --tokenizer_dir hf_model_dir --input_text "Born in north-east France, Soyer trained as a" -# run summarization task -python ../summarize.py --engine_dir tllm_engine_dir --hf_model_dir hf_model_dir --test_trt_llm -``` - -## Reference - -It's recommended to read the workflow[./workflow.md] and checkpoint[./checkpoint.md] documents for more details. diff --git a/docs/source/legacy/architecture/checkpoint.md b/docs/source/legacy/architecture/checkpoint.md deleted file mode 100644 index 81d0955e1eb9..000000000000 --- a/docs/source/legacy/architecture/checkpoint.md +++ /dev/null @@ -1,251 +0,0 @@ -# TensorRT LLM Checkpoint - -> [!WARNING] -> This page describes the **legacy** TensorRT engine-build workflow. -> For new projects, use [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -## Overview - -The earlier versions (pre-0.8 version) of TensorRT LLM were developed with a very aggressive timeline. For those versions, emphasis was not put on defining a unified workflow. Now that TensorRT LLM has reached some level of feature richness, the development team has decided to put more effort into unifying the APIs and workflow of TensorRT LLM. This file documents the workflow around TensorRT LLM checkpoint and the set of CLI tools to generate checkpoint, build engines, and evaluate engines. - -There are three steps in the workflow: - -1. Convert weights from different source frameworks into TensorRT LLM checkpoint. -2. Build the TensorRT LLM checkpoint into TensorRT engines with a unified build command. -3. Load the engines to TensorRT LLM model runner and evaluate with different evaluation tasks. - -``` -NeMo ------------- - | -HuggingFace ------ - | convert build load -Modelopt --------- ----------> TensorRT LLM Checkpoint --------> TensorRT Engine ------> TensorRT LLM ModelRunner - | -JAX -------------- - | -DeepSpeed -------- -``` - -## Prepare the TensorRT LLM Checkpoint - -TensorRT LLM aims at supporting different sources: - -1. Trained models from NVIDIA NeMo, Microsoft DeepSpeed, and JAX -2. Quantized models from NVIDIA Modelopt -3. Popular models from HuggingFace - -TensorRT LLM defines its own checkpoint format. A checkpoint directory includes: - -1. One config `json` file, which contains several model hyper-parameters. -2. One or several rank weights files, each file contains a dictionary of tensors (weights). -The different files are loaded by different ranks in a multi-GPU (multi-process) scenario. - -### Config - -| Field | Type | Default Value | -| :------------------------------------- | :--------- | :------------------ | -| architecture | string | mandatory | -| dtype | string | mandatory | -| logits_dtype | string | 'float32' | -| vocab_size | int | mandatory | -| max_position_embeddings | int | null | -| hidden_size | int | mandatory | -| num_hidden_layers | int | mandatory | -| num_attention_heads | int | mandatory | -| num_key_value_heads | int | num_attention_heads | -| hidden_act | string | mandatory | -| intermediate_size | int | null | -| norm_epsilon | float | 1e-5 | -| position_embedding_type | string | 'learned_absolute' | -| mapping.world_size | int | 1 | -| mapping.tp_size | int | 1 | -| mapping.pp_size | int | 1 | -| quantization.quant_algo | str | null | -| quantization.kv_cache_quant_algo | str | null | -| quantization.group_size | int | 64 | -| quantization.has_zero_point | bool | False | -| quantization.pre_quant_scale | bool | False | -| quantization.exclude_modules | list | null | - -`mapping.world_size` means `mapping` is a dictionary containing the `world_size` sub field. - -```json -{ - "architecture": "OPTForCausalLM", - "mapping": { - "world_size": 1 - } -} -``` - -Supported quantization algorithm list: - -- W8A16 -- W4A16 -- W4A16_AWQ -- W4A8_AWQ -- W4A16_GPTQ -- FP8 -- W8A8_SQ_PER_CHANNEL - -Supported KV cache quantization algorithm list: - -- FP8 -- INT8 - -The config field is extensible, a model could add its own specific config fields. -For example, OPT model has a `do_layer_norm_before` field. - -Here is the model specific config list: - -| Field | Type | Default Value | -| :------------------------------------- | :--------- | :------------------ | -| OPT | | | -| do_layer_norm_before | bool | False | -| | | | -| Falcon | | | -| bias | bool | True | -| new_decoder_architecture | bool | False | -| parallel_attention | bool | False | - -### Rank Weights - -Like PyTorch, the tensor (weight) name is a string containing hierarchical information, -which is uniquely mapped to a certain parameter of a TensorRT LLM model. - -For example, each transformer layer of the OPT model contains an `Attention` layer, an `MLP` layer, and two `LayerNorm` layers. - -#### Attention Weights - -The `Attention` layer contains two `Linear` layers, qkv and dense; each `Linear` layer contains one weight and one bias. -There are four tensors (weights) in total, whose names are: - -- `transformer.layers.0.attention.qkv.weight` -- `transformer.layers.0.attention.qkv.bias` -- `transformer.layers.0.attention.dense.weight` -- `transformer.layers.0.attention.dense.bias` - -where `transformer.layers.0.attention` is the prefix name, indicating that the weights/biases are in the Attention module of the 0-th transformer layer. - -#### MLP Weights - -The `MLP` layer also contains two `Linear` layers, fc and proj; each `Linear` layer contains one weight and one bias. -There are four tensors (weights) in total, whose names are: - -- `transformer.layers.0.mlp.fc.weight` -- `transformer.layers.0.mlp.fc.bias` -- `transformer.layers.0.mlp.proj.weight` -- `transformer.layers.0.mlp.proj.bias` - -where `transformer.layers.0.mlp` is the prefix name, indicating that the weights/biases are in the MLP module of the 0-th transformer layer. - -#### LayerNorm Weights - -Each of the two `LayerNorm` layers, namely `input_layernorm` and `post_layernorm`, contains one weight and one bias. -There are four tensors (weights) in total, whose names are: - -- `transformer.layers.0.input_layernorm.weight` -- `transformer.layers.0.input_layernorm.bias` -- `transformer.layers.0.post_layernorm.weight` -- `transformer.layers.0.post_layernorm.bias` - -where `transformer.layers.0.input_layernorm` and `transformer.layers.0.post_layernorm` are prefix names for the two `layernorm` modules. - -#### KV Cache Quantization Scaling Factors - -If we quantize the model, there will be different tensors (depending on the quantization method applied). -For example, if we quantize the KV cache, the `Attention` layer will have this extra scaling factor: - -- `transformer.layers.0.attention.kv_cache_scaling_factor` - -#### FP8 Quantization Scaling Factors - -Here is the FP8 scaling factors of `attention.qkv` linear layer: - -- `transformer.layers.0.attention.qkv.activation_scaling_factor` -- `transformer.layers.0.attention.qkv.weights_scaling_factor` - -#### AWQ Quantization Scaling Factors - -Here is the AWQ scaling factors of `mlp.fc` linear layer: - -- `transformer.layers.0.mlp.fc.weights_scaling_factor` -- `transformer.layers.0.mlp.fc.prequant_scaling_factor` - - ```{note} - The linear weights in TensorRT LLM checkpoint always follows (`out_feature`, `in_feature`) shape, whereas some quantized linear in TensorRT LLM implemented by plugin may use (`in_feature`, `out_feature`) shape. The `trtllm-build` command adds a transpose operation to post-process it. - -### Example - -Let's take OPT as an example and deploy the model with tensor parallelism 2: - -```bash -cd examples/opt -python3 convert_checkpoint.py --model_dir ./opt-125m \ - --dtype float16 \ - --tp_size 2 \ - --output_dir ./opt/125M/trt_ckpt/fp16/2-gpu/ -``` - -Here is the checkpoint directory: - -``` -./opt/125M/trt_ckpt/fp16/1-gpu/ - config.json - rank0.safetensors - rank1.safetensors -``` - -Here is the `config.json`: - -```json -{ - "architecture": "OPTForCausalLM", - "dtype": "float16", - "logits_dtype": "float32", - "num_hidden_layers": 12, - "num_attention_heads": 12, - "hidden_size": 768, - "vocab_size": 50272, - "position_embedding_type": "learned_absolute", - "max_position_embeddings": 2048, - "hidden_act": "relu", - "mapping": { - "world_size": 2, - "tp_size": 2 - }, - "use_parallel_embedding": false, - "embedding_sharding_dim": 0, - "do_layer_norm_before": true, -} -``` - -## Build Checkpoint into TensorRT Engine - -TensorRT LLM provides a unified build command: `trtllm-build`. Before using it, -you may need to add it to the `PATH`. - -```bash -export PATH=/usr/local/bin:$PATH - -trtllm-build --checkpoint_dir ./opt/125M/trt_ckpt/fp16/2-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./opt/125M/trt_engines/fp16/2-gpu/ -``` - -## Make Evaluation - -```bash -mpirun -n 2 --allow-run-as-root \ - python3 ../summarize.py --engine_dir ./opt/125M/trt_engines/fp16/2-gpu/ \ - --batch_size 1 \ - --test_trt_llm \ - --hf_model_dir opt-125m \ - --data_type fp16 \ - --check_accuracy \ - --tensorrt_llm_rouge1_threshold=14 -``` diff --git a/docs/source/legacy/architecture/core-concepts.md b/docs/source/legacy/architecture/core-concepts.md deleted file mode 100644 index 26b40dbef8fa..000000000000 --- a/docs/source/legacy/architecture/core-concepts.md +++ /dev/null @@ -1,389 +0,0 @@ -(core-concepts)= - -# Model Definition - -> [!WARNING] -> This page describes the **legacy** TensorRT engine-build workflow. -> For new projects, use [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -TensorRT-LLM has a Model Definition API that can be used to define -Large Language Models. This API is built on top of the powerful -[TensorRT Python API](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/index.html) -to create graph representations of deep neural networks in TensorRT. To become -familiar with the core concepts of the TensorRT API, refer to the -[Core Concepts](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/coreConcepts.html) -section of the TensorRT documentation before proceeding further. - -In TensorRT-LLM, the [`tensorrt_llm.Builder`](source:tensorrt_llm/builder.py) class -contains a -[`tensorrt.Builder`](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Core/Builder.html#id1) -object. That instance is used in the `tensorrt_llm.Builder.create_network` -method to create an instance of the -[`tensorrt.INetworkDefinition`](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Graph/Network.html#tensorrt.INetworkDefinition) -class. The `INetworkDefinition` object can then be populated using the free -functions defined in the -[`tensorrt_llm.functional`](source:tensorrt_llm/functional.py). - -A simple example of such a free function is `tensorrt_llm.activation` that inserts a -[`tensorrt.IActivationLayer`](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Graph/Layers.html#tensorrt.IActivationLayer) -node in the graph of the model: - -```python -# In tensorrt_llm.functional: - -def activation(input: Tensor, act_type: trt.ActivationType) -> Tensor: - layer = default_trtnet().add_activation(input.trt_tensor, act_type) # default_trtnet() -> INetworkDefinition - return _create_tensor(layer.get_output(0), layer) -``` - -To make it even easier for users, a few of the most standard activation -functions found in LLMs are derived from that function: - -```python -# In tensorrt_llm.functional: - -relu = partial(activation, act_type=trt.ActivationType.RELU) -sigmoid = partial(activation, act_type=trt.ActivationType.SIGMOID) - -``` - -Specialized activation functions can be used to assemble more advanced -functions such as the `silu` activation: - -```python -# In tensorrt_llm.functional: - -def silu(input: Tensor) -> Tensor: - return input * sigmoid(input) -``` - -When the TensorRT-LLM's Model Definition API is utilized, a graph of the network is -assembled. The graph can later be traversed or transformed using the graph -traversal API exposed by the -[`tensorrt.ILayer`](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Graph/LayerBase.html#tensorrt.ILayer) -class. That graph will also be optimized by TensorRT during the compilation of -the engine, as explained in the next section. - -# Compilation - -Once populated, the instance of the -[`tensorrt.INetworkDefinition`](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Graph/Network.html#tensorrt.INetworkDefinition), -can be compiled into an efficient engine by the -[`tensorrt.Builder`](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Core/Builder.html#id1) -In TensorRT-LLM, it is done through the `build_engine` member function of the -`tensorrt_llm.Builder` class that calls the -[`build_serialized_network`](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Core/Builder.html#tensorrt.Builder.build_serialized_network -method of the -[`tensorrt.Builder`](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Core/Builder.html#id1) -object. That call, if everything works as expected, produces an instance of the -[`tensorrt.IHostMemory`](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/FoundationalTypes/HostMemory.html#tensorrt.IHostMemory) -class. That object is an optimized TensorRT engine that can be stored as a -binary file. - -## TensorRT Compiler - -The TensorRT compiler can sweep through the graph to choose the best kernel for each operation and available GPU. Crucially, it can also identify patterns in the graph where multiple operations are good candidates for being fused into a single kernel. This reduces the required amount of memory movement and the overhead of launching multiple GPU kernels. - -TensorRT also compiles the graph of operations into a single [CUDA Graph](https://developer.nvidia.com/blog/cuda-graphs/) that can be launched all at one time, further reducing the kernel launch overhead. - -The TensorRT compiler is extremely powerful for fusing layers and increasing execution speed, but there are some complex layer fusions—like [FlashAttention](https://arxiv.org/abs/2307.08691) — that involve interleaving many operations together and which can’t be automatically discovered. For those, you can explicitly replace parts of the graph with [plugins](#plugins) at compile time. - -## Model Engine - -The engine file contains the information that you need for executing the model, but LLM usage in practice requires much more than a single forward pass through the model. TensorRT-LLM includes a highly optimized C++ runtime for executing built LLM engines and managing processes like sampling tokens from the model output, managing the KV cache, and batching requests together. - -You can use that runtime directly to execute the model locally, or you can use the TensorRT-LLM runtime backend for NVIDIA Triton Inference Server to serve the model for multiple users. - -## Weight Bindings - -TensorRT engines embed the network weights, that must be known for compilation. -For that reason, the weights must be bound to parameters in the model -definition before calling `tensorrt_llm.Builder.build_engine`. It leads to code like: - -```python -# The Linear operator exposes two parameters (see tensorrt_llm/layers/linear.py): -class Linear(Module): - def __init__(self, ...): - self.weight = Parameter(shape=(self.out_features, self.in_features), dtype=dtype) - self.bias = Parameter(shape=(self.out_features, ), dtype=dtype) - -# The parameters are bound to the weights before compiling the model. See examples/models/core/gpt/weight.py: -tensorrt_llm_gpt.layers[i].mlp.fc.weight.value = fromfile(...) -tensorrt_llm_gpt.layers[i].mlp.fc.bias.value = fromfile(...) -``` - -Note that TensorRT can also -[refit](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#refitting-engine-c) -engines to update the weights after compilation. This feature is available to -TensorRT-LLM users through the `refit_engine` method in the -`tensorrt_llm.Builder` class. - -## Pattern-Matching and Fusion - -One of the key steps performed by TensorRT when it compiles the network graph -is the fusion of operations. Fusion is a well-known technique to improve the -efficiency when executing LLMs. It helps reduce the amount of data transferred -between the memory (DRAM) and the compute cores (CUDA cores as well as Tensor -Cores located on the [Streaming -Multiprocessors](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#introduction) -of a GPU). It also removes kernel launch overhead (each time a kernel is -launched on the GPU, there is a small additional CPU cost that is called the -launch overhead). A classical example is the fusion of the activation function -with the matrix multiplication (matmul) that usually precedes it in the -network. - -In TensorRT-LLM, when defining the model, such a sequence can be written as: - -```python -c = tensorrt_llm.functional.matmul(a, b) -c = tensorrt_llm.functional.relu(c) -``` - -During inference, if the above sequence is executed without fusion, the `c` -tensor has to be written to global memory at the end of the `matmul`, read from -that same memory in `relu` and written again after `relu`. If no other -operation uses the intermediate values between `matmul` and `relu`, it is -suboptimal. That is why, during compilation, TensorRT will identify that -pattern and automatically produce a GPU kernel that applies `relu` at the end -of `matmul` without an intermediate step through global memory. With that -optimization, the `c` tensor is written only once (after `relu`) instead of -twice, and is not read between the two operations. - -The process of identifying the sequences of operations that can be fused is -called _pattern-matching_. TensorRT has a powerful pattern-matching algorithm -that can identify a lot of possible fusions. All the identified patterns are -converted into more efficient kernels by an advanced kernel compiler. - -## Plugins - -The number of possible fusions is almost infinite and some useful fusions -involve very advanced modifications of the graph. A well-known example -is the [Flash-Attention](https://arxiv.org/abs/2205.14135) technique to -optimize the [Multihead-Attention](https://arxiv.org/abs/1706.03762) block -found in many LLMs. Flash-Attention requires modifications to the arithmetic -performed in the sequence `BMM-Softmax-BMM` (where `BMM` stands for Batched -Matrix-Matrix product) and the interleaving of the `for`-loops of the two -batched matrix products. That's non-trivial and not necessarily something -you can expect a compiler to "discover" on its own (or it might require the -support for a [polyhedral -model](https://en.wikipedia.org/wiki/Polytope_model)). - -As a result, even if TensorRT has a powerful pattern-matching algorithm and -supports a lot of possible fusions, there is always the risk that it cannot -identify uncommon and/or very advanced patterns. To overcome that inevitable -limitation, TensorRT offers a powerful mechanism known as -[plugins](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Plugin/pyPlugin.html). - -The plugins are nodes inserted in the network graph definition that map to user-defined -GPU kernels. TensorRT-LLM uses a number of such plugins. They can be found in -the [`cpp/tensorrt_llm/plugins`](source:/cpp/tensorrt_llm/plugins) directory. - -Plugins are written in C++ and follow a well-defined interface described in the -[Extending TensorRT with Custom Layers](https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/extending-custom-layers.html) -section of the TensorRT -[Developer Guide](https://docs.nvidia.com/deeplearning/tensorrt/latest/index.html). -When executed within a TensorRT engine, plugins trigger the execution of -their encapsulated GPU kernels. A fairly simple example of plugins is the -[`QuantizeTensorPlugin`](source:/cpp/tensorrt_llm/plugins/quantizeTensorPlugin) that -triggers a CUDA kernel in the `QuantizeTensorPlugin::enqueue` member function: - -```cpp -// In cpp/tensorrt_llm/plugins/quantizeTensorPlugin/quantizeTensorPlugin.cpp: - -int QuantizeTensorPlugin::enqueue(...) { - if (inputDesc[0].type == DataType::kFLOAT) { - invokeQuantization(...); - } else { - invokeQuantization(...); - } - return 0; -} - -// In cpp/tensorrt_llm/kernels/quantization.cu: - -template -void invokeQuantization(...) { - // The standard <<< >>> construct to launch CUDA kernels - quantizedKernel<<>>(...); -} -``` - -For more details on how TensorRT-LLM implements the GPT Attention operator, see -the [Multi-head, Multi-query and Group-query Attention](../advanced/gpt-attention.md) document. - -# Runtime - -TensorRT-LLM includes an API to implement Python and C++ runtimes. The role of -the runtime components is to load the TensorRT engines and drive their -execution. Typically, for an auto-regressive model like GPT, the runtime is in -charge of loading the engine that implements both the processing of the input -sequence as well as the body of the generation loop. See the [GPT C++ -Runtime](../advanced/gpt-runtime.md) document for details on the C++ Runtime. - -(multi-gpu-multi-node)= - -# Multi-GPU and Multi-Node Support - -Even if TensorRT is designed for single-GPU systems, TensorRT-LLM adds the -support for systems with multiple GPUs and nodes. It is enabled -using TensorRT plugins that wrap communication primitives from the -[NCCL](https://developer.nvidia.com/nccl) library as well as a custom -plugin that optimize the All-Reduce primitive in the presence of All-to-all -connections between GPUs (through NVSwitch in DGX systems). - -The communication plugins can be found in -[cpp/tensorrt_llm/plugins/ncclPlugin](source:cpp/tensorrt_llm/plugins/ncclPlugin) -and the multi-GPU functions are exposed in the TensorRT-LLM Model Definition API -as: - -```python -# In tensorrt_llm/functional.py: - -# Collectives. -def allreduce(tensor: Tensor, group: List[int]) -> Tensor -def allgather(tensor: Tensor, group: List[int], gather_dim: int = 0) -> Tensor - -# Point-to-point communication primitives. -def send(tensor: Tensor, tgt: int) -> Tensor -def recv(tensor: Tensor, src: int) -> Tensor -``` - -The multi-GPU support can be enabled through two different modes of model -parallelism: Tensor Parallelism and Pipeline Parallelism. The former mode -splits the different layers of a model across the GPUs. Each GPU runs the -entire network and synchronizes with its siblings when needed. The Pipeline -Parallelism distributes the different layers to the GPUs. Each GPU runs a -subset of the entire model and communications happen at the boundary of those -subsets of layers. Tensor Parallelism usually leads to more balanced executions -but requires more memory bandwidth between the GPUs. Pipeline Parallelism -reduces the need for high-bandwidth communication but may incur load-balancing -issues and may be less efficient in terms of GPU utilization. - -## Examples - -Here are examples of Llama 3.1 70B and Llama 3.1 405B showing how to perform multi-GPU and multi-node inference in TensorRT-LLM. The example of Llama 3.1 70B performs multi-GPU inference on a single node, while the example of Llama 3.1 405B performs multi-node inference. - -### Llama 3.1 70B - -The following sample commands build an engine for running the Llama 3.1 70B model with tensor parallelism (TP=4) using 4 GPUs on a single node. - -```bash -folder_trt_llm=../TensorRT-LLM -model_dir=Llama-3.1-70B -ckpt_dir=ckpt_llama_3.1_70b -engine_dir=engine_llama_3.1_70b -dtype=bfloat16 -tp_size=4 -pp_size=1 -kv_cache_type=paged -max_input_len=128 -max_output_len=128 -max_batch_size=4 -workers=$(( tp_size * pp_size )) - -python ${folder_trt_llm}/examples/models/core/llama/convert_checkpoint.py \ - --output_dir ${ckpt_dir} \ - --model_dir ${model_dir} \ - --dtype ${dtype} \ - --tp_size ${tp_size} \ - --pp_size ${pp_size} \ - --workers ${workers} \ - --use_parallel_embedding - -trtllm-build \ - --output_dir ${engine_dir} \ - --checkpoint_dir ${ckpt_dir} \ - --gemm_plugin ${dtype} \ - --gpt_attention_plugin ${dtype} \ - --kv_cache_type ${kv_cache_type} \ - --max_input_len ${max_input_len} \ - --max_seq_len $(( max_input_len + max_output_len )) \ - --max_batch_size ${max_batch_size} \ - --workers ${workers} -``` - -The following sample commands perform inference using 4 GPUs on a single node by running `examples/run.py`. - -```bash -input_text="Born in north-east France, Soyer trained as a" - -mpirun -n $(( tp_size * pp_size )) \ - python ${folder_trt_llm}/examples/run.py \ - --engine_dir ${engine_dir} \ - --tokenizer_dir ${model_dir} \ - --input_text "${input_text}" \ - --max_output_len ${max_output_len} -``` - -### Llama 3.1 405B - -The following sample commands build an engine for running the Llama 3.1 405B model with tensor parallelism (TP=16) on 2 nodes that each have 8 GPUs. Although the model runs on multiple nodes, you can build the engine on a single node. - -```bash -folder_trt_llm=../TensorRT-LLM -model_dir=Llama-3.1-405B -ckpt_dir=ckpt_llama_3.1_405b -engine_dir=engine_llama_3.1_405b -dtype=bfloat16 -tp_size=16 -pp_size=1 -kv_cache_type=paged -max_input_len=128 -max_output_len=128 -max_batch_size=4 -workers=8 - -python ${folder_trt_llm}/examples/models/core/llama/convert_checkpoint.py \ - --output_dir ${ckpt_dir} \ - --model_dir ${model_dir} \ - --dtype ${dtype} \ - --tp_size ${tp_size} \ - --pp_size ${pp_size} \ - --workers ${workers} \ - --use_parallel_embedding - -trtllm-build \ - --output_dir ${engine_dir} \ - --checkpoint_dir ${ckpt_dir} \ - --gemm_plugin ${dtype} \ - --gpt_attention_plugin ${dtype} \ - --kv_cache_type ${kv_cache_type} \ - --max_input_len ${max_input_len} \ - --max_seq_len $(( max_input_len + max_output_len )) \ - --max_batch_size ${max_batch_size} \ - --workers ${workers} -``` - -The following sample script, `launch_llama_3.1_405b.sh`, shows how to perform inference with Slurm on 2 nodes that each have 8 GPUs. If you use a different workload management software, the key concern is to run the `examples/run.py` command. - -```bash -#!/bin/bash -#SBATCH --account account -#SBATCH --partition partition -#SBATCH --job-name job-name -#SBATCH --time 1:00:00 -#SBATCH --nodes 2 - -folder_trt_llm=../TensorRT-LLM -engine_dir=engine_llama_3.1_405b -model_dir=Llama-3.1-405B -max_output_len=128 - -input_text="Born in north-east France, Soyer trained as a" - -srun \ - --ntasks-per-node 8 \ - --mpi pmix \ - python ${folder_trt_llm}/examples/run.py \ - --engine_dir ${engine_dir} \ - --tokenizer_dir ${model_dir} \ - --input_text "${input_text}" \ - --max_output_len ${max_output_len} -``` - -You can perform inference by running the script on the Slurm cluster. - -```bash -sbatch launch_llama_3.1_405b.sh -``` diff --git a/docs/source/legacy/architecture/model-weights-loader.md b/docs/source/legacy/architecture/model-weights-loader.md deleted file mode 100644 index d41191364f1e..000000000000 --- a/docs/source/legacy/architecture/model-weights-loader.md +++ /dev/null @@ -1,256 +0,0 @@ -# TensorRT-LLM Model Weights Loader - -## Overview - -The weights loader is designed for easily converting and loading external weight checkpoints into TensorRT-LLM models. - -## Workflow - -Weight checkpoints can be generated from all sources, and may have different naming and data layouts compared to TRT-LLM's requirements. E.g.: - -```bash -# HuggingFace LLaMA checkpoints -{ - "model.embed_tokens.weight": torch.Tensor([vocab_size, hidden_size]) - "model.layers.0.input_layernorm.weight": torch.Tensor([hidden_size]), - "model.layers.0.mlp.down_proj.weight": torch.Tensor([hidden_size, inter_size]), - "model.layers.0.mlp.gate_proj.weight": torch.Tensor([inter_size, hidden_size]), - "model.layers.0.mlp.up_proj.weight": torch.Tensor([inter_size, hidden_size]), - "model.layers.0.post_attention_layernorm.weight": torch.Tensor([hidden_size]), - "model.layers.0.self_attn.q_proj.weight": torch.Tensor([hidden_size, hidden_size]), - "model.layers.0.self_attn.k_proj.weight": torch.Tensor([hidden_size, hidden_size]), - "model.layers.0.self_attn.v_proj.weight": torch.Tensor([hidden_size, hidden_size]), - "model.layers.0.self_attn.o_proj.weight": torch.Tensor([hidden_size, hidden_size]), - ..., -} -# TensorRT-LLM expected weights -{ - "transformer.vocab_embedding.weight": torch.Tensor([vocab_size, hidden_size]) - "transformer.layers.0.input_layernorm.weight": torch.Tensor([hidden_size]), - "transformer.layers.0.mlp.down_proj.weight": torch.Tensor([hidden_size, inter_size]), - "transformer.layers.0.mlp.gate_proj.weight": torch.Tensor([inter_size, hidden_size]), - "transformer.layers.0.mlp.up_proj.weight": torch.Tensor([inter_size, hidden_size]), - "transformer.layers.0.post_layernorm.weight": torch.Tensor([hidden_size]), - "transformer.layers.0.attention.qkv.weight": torch.Tensor([hidden_size * 3, hidden_size]), # Different layout - "transformer.layers.0.attention.dense.weight": torch.Tensor([hidden_size, hidden_size]), - ..., -} -``` - -Conversion means converting the dictionary of `{external_keys:external_weights}` into `{tllm_keys:tllm_weights}`, it includes changing the naming logic and data layouts, and consists of the following parts: - -1. Translate a TRT-LLM parameter name into external-format name(s). -2. Loading tensor slice(s) according to the translated names. -3. Postprocess the tensor(s) into target layout. - -### Translator - -TRT-LLM parameter names are translated in units of sections divided by dots. E.g.: - -| TensorRT-LLM key | `transformer` |.| `layers` |.| `0` |.| `attention` |.| `dense` |.| `weight` | -| :---------------------: | :-----------: |-| :------: |-|:---:|-| :---------: |-| :------: |-| :------: | -| Translated external key | `model` |.| `layers` |.| `0` |.| `self_attn` |.| `o_proj` |.| `weight` | - -The mapping between TRT-LLM keywords and HF keywords are described in `tllm_to_externel_key_dict` of `ModelWeightsLoader` class object. \ -If any of the mappings has one-to-multiple corresponding, the translated key will get multiplied accordingly. E.g.: - -| TensorRT-LLM key and related keyword mapping | Translated external keys | -| :----------------------------------------------------------: | :----------------------: | -| `transformer.layers.0.attention.qkv.weight`
`{"qkv":[q_proj, k_proj, v_proj]}` | `model.layers.0.self_attn.q_proj.weights`
`model.layers.0.self_attn.k_proj.weights`
`model.layers.0.self_attn.v_proj.weights`| -| `transformer.layers.0.mlp.fc.weight`
`{"weight":[qweight, qzeros, scales]}` | `model.layers.0.mlp.gate_proj.qweight`
`model.layers.0.mlp.gate_proj.qzeros`
`model.layers.0.mlp.gate_proj.scales`| - -The default `tllm_to_externel_key_dict` is based on HF LLaMA as: - -```python -class ModelWeightsLoader: - def __init__(self, model_dir, customized_key_dict: dict = {}) -> None: - ... - self.tllm_to_externel_key_dict = { - "transformer": "model", - "vocab_embedding": "embed_tokens", - "lm_head": "lm_head", - "ln_f": "norm", - "attention": "self_attn", - "qkv": ["q_proj", "k_proj", "v_proj"], - "dense": "o_proj", - "gate": "up_proj", - "proj": "down_proj", - "fc": "gate_proj", - "input_layernorm": "input_layernorm", - "post_layernorm": "post_attention_layernorm", - } - self.tllm_to_externel_key_dict.update(customized_key_dict) - ... -``` - -It can be updated through passing `customized_key_dict` when initializing `ModelWeightsLoader`. - -The dictionary will also get updated according to the layer classes. When iterating over parameters, -if the layer class has attribute `tllm_to_externel_key_dict`, for keywords exist both in the default one and the layer-specified one, -the weight loader will translate according to the layer attribute with higher priority. -This can enable the support for different quantization precisions automatically. - - -### Loading function - -The loading function can load an arbitrary tensor slice according to its `key`, `tp_size`, `tp_dim` and `tp_rank`. - -The template for loading function is as following. - -```python -def load_tensor(self, key, tp_size, tp_dim, tp_rank): - # Retrieve file pointer index - if key in self.shard_map: - ptr_idx = self.shard_map[key] - else: - return None - - # Load tensor from the corresponding shard - if self.format == ModelWeightsFormat.SAFETENSORS: - tensor = self.shards[ptr_idx].get_slice(key) - tensor_shape = tensor.get_shape() - else: - ... - - # Shard and return a tensor slice - slice_shape = ... - return tensor[slice_shape] -``` - -When initializing the `ModelWeightsLoader` object, the file format will be derived from `model_dir` through `detect_format`. The following formats are supported for now: - - * Directory contains or file named `*.safetensors` (Recommended, has better performance) - * Directory contains or file named `*.bin` - * Directory contains or file named `*.pth` - -To support other formats or in-memory loaded models, the format need to be claimed in `ModelWeightsFormat`, `detect_format()`, `preload()` and `load_tensor()`. - -### Postprocessing functions - -After translation and loading, a TRT-LLM key will become a tensor or a list of tensors, which is the input of postprocessing functions. \ -Operations including QKV concatenating, MoE weight stacking and weight-only quantization can be handled here. -The template of postprocessing function is: - -```python -# Example for 1-1 weights mapping -class CustomizedModuleA(Module): - def __init__(...): - super().__init__(...) - ... - self.tp_dim = 0 # Need to set or inherit from parent class - - def postprocess(self, tllm_key, weights, **kwargs): - weights = proc(weights) - return {tllm_key: weights} - -# Example for multiple-multiple weights mapping -class CustomizedModuleB(Module): - def __init__(...): - super().__init__(...) - ... - self.tp_dim = 0 # Need to set or inherit from parent class - # The default value of "weights" in tllm_to_externel_key_dict will be override - self.tllm_to_externel_key_dict = {"weight": ["qweight", "qzeros", "scales"]} - - def postprocess(self, tllm_key, weights, **kwargs): - # Skipped the postprocess of zeros and weights_scaling_factor - # They are loaded in the postprocess of weight - config = kwargs.get("config", None) # Passed through kwargs by default - if not tllm_key.endswith("weight"): - return {} - # The order in weights is defined in tllm_to_externel_key_dict - qweight, qzeros, scales = weights - proccessed_weight, proccessed_zeros = proc(qweight, qzeros, config.num_heads) - return { - tllm_key: proccessed_weight, - tllm_key.replace("weight", "zeros"): proccessed_zeros, - tllm_key.replace("weight", "weights_scaling_factor"): scales, - } -``` - -## Examples - -The `ModelWeightsLoader` class can support different models with the following levels: - -### Natively supported models -For models with native support, users can call the default weight loader without any other operations. -```python -# Using the model weights loader for LLaMA -from tensorrt_llm.models.model_weights_loader import ModelWeightsLoader -loader = ModelWeightsLoader(external_checkpoint_dir) -loader.generate_tllm_weights(trtllm_model) -``` -For calibration-free quantization precisions, passing a properly quantized `trtllm_model` will let the weight loader load at the given precision accordingly. The configurations will be read from `trtllm_model.config` automatically. For now, LLaMA family models using the default `tllm_to_externel_key_dict` is supported natively. - -### Models with customized key names -For models with different naming logic, users can still call the default weight loader with `customized_key_dict` specified. -```python -# Using the model weights loader for the LLM part of LLaVA -from tensorrt_llm.models.model_weights_loader import ModelWeightsLoader -llava_dict = { - "transformer": "language_model.model", - "lm_head": "language_model.lm_head" -} -loader = ModelWeightsLoader(external_checkpoint_dir, llava_dict) -loader.generate_tllm_weights(trtllm_model) -``` -Users need to specify the different part from the default `tllm_to_externel_key_dict`. The loader still have support across different precisions. -The support for LLaVA and Exaone is in `LLaMAForCausalLM.from_hugging_face()` of [model.py](../../../tensorrt_llm/models/llama/model.py), and can also be taken as examples. - -### Models with customized weight layout -For models with different weight layout, users can write the conversion loop explicitly and do customized operations. -```python -# Using the model weights loader for BLOOM -from tensorrt_llm.models.model_weights_loader import ModelWeightsLoader -bloom_dict = { - "transformer": "", - "layers": "h", - "ln_f": "ln_f", - "lm_head": "word_embeddings", - "ln_embed": "word_embeddings_layernorm", - "vocab_embedding": "word_embeddings", - "attention": "self_attention", - "qkv": "query_key_value", - "dense": "dense", - "fc": "dense_h_to_4h", - "proj": "dense_4h_to_h", - "post_layernorm": "post_attention_layernorm", -} -loader = ModelWeightsLoader(external_checkpoint_dir, bloom_dict) -# See ModelWeightsLoader.generate_tllm_weights() -loader.update_key_mapping(trtllm_model) -tllm_weights = {} -for tllm_key, _ in tqdm(trtllm_model.named_parameters()): - if tllm_key.endswith("qkv"): - # Passing the callable handle - tllm_weights.update(loader.load(tllm_key, preprocess=customized_preprocess)) - else: - tllm_weights.update(loader.load(tllm_key)) -loader.fill(tllm_weights) -``` -This will apply `preprocess` after `load_tensor()` and before `postprocess`, and demonstrates how to convert the loaded shard into default HF layout. The loader still have support for precisions quantized from FP16/BF16 (e.g. INT8-wo/INT4-wo), the other precisions may require special operations, and can be addressed inside the `preprocess` function. -The support for Qwen-1 is in `QWenForCausalLM.from_hugging_face()` of [model.py](../../../tensorrt_llm/models/qwen/model.py), and can also be taken as example. - -### Fully customized -If the model weights loader cannot satisfy the requirements, users can write the conversion loop totally on their own. -```python -tllm_weights = {} -for tllm_key, param in tqdm(trtllm_model.named_parameters()): - # Load from external checkpoints - # The load_tensor() function can also be called here - tensor = ... - # Convert tensor and set the values according to the config - if trtllm_model.config.quantization.quant_algo == xxx: - ... - else: - ... - param.value = tensor -``` -In this mode, every precision require user's own support. - -## Troubleshooting -The weights loader is enabled for LLaMA family models and Qwen models by default with TensorRT flow only. - -If users encounter failure caused by `ModelWeightsLoader`, a workaround is passing environmental variable `TRTLLM_DISABLE_UNIFIED_CONVERTER=1` to disable the model weights loader and fallback to the legacy path. - -This workaround will be removed in future version after the LLaMA/Qwen weights conversion is stable. diff --git a/docs/source/legacy/architecture/workflow.md b/docs/source/legacy/architecture/workflow.md deleted file mode 100644 index 3c880e9723fb..000000000000 --- a/docs/source/legacy/architecture/workflow.md +++ /dev/null @@ -1,241 +0,0 @@ -# TensorRT-LLM Build Workflow - -> [!WARNING] -> This page describes the **legacy** TensorRT engine-build workflow. -> For new projects, use [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -## Overview - - -The build workflow contains two major steps. - -1. Create TensorRT-LLM models from existing model checkpoints exported by the training framework. -2. Build the TensorRT-LLM models to TensorRT-LLM engines. - -To generalize the TensorRT-LLM optimization features to all models, and to share the same workflow between different models for TensorRT-LLM users, TensorRT-LLM has conventions about how the models shall be defined and how the models shall be imported. - -TensorRT-LLM checkpoint convention is documented in [](checkpoint.md) and all decoder-only models had been migrated to adopt the convention. Model-specific convert_checkpoint.py scripts are shipped as source code in example directories, and a trtllm-build CLI tool had been added. However, there are some disadvantages of providing convert checkpoint scripts outside the core TensorRT-LLM lib as example: - -1. TensorRT-LLM evolves so quickly that the model's definition code might have changed for better performance; which means the `convert_checkpoint.py` is out of date. - - -2. TensorRT-LLM is creating a new set of high-level APIs which handle model conversion, engine building, and inference in one class for easier-of-use. Thus, the high-level APIs need to call the weights conversion code, which shall be part of TensorRT-LLM core lib, not the example. And the conversion code of different models shall have same interface such that the high-level APIs do not need to add many ad-hoc code for different models. - -To mitigate these issues, the model specific `convert_checkpoint.py` scripts are being refactored. Most of the conversion code will be moved into core lib, sitting next to the model definition. Refer to `tensorrt_llm/models/llama/` as an example. There is a new set of APIs for importing models and converting weights. The 0.9 release refactored the LLaMA model class to adopt the new APIs, others models' refactor work is ongoing. - - -## Conversion APIs - -The API for weight conversion of the LLaMA model looks like this. A `TopModelMixin` class is introduced, `from_hugging_face()` interface is declared, the `LLaMAForCausalLM` class inherits `TopModelMixin` (not direct parent class, but in its base class hierarchy), and implements the interface. - -```python -class TopModelMixin - @classmethod - def from_hugging_face(cls, - hf_model_dir: str, - dtype: Optional[str] = 'float16', - mapping: Optional[Mapping] = None, - **kwargs): - raise NotImplementedError("Subclass shall override this") - -# TopModelMixin is in the part of base class hierarchy -class LLaMAForCausalLM (DecoderModelForCausalLM): - @classmethod - def from_hugging_face(cls, - hf_model_dir, - dtype='float16', - mapping: Optional[Mapping] = None) -> LLaMAForCausalLM: - # creating a TensorRT-LLM llama model object - # converting HuggingFace checkpoint to TensorRT-LLM expected weights dict - # Load the weights to llama model object -``` - - -Then, in the convert_checkpoint.py script in the -[`examples/models/core/llama/`](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/llama/) directory of the GitHub repo, -the logic can be greatly simplified. Even if the model definition code of TensorRT-LLM LLaMA class is changed due to some reason, the `from_hugging_face` API will keep the same, thus the existing workflow using this interface will not be affected. - - -```python -#other args omitted for simplicity here. -llama = LLaMAForCausalLM.from_hugging_face(model_dir, dtype, mapping=mapping) -llama.save_checkpoint(output_dir, save_config=(rank==0)) -``` - -The `from_hugging_face` API does not save the checkpoint into disk intentionally, instead it returns an in-memory object. Call `save_checkpoint` to save the models. This keeps the flexibility and makes the flow of convert->build in one process faster. Typically, saving and loading disk for large models are slower and thus should be avoided. - - -Since LLaMA models were also released with different formats, such as the Meta checkpoint, the `LLaMAForCausalLM` class has a `from_meta_ckpt` function for that. This function is not declared in the `TopModelMixin` class due to it being LLaMA specific, and therefore, other models don't use it. - - -In the 0.9 release, only LLaMA is refactored. Since popular LLaMA (and its variants) models are released by Hugging Face and Meta checkpoint formats, only these two functions are implemented. - - -In future releases, there might be `from_jax`, `from_nemo`, `from_keras` or other factory methods for different training checkpoints added. -For example, the Gemma 2B model and the convert_checkpoint.py file in the [`examples/models/core/gemma`](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/gemma/) -directory support JAX and Keras formats in addition to Hugging Face. The model developers can choose to implement **any subset** of these factory methods for the models they contributed to TensorRT-LLM. - - -For some formats which are not supported by TensorRT-LLM model developers, you still have the freedom to implement your own weights conversion outside the core lib; the flow will look like this: - - -```python -config = read_config_from_the_custom_training_checkpoint(model_dir) -llama = LLaMAForCausalLM(config) - -# option 1: -# Create a weights dict and then calls LLaMAForCausalLM.load -weights_dict = convert_weights_from_custom_training_checkpoint(model_dir) -llama.load(weights_dict) - -# option 2: -# Internally assign the model parameters directly -convert_and_load_weights_into_trtllm_llama(llama, model_dir) -# Use the llama object as usual, to save the checkpoint or build engines -``` - -Though there are some limitations and pitfalls of doing these custom weights loading, if the model definition is inside TensorRT-LLM core lib, and the weights loading/conversion are outside the core lib, the conversion code might need to be updated when new TensorRT-LLM is released. - - -## Quantization APIs - -TensorRT-LLM relies on NVIDIA Modelopt toolkit to support some of the quantization like: FP8, W4A16_AWQ, W4A8_AWQ, while it also has some its own quantization implementation for Smooth Quant, INT8 KV cache, and INT4/INT8 weight only. - - -In TensorRT-LLM 0.8 version: - -* For Modelopt-supported quantization algorithms, a standalone script, - [example/quantization/quantize.py](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/quantization/quantize.py) - can export TensorRT-LLM checkpoints, and the trtllm-build command needs to be executed to build the checkpoints to engines. - -* For the non-Modelopt quantization algorithms, users need to use the per-model convert_checkpoint.py scripts to export TensorRT-LLM checkpoints. - -Use the `quantize()` interface to unify the different quantization flows. The default implementation is added in the `PretrainedModel` class. - - -```python -class PretrainedModel: - @classmethod - def quantize( - cls, - hf_model_dir, - output_dir, - quant_config: QuantConfig, - mapping: Optional[Mapping] = None): #some args are omitted here - # Internally quantize the given hugging face models using Modelopt - # and save the checkpoint to output_dir -``` - -* The default implementation only handles the Modelopt supported quantization. The LLaMA class then inherits this `PretrainedModel` and dispatches the Modelopt quantization to the super class's default implementation. -* The model developer raises errors in the sub-class implementation if the new model is not supported by Modelopt yet. - - -```python -class LLaMAForCausalLM: - @classmethod - def quantize( - cls, - hf_model_dir, - output_dir, - quant_config: QuantiConfig, - mapping: Optional[Mapping] = None): #some args are omitted here - use_modelopt_quantization = ... # determine if to use Modelopt or use native - if use_modelopt_quantization: - super().quantize(hf_model_dir, - output_dir, - quant_config) - else: - # handles TensorRT-LLM native model specific quantization - # or raise exceptions if not supported -``` - - -The `quantize` API is designed to take multi-GPU resources internally to make quantization. For example, a LLaMA 70B BF16 takes 140G memory, if we make FP8 quantization, then, another 70G is needed. So, we need at least 210G, 4 * A100(H100) is needed to quantize the LLaMA 70B model. If you want to call `quantize` API inside a MPI program, be cautious and ensure the quantize API is only called by rank 0. - - -Usage of the `quantize` API in an MPI program looks like this, only rank 0 calls it. In an non-MPI program, the `if rank == 0` and the `mpi_barrier()` are not needed. - -```python -quant_config = QuantConfig() -quant_config.quant_algo = quant_mode.W4A16_AWQ -mapping = Mapping(world_size=tp_size, tp_size=tp_size) -if rank == 0: - LLaMAForCausalLM.quantize(hf_model_dir, - checkpoint_dir, - quant_config=quant_config) -mpi_barrier() # wait for rank-0 to finish the quantization -llama = LLaMAForCausalLM.from_checkpoint(checkpoint_dir, rank) -engine = build(llama, build_config) -engine.save(engine_dir) -``` - - -The `examples/quantization/quantize.py` is kept for backward compatibility. - - -## Build APIs - - -The `tensorrt_llm.build` API builds the TensorRT-LLM model object to TensorRT-LLM engine. This new API replaced the older flow: creating a builder, creating a network object, tracing the model to the network, and building TensorRT engines. -The usage of this API looks like this: - -```python -llama = ... # create LLaMAForCausalLM object -build_config = BuildConfig(max_batch_size=1) -engine = tensorrt_llm.build(llama, build_config) -engine.save(engine_dir) -``` - - -The Llama object can be created by any method mentioned in the [](#conversion-apis) or [](#quantization-apis) sections. - - -The `trtllm-build` CLI tool is a thin wrapper around this `tensorrt_llm.build` API. The flags of the CLI tool are kept close to the fields of the `BuildConfig` class. - - -If a model were to be saved into disk and then built to the engine later, TensorRT-LLM provides a `from_checkpoint` API to deserialize the checkpoint. - -```python -## TensorRT-LLM code -class PretrainedModel: - @classmethod - def from_checkpoint(cls, - ckpt_dir: str, - rank: int = 0, - config: PretrainedConfig = None): - # Internally load the model weights from a given checkpoint directory -``` - - -The `from_checkpoint` API is called to deserialize the checkpoint to a model object. The `tensorrt_llm.build` API can be called to build the engine. - - -```python -llama = LLaMAForCausalLM.from_checkpoint(checkpoint_dir) -engine = build(llama, build_config) -engine.save(engine_dir) -``` - -## CLI Tools - -All the weights conversion, quantization, and build APIs mentioned above have corresponding CLI tools for convenience. - -* Model specific `convert_checkpoint.py` scripts are inside the `examples//` folder. -* A unified quantization script is inside the `examples/quantization/quantize.py` and can be shared by all **supported** models. -* A `trtllm-build` CLI tool builds all models from TensorRT-LLM checkpoint. - -Refer to the following considerations for the CLI tools: - -* These scripts and tools should be used for scripting. Do not import the Python functions/class defined in these tools. TensorRT-LLM does not promise the content of these scripts can be compatible with previous versions. The options of these tools may also be changed when it’s not avoidable. - -* These scripts in the example folder may use TensorRT-LLM internal/unstable APIs, which is not guaranteed to work if the examples’ version and the TensorRT-LLM install version are mismatched. There are some GitHub issues caused by version mismatch. - - https://github.com/NVIDIA/TensorRT-LLM/issues/1293 - - https://github.com/NVIDIA/TensorRT-LLM/issues/1252 - - https://github.com/NVIDIA/TensorRT-LLM/issues/1079 - - You should always install the same TensorRT-LLM version specified in `examples//requirements.txt`. - -* In the future, the per-model conversion script may or may not be unified to one single script shared by models, given the nature of different models’ attributes may be different. However, the TensorRT-LLM team will try to make sure the flags for the same feature are consistent between different scripts. - -* The TensorRT-LLM team encourages use of the new low-level conversion/quantization/build API instead of these scripts. The conversion APIs will be added model-by-model gradually, which may span a few releases. diff --git a/docs/source/legacy/dev-on-cloud/build-image-to-dockerhub.md b/docs/source/legacy/dev-on-cloud/build-image-to-dockerhub.md deleted file mode 100644 index a07fbfdbc819..000000000000 --- a/docs/source/legacy/dev-on-cloud/build-image-to-dockerhub.md +++ /dev/null @@ -1,56 +0,0 @@ -(build-image-to-dockerhub)= - -# Build the TensorRT LLM Docker Image -When you develop trt-llm on cloud platform such as runpod, you may need to provide a docker image for the platform. So you firstly need to upload the image to dockerhub. - -## Build the TensorRT LLM Docker Image and Upload to DockerHub - -```bash -make -C docker build -``` -Then we can get the docker image named `tensorrt_llm/devel:latest` - -### Enable ssh access to the container -Since the default docker image doesn’t have ssh support, we can’t ssh into it. We need to add ssh support to the container. -Let’s first create a new Dockerfile with below content: - -```Dockerfile -FROM tensorrt_llm/devel:latest - -RUN apt update && apt install openssh-server -y -RUN mkdir -p /run/sshd && chmod 755 /run/sshd -RUN mkdir -p /root/.ssh && chmod 700 /root/.ssh && touch /root/.ssh/authorized_keys && chmod 600 /root/.ssh/authorized_keys -# add sshd to entrypoint script -RUN echo "sshd -E /opt/sshd.log" >> /opt/nvidia/entrypoint.d/99-start-sshd.sh -``` - -If we save this Dockerfile as `Dockerfile.ssh`. Then we can build the docker image with below command: - -```bash -docker build -t tensorrt_llm/devel:with_ssh -f Dockerfile.ssh . -``` - -Then we can get the docker image named `tensorrt_llm/devel:with_ssh` - -## Upload the Docker Image to DockerHub - -You need to register a [dockerhub](https://hub.docker.com) account first if you don't have one. - -Then you can click 'Personal Access Tokens' in the user menu and create a new token. - -With the token, you can login to dockerhub with below command: - -```bash -docker login -u -``` - -Enter the token to the console. - -After login, you can tag and push the docker image to dockerhub with below command: - -```bash -docker tag tensorrt_llm/devel:with_ssh /tensorrt_llm:devel -docker push /tensorrt_llm:devel -``` - -Finally, you can see the docker image in your dockerhub repository and can use it with the link such as `docker.io//tensorrt_llm:devel`. diff --git a/docs/source/legacy/dev-on-cloud/dev-on-runpod.md b/docs/source/legacy/dev-on-cloud/dev-on-runpod.md deleted file mode 100644 index 254f8fe14895..000000000000 --- a/docs/source/legacy/dev-on-cloud/dev-on-runpod.md +++ /dev/null @@ -1,47 +0,0 @@ -(dev-on-runpod)= - -# Develop TensorRT LLM on Runpod -[Runpod](https://runpod.io) is a popular cloud platform among many researchers. This doc describes how to develop TensorRT LLM on Runpod. - -## Prepare - -### Create a Runpod account -Please refer to the [Runpod Getting Started](https://docs.runpod.io/get-started/). - -### Configure SSH Key -Please refer to the [Configure SSH Key](https://docs.runpod.io/pods/configuration/use-ssh). - -Note that we can skip the step of "Start your Pod. Make sure of the following things" here as we will introduce it below. - -## Build the TensorRT LLM Docker Image and Upload to DockerHub -Please refer to the [Build Image to DockerHub](build-image-to-dockerhub.md). - -Note that the docker image must enable ssh access. See on [Enable ssh access to the container](build-image-to-dockerhub.md#enable-ssh-access-to-the-container). - -## Create a Pod Template -Click the "Template" button in the menu, then click "Create Template". - -Fill the docker image link of DockerHub such as `docker.io//tensorrt_llm:devel` on "Docker Image" field. - -Fill "22" into "Expose TCP Ports" field. - -Fill -```bash -sleep infinity -``` -into 'Container Start Command' field. - -## Connect to the Pod -Please refer to the [Connect to the Pod](https://docs.runpod.io/pods/connect-to-a-pod). - -You can connect the pod with SSH or Web Terminal. - -If you want to connect the pod with SSH, you can copy the command from "SSH over exposed TCP" field and run it on your host. - -In some scenarios such as using a team account, your public key has not been added to the pod successfully. You can directly add this command to the 'Container Start Command' field as: - -```bash -bash -c 'echo "" >> ~/.ssh/authorized_keys;sleep infinity' -``` - -Enjoy your development! diff --git a/docs/source/legacy/key-features.md b/docs/source/legacy/key-features.md deleted file mode 100644 index 6fac19262df4..000000000000 --- a/docs/source/legacy/key-features.md +++ /dev/null @@ -1,10 +0,0 @@ -# Key Features - -This document lists key features supported in TensorRT-LLM. - -- [Quantization](reference/precision.md) -- [Inflight Batching](advanced/gpt-attention.md#in-flight-batching) -- [Chunked Context](advanced/gpt-attention.md#chunked-context) -- [LoRA](advanced/lora.md) -- [KV Cache Reuse](advanced/kv-cache-reuse.md) -- [Speculative Sampling](advanced/speculative-decoding.md) diff --git a/docs/source/legacy/performance/perf-analysis.md b/docs/source/legacy/performance/perf-analysis.md deleted file mode 100644 index 71f5320ee3a5..000000000000 --- a/docs/source/legacy/performance/perf-analysis.md +++ /dev/null @@ -1,91 +0,0 @@ -(perf-analysis)= - -# Performance Analysis - -NVIDIA Nsight Systems reports at the application level are highly informative. Metric sampling capabilities have increased over generations and provide a clean middle-ground between timing analysis and kernel-level deep dives with NVIDIA Nsight Compute. - -Given the potential long runtimes of Large Languages Models (LLMs) and the diversity of workloads a model may experience during a single inference pass or binary execution, we have added features to TensorRT-LLM to get the most out of Nsight Systems capabilities. This document outlines those features as well as provides examples of how to best utilize them to understand your application. - - -## Feature Descriptions - -The main functionality here: - * Relies on toggling the CUDA profiler runtime API on and off. - * (PyTorch workflow only) Toggling the PyTorch profiler on and off. - * Provides a means to understand which regions a user may want to focus on. - -Toggling the CUDA profiler runtime API on and off: - * Allows users to know specifically what the profiled region corresponds to. - * Results in smaller files to post-process (for metric extraction or similar). - -(PyTorch workflow only) Toggling the PyTorch profiler on and off: - * Help users to analyze the performance breakdown in the model. - * Results in smaller files to post-process (for metric extraction or similar). - - -## Coordinating with NVIDIA Nsight Systems Launch - -Consult the Nsight Systems User Guide for full overview of options. - -On the PyTorch workflow, basic NVTX markers are by default provided. On the C++/TensorRT workflow, append `--nvtx` when calling `scripts/build_wheel.py` script to compile, and clean build the code. - -### Only collect specific iterations - -To reduce the Nsight Systems profile size, and to control that only specific iterations are collected, set environment variable `TLLM_PROFILE_START_STOP=A-B`, and append `-c cudaProfilerApi` to `nsys profile` command. - - -### Enable more NVTX markers for debugging -Set environment variable `TLLM_NVTX_DEBUG=1`. - -### Enable garbage collection (GC) NVTX markers -Set environment variable `TLLM_PROFILE_RECORD_GC=1`. - - -### Enable GIL information in NVTX markers -Append “python-gil” to Nsys “-t” option. - - -## Coordinating with PyTorch profiler (PyTorch workflow only) - -### Collect PyTorch profiler results -1. Set environment variable `TLLM_PROFILE_START_STOP=A-B` to specify the range of the iterations to be collected. -2. Set environment variable `TLLM_TORCH_PROFILE_TRACE=`, and the results will be saved to ``. - -### Visualize the PyTorch profiler results -Use to inspect the saved profile. - - -## Examples -Consult the Nsight Systems User Guide for full overview of MPI-related options. - -### Profiling specific iterations on a trtllm-bench/trtllm-serve run - -Say we want to profile iterations 100 to 150 on a trtllm-bench/trtllm-serve run, we want to collect as much information as possible for debugging, such as GIL, debugging NVTX markers, etc: - -```bash -#!/bin/bash - -# Prepare dataset for the benchmark -trtllm-bench \ - --model=${MODEL_PATH} prepare-dataset \ - --output /tmp/dataset.txt token-norm-dist --num-requests=${NUM_SAMPLES} \ - --input-mean=1000 --output-mean=1000 --input-stdev=0 --output-stdev=0 - -# Benchmark and profile -TLLM_PROFILE_START_STOP=100-150 nsys profile \ - -o trace -f true \ - -t 'cuda,nvtx,python-gil' -c cudaProfilerApi \ - --cuda-graph-trace node \ - -e TLLM_PROFILE_RECORD_GC=1,TLLM_LLMAPI_ENABLE_NVTX=1,TLLM_TORCH_PROFILE_TRACE=trace.json \ - --trace-fork-before-exec=true \ - trtllm-bench \ # or trtllm-serve command - --model deepseek-ai/DeepSeek-V3 \ - --model_path ${MODEL_PATH} \ - throughput \ - --dataset /tmp/dataset.txt --warmup 0 \ - --streaming -``` - -The Nsight Systems reports will be saved to `trace.nsys-rep`. Use NVIDIA Nsight Systems application to open it. - -The PyTorch profiler results will be saved to `trace.json`. Use to inspect the saved profile. diff --git a/docs/source/legacy/performance/perf-benchmarking.md b/docs/source/legacy/performance/perf-benchmarking.md deleted file mode 100644 index 4fc460596f22..000000000000 --- a/docs/source/legacy/performance/perf-benchmarking.md +++ /dev/null @@ -1,871 +0,0 @@ -(perf-benchmarking)= - -# TensorRT-LLM Benchmarking - -```{important} -This benchmarking suite is a work in progress. -Expect breaking API changes. -``` - -TensorRT-LLM provides the `trtllm-bench` CLI, a packaged benchmarking utility that aims to make it -easier for users to reproduce our officially published [performance overview](../../developer-guide/perf-overview.md#throughput-measurements). `trtllm-bench` provides the follows: - -- A streamlined way to build tuned engines for benchmarking for a variety of models and platforms. -- An entirely Python workflow for benchmarking. -- Ability to benchmark various flows and features within TensorRT-LLM. - -`trtllm-bench` executes all benchmarks using [in-flight batching] -- for more information see -the [this section](../advanced/gpt-attention.md#in-flight-batching) that describes the concept -in further detail. - -## Before Benchmarking - -For rigorous benchmarking where consistent and reproducible results are critical, proper GPU configuration is essential. These settings help maximize GPU utilization, eliminate performance variability, and ensure optimal conditions for accurate measurements. While not strictly required for normal operation, we recommend applying these configurations when conducting performance comparisons or publishing benchmark results. - -### Persistence mode -Ensure persistence mode is enabled to maintain consistent GPU state: -```shell -sudo nvidia-smi -pm 1 -``` - -### GPU Clock Management -Allow the GPU to dynamically adjust its clock speeds based on workload and temperature. While locking clocks at maximum frequency might seem beneficial, it can sometimes lead to thermal throttling and reduced performance. Reset GPU clocks using: -```shell -sudo nvidia-smi -rgc -``` - -### Set power limits -First query the maximum power limit: -```shell -nvidia-smi -q -d POWER -``` -Then configure the GPU to operate at its maximum power limit for consistent performance: -```shell -sudo nvidia-smi -pl -``` - -### Boost settings -Potentially a GPU may support boost levels. First query available boost levels: -```shell -sudo nvidia-smi boost-slider -l -``` -If supported, enable the boost slider using one of the available levels for maximum performance: -```shell -sudo nvidia-smi boost-slider --vboost -``` - - -## Throughput Benchmarking - -### Limitations and Caveats - -#### Validated Networks for Benchmarking - -While `trtllm-bench` should be able to run any network that TensorRT-LLM supports, the following are the list -that have been validated extensively and is the same listing as seen on the -[Performance Overview](../../developer-guide/perf-overview.md) page. - -- [meta-llama/Llama-2-7b-hf](https://huggingface.co/meta-llama/Llama-2-7b-hf) -- [meta-llama/Llama-2-70b-hf](https://huggingface.co/meta-llama/Llama-2-70b-hf) -- [tiiuae/falcon-180B](https://huggingface.co/tiiuae/falcon-180B) -- [EleutherAI/gpt-j-6b](https://huggingface.co/EleutherAI/gpt-j-6b) -- [meta-llama/Meta-Llama-3-8B](https://huggingface.co/meta-llama/Meta-Llama-3-8B) -- [meta-llama/Llama-3.1-8B](https://huggingface.co/meta-llama/Llama-3.1-8B) -- [meta-llama/Meta-Llama-3-70B](https://huggingface.co/meta-llama/Meta-Llama-3-70B) -- [meta-llama/Llama-3.1-70B](https://huggingface.co/meta-llama/Llama-3.1-70B) -- [meta-llama/Llama-3.1-405B](https://huggingface.co/meta-llama/Llama-3.1-405B) -- [mistralai/Mixtral-8x7B-v0.1](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1) -- [mistralai/Mistral-7B-v0.1](https://huggingface.co/mistralai/Mistral-7B-v0.1) -- [meta-llama/Llama-3.1-8B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct) -- [meta-llama/Llama-3.1-70B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-70B-Instruct) -- [meta-llama/Llama-3.1-405B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-405B-Instruct) -- [mistralai/Mixtral-8x7B-Instruct-v0.1](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1) - -```{tip} -`trtllm-bench` can automatically download the model from Hugging Face Model Hub. -Export your token in the `HF_TOKEN` environment variable. -``` - -#### Supported Quantization Modes - -`trtllm-bench` supports the following quantization modes: - -- None (no quantization applied) -- `FP8` -- `NVFP4` - -For more information about quantization, refer to [](../reference/precision.md) and -the [support matrix](../reference/precision.md#support-matrix) of the supported quantization methods for each network. - -```{tip} -Although TensorRT-LLM supports more quantization modes than listed above, `trtllm-bench` currently only configures for -a smaller subset. -``` - -### Quickstart - -This quick start focuses on running a short max throughput benchmark on -`meta-llama/Llama-3.1-8B` on a synthetic dataset with a uniform distribution of prompts with ISL:OSL -of 128:128. -To run the benchmark from start to finish, run the following commands: - -```shell -trtllm-bench --tokenizer meta-llama/Llama-3.1-8B prepare-dataset --output /tmp/synthetic_128_128.txt token-norm-dist --input-mean 128 --output-mean 128 --input-stdev 0 --output-stdev 0 --num-requests 3000 -trtllm-bench --model meta-llama/Llama-3.1-8B build --dataset /tmp/synthetic_128_128.txt --quantization FP8 -trtllm-bench --model meta-llama/Llama-3.1-8B throughput --dataset /tmp/synthetic_128_128.txt --engine_dir /tmp/meta-llama/Llama-3.1-8B/tp_1_pp_1 -``` - -After the benchmark completes, `trtllm-bench` prints a summary with summary metrics. - -```shell -=========================================================== -= ENGINE DETAILS -=========================================================== -Model: meta-llama/Llama-3.1-8B -Engine Directory: /tmp/meta-llama/Llama-3.1-8B/tp_1_pp_1 -TensorRT-LLM Version: 0.17.0 -Dtype: bfloat16 -KV Cache Dtype: FP8 -Quantization: FP8 -Max Input Length: 256 -Max Sequence Length: 256 - -=========================================================== -= WORLD + RUNTIME INFORMATION -=========================================================== -TP Size: 1 -PP Size: 1 -Max Runtime Batch Size: 4096 -Max Runtime Tokens: 8192 -Scheduling Policy: Guaranteed No Evict -KV Memory Percentage: 90.00% -Issue Rate (req/sec): 5.0689E+14 - -=========================================================== -= PERFORMANCE OVERVIEW -=========================================================== -Number of requests: 3000 -Average Input Length (tokens): 128.0000 -Average Output Length (tokens): 128.0000 -Token Throughput (tokens/sec): 28390.4265 -Request Throughput (req/sec): 221.8002 -Total Latency (ms): 13525.6862 - -=========================================================== -``` - -### Workflow - -The workflow for `trtllm-bench` is composed of the following steps: - -1. Prepare a dataset to drive the inflight batching benchmark. -2. Build a benchmark engine using `trtllm-bench build` subcommand (not required for [PyTorch flow](#running-with-the-pytorch-workflow)). -3. Run the max throughput benchmark using the `trtllm-bench throughput` subcommand or low latency benchmark using the `trtllm-bench latency` subcommand. - - -#### Preparing a Dataset - -The throughput benchmark utilizes a fixed JSON schema to specify requests. The schema is defined as follows: - -| Key | Required | Type | Description | -| :-------------- | :------: | :-----------: | :---------------------------------------------- | -| `task_id` | Y | String | Unique identifier for the request. | -| `prompt` | N* | String | Input text for a generation request. | -| `input_ids` | Y* | List[Integer] | List of token IDs that make up the request prompt. | -| `output_tokens` | Y | Integer | Number of generated tokens for this request. | - -```{tip} -\* Specifying `prompt` or `input_ids` is required. However, you cannot have both prompts and token IDs (`input_ids`) -defined at the same time. If you specify `input_ids`, the `prompt` entry is ignored for request generation. -``` - -Refer to the following examples of valid entries for the benchmark: - -- Entries with a human-readable prompt and no logits. - - ```json - {"task_id": 1, "prompt": "Generate an infinite response to the following: This is the song that never ends, it goes on and on my friend.", "output_tokens": 1000} - {"task_id": 2, "prompt": "Generate an infinite response to the following: Na, na, na, na", "output_tokens": 1000} - ``` - -- Entries which contain logits. - - ```json - {"task_id":0,"input_ids":[863,22056,25603,11943,8932,13195,3132,25032,21747,22213],"output_tokens":128} - {"task_id":1,"input_ids":[14480,13598,15585,6591,1252,8259,30990,26778,7063,30065,21764,11023,1418],"output_tokens":128} - ``` - -```{tip} -Specify each entry on one line. -To simplify passing the data, a complete JSON entry is on each line so that the benchmarker -can simply read a line and assume a complete entry. When creating a dataset, be sure that a complete -JSON entry is on every line. -``` - -In order to prepare a synthetic dataset, you can use the provided script in the `benchmarks/cpp` -directory. For example, to generate a synthetic dataset of 1000 requests with a uniform ISL/OSL of -128/128 for [meta-llama/Llama-3.1-8B](https://huggingface.co/meta-llama/Llama-3.1-8B), run: - -```shell -trtllm-bench --tokenizer meta-llama/Llama-3.1-8B prepare-dataset --output /tmp/synthetic_128_128.txt token-norm-dist --input-mean 128 --output-mean 128 --input-stdev 0 --output-stdev 0 --num-requests 3000 -``` - -### Building a Benchmark Engine - -#### Default Build Behavior -The `trtllm-bench` CLI tool provides the `build` subcommand to build the TRT-LLM engines for max throughput benchmark. -To build an engine for benchmarking, you can specify the dataset generated with `prepare_dataset.py` through `--dataset` option. -By default, `trtllm-bench`'s tuning heuristic uses the high-level statistics of the dataset (average ISL/OSL, max sequence length) -to optimize engine build settings. The following command builds an FP8 quantized engine optimized using the dataset's ISL/OSL. - -```shell -trtllm-bench --model meta-llama/Llama-3.1-8B build --quantization FP8 --dataset /tmp/synthetic_128_128.txt -``` - -#### Other Build Modes - -The build subcommand also provides other ways to build the engine where users have larger control over the tuning values. - -- Build engine with self-defined tuning values: -You specify the tuning values to build the engine with by setting `--max_batch_size` and `--max_num_tokens` directly. -`max_batch_size` and `max_num_tokens` control the maximum number of requests and tokens that can be scheduled in each iteration. -If no value is specified, the default `max_batch_size` and `max_num_tokens` values of `2048` and `8192` are used. -The following command builds an FP8 quantized engine by specifying the engine tuning values. - -```shell -trtllm-bench --model meta-llama/Llama-3.1-8B build --quantization FP8 --max_seq_len 4096 --max_batch_size 1024 --max_num_tokens 2048 -``` - - -#### Parallelism Mapping Support -The `trtllm-bench build` subcommand supports combinations of tensor-parallel (TP) and pipeline-parallel (PP) mappings as long as the world size (`tp_size x pp_size`) `<=` `8`. The parallelism mapping in build subcommad is controlled by `--tp_size` and `--pp_size` options. The following command builds an engine with TP2-PP2 mapping. - -```shell -trtllm-bench --model meta-llama/Llama-3.1-8B build --quantization FP8 --dataset /tmp/synthetic_128_128.txt --tp_size 2 --pp_size 2 -``` - - -#### Example of Build Subcommand Output: -The output of the `build` subcommand looks similar to the snippet below (for `meta-llama/Llama-3.1-8B`): - -```shell -user@387b12598a9e:/scratch/code/trt-llm/tekit_2025$ trtllm-bench --model meta-llama/Llama-3.1-8B build --dataset /tmp/synthetic_128_128.txt --quantization FP8 -[TensorRT-LLM] TensorRT-LLM version: 0.17.0 -[01/18/2025-00:55:14] [TRT-LLM] [I] Found dataset. -[01/18/2025-00:55:14] [TRT-LLM] [I] -=========================================================== -= DATASET DETAILS -=========================================================== -Max Input Sequence Length: 128 -Max Output Sequence Length: 128 -Max Sequence Length: 256 -Target (Average) Input Sequence Length: 128 -Target (Average) Output Sequence Length: 128 -Number of Sequences: 3000 -=========================================================== - - -[01/18/2025-00:55:14] [TRT-LLM] [I] Max batch size and max num tokens are not provided, use tuning heuristics or pre-defined setting from trtllm-bench. -[01/18/2025-00:55:14] [TRT-LLM] [I] Estimated total available memory for KV cache: 132.37 GB -[01/18/2025-00:55:14] [TRT-LLM] [I] Estimated total KV cache memory: 125.75 GB -[01/18/2025-00:55:14] [TRT-LLM] [I] Estimated max number of requests in KV cache memory: 8048.16 -[01/18/2025-00:55:14] [TRT-LLM] [I] Estimated max batch size (after fine-tune): 4096 -[01/18/2025-00:55:14] [TRT-LLM] [I] Estimated max num tokens (after fine-tune): 8192 -[01/18/2025-00:55:14] [TRT-LLM] [I] Set dtype to bfloat16. -[01/18/2025-00:55:14] [TRT-LLM] [I] Set multiple_profiles to True. -[01/18/2025-00:55:14] [TRT-LLM] [I] Set use_paged_context_fmha to True. -[01/18/2025-00:55:14] [TRT-LLM] [I] Set use_fp8_context_fmha to True. -[01/18/2025-00:55:14] [TRT-LLM] [I] -=========================================================== -= ENGINE BUILD INFO -=========================================================== -Model Name: meta-llama/Llama-3.1-8B -Model Path: None -Workspace Directory: /tmp -Engine Directory: /tmp/meta-llama/Llama-3.1-8B/tp_1_pp_1 - -=========================================================== -= ENGINE CONFIGURATION DETAILS -=========================================================== -Max Sequence Length: 256 -Max Batch Size: 4096 -Max Num Tokens: 8192 -Quantization: FP8 -KV Cache Dtype: FP8 -=========================================================== - -Loading Model: [1/3] Downloading HF model -Downloaded model to /data/models--meta-llama--Llama-3.1-8B/snapshots/d04e592bb4f6aa9cfee91e2e20afa771667e1d4b -Time: 0.321s -Loading Model: [2/3] Loading HF model to memory -Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:59<00:00, 14.79s/it] -Generating train split: 100%|████████████████████████████████████████████████████████████████████████████████████| 287113/287113 [00:06<00:00, 41375.57 examples/s] -Generating validation split: 100%|█████████████████████████████████████████████████████████████████████████████████| 13368/13368 [00:00<00:00, 41020.63 examples/s] -Generating test split: 100%|███████████████████████████████████████████████████████████████████████████████████████| 11490/11490 [00:00<00:00, 41607.11 examples/s] -Inserted 675 quantizers -/usr/local/lib/python3.12/dist-packages/modelopt/torch/quantization/model_quant.py:71: DeprecationWarning: forward_loop should take model as argument, but got forward_loop without any arguments. This usage will be deprecated in future versions. - warnings.warn( -Disable lm_head quantization for TRT-LLM export due to deployment limitations. -current rank: 0, tp rank: 0, pp rank: 0 -Time: 122.568s -Loading Model: [3/3] Building TRT-LLM engine -/usr/local/lib/python3.12/dist-packages/tensorrt/__init__.py:85: DeprecationWarning: Context managers for TensorRT types are deprecated. Memory will be freed automatically when the reference count reaches 0. - warnings.warn( -Time: 53.820s -Loading model done. -Total latency: 176.709s - - - -=========================================================== -ENGINE SAVED: /tmp/meta-llama/Llama-3.1-8B/tp_1_pp_1 -=========================================================== -``` - -The engine in this case will be written to `/tmp/meta-llama/Llama-3.1-8B/tp_1_pp_1` (the end of the log). - - -### Max Throughput Benchmark - -The `trtllm-bench` command line tool provides a max throughput benchmark that is accessible via the -`throughput` subcommand. This benchmark tests a TensorRT-LLM engine or PyTorch backend under maximum load to provide an -upper bound throughput number. - -#### How the Benchmarker Works - -The benchmarker reads a data file where a single line contains -a complete JSON request entry as specified in [](#preparing-a-dataset). -The process that the benchmarker follows is: - -1. Iterate over all input requests. If `logits` is specified, construct the request using the specified -list of logits. Otherwise, tokenize the `prompt` as specified by `--model $HF_MODEL_NAME`. -1. Submit the dataset to the TensorRT-LLM `Executor` API as fast as possible (offline mode). -1. Wait for all requests to return, compute statistics, and then report results. - -To run the benchmarker, run the following commands with the [engine](#building-a-benchmark-engine) and -[dataset](#preparing-a-dataset) generated from previous steps: - -```shell -trtllm-bench --model meta-llama/Llama-3.1-8B throughput --dataset /tmp/synthetic_128_128.txt --engine_dir /tmp/meta-llama/Llama-3.1-8B/tp_1_pp_1 -[TensorRT-LLM] TensorRT-LLM version: 0.17.0 -[01/18/2025-01:01:13] [TRT-LLM] [I] Preparing to run throughput benchmark... -[01/18/2025-01:01:13] [TRT-LLM] [I] Setting up throughput benchmark. - - - -[01/18/2025-01:01:26] [TRT-LLM] [I] Setting up for warmup... -[01/18/2025-01:01:26] [TRT-LLM] [I] Running warmup. -[01/18/2025-01:01:26] [TRT-LLM] [I] Starting benchmarking async task. -[01/18/2025-01:01:26] [TRT-LLM] [I] Starting benchmark... -[01/18/2025-01:01:26] [TRT-LLM] [I] Request submission complete. [count=2, time=0.0000s, rate=121847.20 req/s] -[01/18/2025-01:01:28] [TRT-LLM] [I] Benchmark complete. -[01/18/2025-01:01:28] [TRT-LLM] [I] Stopping LLM backend. -[01/18/2025-01:01:28] [TRT-LLM] [I] Cancelling all 0 tasks to complete. -[01/18/2025-01:01:28] [TRT-LLM] [I] All tasks cancelled. -[01/18/2025-01:01:28] [TRT-LLM] [I] LLM Backend stopped. -[01/18/2025-01:01:28] [TRT-LLM] [I] Warmup done. -[01/18/2025-01:01:28] [TRT-LLM] [I] Starting benchmarking async task. -[01/18/2025-01:01:28] [TRT-LLM] [I] Starting benchmark... -[01/18/2025-01:01:28] [TRT-LLM] [I] Request submission complete. [count=3000, time=0.0012s, rate=2590780.97 req/s] -[01/18/2025-01:01:42] [TRT-LLM] [I] Benchmark complete. -[01/18/2025-01:01:42] [TRT-LLM] [I] Stopping LLM backend. -[01/18/2025-01:01:42] [TRT-LLM] [I] Cancelling all 0 tasks to complete. -[01/18/2025-01:01:42] [TRT-LLM] [I] All tasks cancelled. -[01/18/2025-01:01:42] [TRT-LLM] [I] LLM Backend stopped. -[01/18/2025-01:01:42] [TRT-LLM] [I] - -=========================================================== -= ENGINE DETAILS -=========================================================== -Model: meta-llama/Llama-3.1-8B -Engine Directory: /tmp/meta-llama/Llama-3.1-8B/tp_1_pp_1 -TensorRT-LLM Version: 0.17.0 -Dtype: bfloat16 -KV Cache Dtype: FP8 -Quantization: FP8 -Max Input Length: 256 -Max Sequence Length: 256 - -=========================================================== -= WORLD + RUNTIME INFORMATION -=========================================================== -TP Size: 1 -PP Size: 1 -Max Runtime Batch Size: 4096 -Max Runtime Tokens: 8192 -Scheduling Policy: Guaranteed No Evict -KV Memory Percentage: 90.00% -Issue Rate (req/sec): 5.0689E+14 - -=========================================================== -= PERFORMANCE OVERVIEW -=========================================================== -Number of requests: 3000 -Average Input Length (tokens): 128.0000 -Average Output Length (tokens): 128.0000 -Token Throughput (tokens/sec): 28390.4265 -Request Throughput (req/sec): 221.8002 -Total Latency (ms): 13525.6862 - -=========================================================== - -[01/18/2025-01:01:42] [TRT-LLM] [I] Thread proxy_dispatch_result_thread stopped. -[TensorRT-LLM][INFO] Refreshed the MPI local session -``` - -### Running with the PyTorch Workflow - -```{eval-rst} -.. include:: ../../_includes/note_sections.rst - :start-after: .. start-note-config-flag-alias - :end-before: .. end-note-config-flag-alias -``` - -To benchmark the PyTorch backend (`tensorrt_llm._torch`), use the following command with [dataset](#preparing-a-dataset) generated from previous steps. With the PyTorch flow, you will not need to -run `trtllm-bench build`; the `throughput` benchmark initializes the backend by tuning against the -dataset provided via `--dataset` (or the other build mode settings described [above](#other-build-modes)). -Note that CUDA graph is enabled by default. You can add additional pytorch config with -`--config` followed by the path to a YAML file. For more details, please refer to the -help text by running the command with `--help`. - -```{tip} -The command below specifies the `--model_path` option. The model path is optional and used only when you want to run a locally -stored checkpoint. When using `--model_path`, the `--model` is still required for reporting reasons and in order to look up parameters -for build heuristics. -``` - -```shell -trtllm-bench --model meta-llama/Llama-3.1-8B --model_path /Ckpt/Path/To/Llama-3.1-8B throughput --dataset /tmp/synthetic_128_128.txt - -# Example output - -=========================================================== -= PyTorch backend -=========================================================== -Model: meta-llama/Llama-3.1-8B -Model Path: /Ckpt/Path/To/Llama-3.1-8B -TensorRT-LLM Version: 0.17.0 -Dtype: bfloat16 -KV Cache Dtype: None -Quantization: FP8 - -=========================================================== -= WORLD + RUNTIME INFORMATION -=========================================================== -TP Size: 1 -PP Size: 1 -Max Runtime Batch Size: 2048 -Max Runtime Tokens: 4096 -Scheduling Policy: Guaranteed No Evict -KV Memory Percentage: 90.00% -Issue Rate (req/sec): 7.6753E+14 - -=========================================================== -= PERFORMANCE OVERVIEW -=========================================================== -Number of requests: 3000 -Average Input Length (tokens): 128.0000 -Average Output Length (tokens): 128.0000 -Token Throughput (tokens/sec): 20685.5510 -Request Throughput (req/sec): 161.6059 -Total Latency (ms): 18563.6825 - -``` - -#### Benchmarking with LoRA Adapters in PyTorch workflow - -The PyTorch workflow supports benchmarking with LoRA (Low-Rank Adaptation) adapters. This requires preparing a dataset with LoRA metadata and configuring the LoRA settings. - -**Preparing LoRA Dataset** - -Use `prepare_dataset.py` with LoRA-specific options to generate requests with LoRA metadata: - -```shell -python3 benchmarks/cpp/prepare_dataset.py \ - --stdout \ - --rand-task-id 0 1 \ - --tokenizer /path/to/tokenizer \ - --lora-dir /path/to/loras \ - token-norm-dist \ - --num-requests 100 \ - --input-mean 128 \ - --output-mean 128 \ - --input-stdev 16 \ - --output-stdev 24 \ - > synthetic_lora_data.json -``` - -Key LoRA options: -- `--lora-dir`: Parent directory containing LoRA adapter subdirectories named by their task IDs (e.g., `0/`, `1/`, etc.) -- `--rand-task-id`: Range of LoRA task IDs to randomly assign to requests -- `--task-id`: Fixed LoRA task ID for all requests (alternative to `--rand-task-id`) - -The generated dataset will include LoRA request metadata. Below is an example of a single such request data entry: - -```json -{ - "task_id": 0, - "input_ids": [3452, 88226, 102415, ...], - "output_tokens": 152, - "lora_request": { - "lora_name": "lora_0", - "lora_int_id": 0, - "lora_path": "/path/to/loras/0" - } -} -``` - -**LoRA Configuration** - -Create a `config.yaml` file with LoRA configuration: - -```yaml -lora_config: - lora_dir: - - /path/to/loras/0 - - /path/to/loras/1 - max_lora_rank: 64 - lora_target_modules: - - attn_q - - attn_k - - attn_v - trtllm_modules_to_hf_modules: - attn_q: q_proj - attn_k: k_proj - attn_v: v_proj -``` - -**Running LoRA Benchmark** - -```shell -trtllm-bench --model /path/to/base/model \ - throughput \ - --dataset synthetic_lora_data.json \ - --config config.yaml -``` - -```{note} -The LoRA directory structure should have task-specific subdirectories named by their task IDs (e.g., `loras/0/`, `loras/1/`). -Each subdirectory should contain the LoRA adapter files for that specific task. -``` - -#### Running multi-modal models in the PyTorch Workflow - -To benchmark multi-modal models with PyTorch workflow, you can follow the similar approach as above. - -First, prepare the dataset: -``` -python ./benchmarks/cpp/prepare_dataset.py \ - --tokenizer Qwen/Qwen2-VL-2B-Instruct \ - --stdout \ - dataset \ - --dataset-name lmms-lab/MMMU \ - --dataset-split test \ - --dataset-image-key image \ - --dataset-prompt-key question \ - --num-requests 10 \ - --output-len-dist 128,5 > mm_data.jsonl -``` -It will download the media files to `/tmp` directory and prepare the dataset with their paths. Note that the `prompt` fields are texts and not tokenized ids. This is due to the fact that -the `prompt` and the media (image/video) are processed by a preprocessor for multimodal files. - -Sample dataset for multimodal: -``` -{"task_id":0,"prompt":"Brahma Industries sells vinyl replacement windows to home improvement retailers nationwide. The national sales manager believes that if they invest an additional $25,000 in advertising, they would increase sales volume by 10,000 units. What is the total contribution margin?","media_paths":["/tmp/tmp9so41y3r.jpg"],"output_tokens":126} -{"task_id":1,"prompt":"Let us compute for the missing amounts under work in process inventory, what is the cost of goods manufactured? ","media_paths":["/tmp/tmpowsrb_f4.jpg"],"output_tokens":119} -{"task_id":2,"prompt":"Tsuji is reviewing the price of a 3-month Japanese yen/U.S. dollar currency futures contract, using the currency and interest rate data shown below. Because the 3-month Japanese interest rate has just increased to .50%, Itsuji recognizes that an arbitrage opportunity exists nd decides to borrow $1 million U.S. dollars to purchase Japanese yen. Calculate the yen arbitrage profit from Itsuji's strategy, using the following data: ","media_paths":["/tmp/tmpxhdvasex.jpg"],"output_tokens":126} -... -``` - -Run the benchmark: -``` -trtllm-bench --model Qwen/Qwen2-VL-2B-Instruct \ - throughput \ - --dataset mm_data.jsonl \ - --num_requests 10 \ - --max_batch_size 4 \ - --modality image -``` - - -Sample output: -``` -=========================================================== -= REQUEST DETAILS -=========================================================== -Number of requests: 10 -Number of concurrent requests: 5.3019 -Average Input Length (tokens): 411.6000 -Average Output Length (tokens): 128.7000 -=========================================================== -= WORLD + RUNTIME INFORMATION -=========================================================== -TP Size: 1 -PP Size: 1 -EP Size: None -Max Runtime Batch Size: 4 -Max Runtime Tokens: 12288 -Scheduling Policy: GUARANTEED_NO_EVICT -KV Memory Percentage: 90.00% -Issue Rate (req/sec): 1.4117E+17 - -=========================================================== -= PERFORMANCE OVERVIEW -=========================================================== -Request Throughput (req/sec): 1.4439 -Total Output Throughput (tokens/sec): 185.8351 -Per User Output Throughput (tokens/sec/user): 38.1959 -Per GPU Output Throughput (tokens/sec/gpu): 185.8351 -Total Token Throughput (tokens/sec): 780.1607 -Total Latency (ms): 6925.4963 -Average request latency (ms): 3671.8441 - --- Request Latency Breakdown (ms) ----------------------- - -[Latency] P50 : 3936.3022 -[Latency] P90 : 5514.4701 -[Latency] P95 : 5514.4701 -[Latency] P99 : 5514.4701 -[Latency] MINIMUM: 2397.1047 -[Latency] MAXIMUM: 5514.4701 -[Latency] AVERAGE: 3671.8441 - -=========================================================== -= DATASET DETAILS -=========================================================== -Dataset Path: /workspaces/tensorrt_llm/mm_data.jsonl -Number of Sequences: 10 - --- Percentiles statistics --------------------------------- - - Input Output Seq. Length ------------------------------------------------------------ -MIN: 167.0000 119.0000 300.0000 -MAX: 1059.0000 137.0000 1178.0000 -AVG: 411.6000 128.7000 540.3000 -P50: 299.0000 128.0000 427.0000 -P90: 1059.0000 137.0000 1178.0000 -P95: 1059.0000 137.0000 1178.0000 -P99: 1059.0000 137.0000 1178.0000 -=========================================================== -``` - -**Notes and Limitations**: -- Only image datasets are supported for now. -- `--output-len-dist` is a required argument for multimodal datasets. -- Tokenizer is unused during the prepare step but it is still a required argument. -- Since the images are converted to tokens when the model is run, `trtllm-bench` uses a default large value for the maximum input sequence length when setting up the execution settings. - You can also modify the behavior by specifying a different value with the flag `--max_input_len` that suits your use-case. - -#### Quantization in the PyTorch Flow - -In order to run a quantized run with `trtllm-bench` utilizing the PyTorch flow, you will need to use a pre-quantized -To run a quantized benchmark with `trtllm-bench` utilizing the PyTorch flow, you will need to use a pre-quantized -checkpoint. For the Llama-3.1 models, TensorRT-LLM provides the following checkpoints via HuggingFace: - -- [`nvidia/Llama-3.1-8B-Instruct-FP8`](https://huggingface.co/nvidia/Llama-3.1-8B-Instruct-FP8) -- [`nvidia/Llama-3.1-70B-Instruct-FP8`](https://huggingface.co/nvidia/Llama-3.1-70B-Instruct-FP8) -- [`nvidia/Llama-3.1-405B-Instruct-FP8`](https://huggingface.co/nvidia/Llama-3.1-405B-Instruct-FP8) - -`trtllm-bench` utilizes the `hf_quant_config.json` file present in the pre-quantized checkpoints above. The configuration -file is present in checkpoints quantized with [Model Optimizer](https://github.com/NVIDIA/Model-Optimizer) -and describes the compute and KV cache quantization that checkpoint was compiled with. For example, from the checkpoints -above: - -```json -{ - "producer": { - "name": "modelopt", - "version": "0.23.0rc1" - }, - "quantization": { - "quant_algo": "FP8", - "kv_cache_quant_algo": null - } -} -``` - -The checkpoints above are quantized to run with a compute precision of `FP8` and default to no KV cache quantization (full -`FP16` cache). When running `trtllm-bench throughput`. The benchmark will select a KV cache quantization that is best suited -for the compute precision in the checkpoint automatically if `kv_cache_quant_algo` is specified as `null`, otherwise it will -be forced to match the specified non-null KV cache quantization. The following are the mappings that `trtllm-bench` will -follow when a checkpoint does not specify a KV cache quantization algorithm: - -| Checkpoint Compute Quant | Checkpoint KV Cache Quant | `trtllm-bench` | Note | -| - | - | - | - | -| `null` | `null` | `null` | In this case, a quantization config doesn't exist. | -| `FP8` | `FP8` | `FP8` | Matches the checkpoint | -| `FP8` | `null` | `FP8` | Set to `FP8` via benchmark | -| `NVFP4` | `null` | `FP8` | Set to `FP8` via benchmark | - -If you would like to force the KV cache quantization, you can specify the following in the YAML file to force the precision -when the checkpoint precision is `null`: - -```yaml -kv_cache_dtype: "fp8" -``` - -```{tip} -The valid values for `kv_cache_dtype` are `auto`, `fp8`, and `nvfp4`. -``` - -## Low Latency Benchmark - -The low latency benchmark follows a similar workflow to the [throughput benchmark](#max-throughput-benchmark) -but requires building the engine separately from `trtllm-bench`. Low latency benchmarks has the following modes: - -- A single-request low-latency engine -- A Medusa-enabled speculative-decoding engine - -### Low Latency TensorRT-LLM Engine for Llama-3 70B - -To build a low-latency engine for the latency benchmark, run the following quantize and build commands. -The `$checkpoint_dir` is the path to the [meta-llama/Meta-Llama-3-70B](https://huggingface.co/meta-llama/Meta-Llama-3-70B) Hugging Face checkpoint in your cache or downloaded to a specific location with the [huggingface-cli](https://huggingface.co/docs/huggingface_hub/en/guides/cli). -To prepare a dataset, follow the same process as specified in [](#preparing-a-dataset). - -#### Benchmarking a non-Medusa Low Latency Engine - -To quantize the checkpoint: - -```shell -cd tensorrt_llm/examples/models/core/llama -python ../quantization/quantize.py \ - --model_dir $checkpoint_dir \ - --dtype bfloat16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir /tmp/meta-llama/Meta-Llama-3-70B/checkpoint \ - --calib_size 512 \ - --tp_size $tp_size -``` - -then build, - -```shell -trtllm-build \ - --checkpoint_dir /tmp/meta-llama/Meta-Llama-3-70B/checkpoint \ - --use_fused_mlp enable \ - --gpt_attention_plugin bfloat16 \ - --output_dir /tmp/meta-llama/Meta-Llama-3-70B/engine \ - --max_batch_size 1 \ - --max_seq_len $(($isl+$osl)) \ - --reduce_fusion enable \ - --gemm_plugin fp8 \ - --workers $tp_size \ - --use_fp8_context_fmha enable \ - --max_num_tokens $isl \ - --use_paged_context_fmha disable \ - --multiple_profiles enable -``` - -After the engine is built, run the low-latency benchmark: - -```shell -env TRTLLM_ENABLE_MMHA_MULTI_BLOCK_DEBUG=1 \ - TRTLLM_MMHA_KERNEL_BLOCK_SIZE=256 \ - TRTLLM_MMHA_BLOCKS_PER_SEQUENCE=32 \ - FORCE_MULTI_BLOCK_MODE=ON \ - TRTLLM_ENABLE_PDL=1 \ - trtllm-bench --model meta-llama/Meta-Llama-3-70B \ - latency \ - --dataset $DATASET_PATH \ - --engine_dir /tmp/meta-llama/Meta-Llama-3-70B/engine -``` - -### Building a Medusa Low-Latency Engine - -To build a Medusa-enabled engine requires checkpoints that contain Medusa heads. -NVIDIA provides TensorRT-LLM checkpoints on the [NVIDIA](https://huggingface.co/nvidia) page on Hugging Face. -The checkpoints are pre-quantized and can be directly built after downloading them with the -[huggingface-cli](https://huggingface.co/docs/huggingface_hub/en/guides/cli). -After you download the checkpoints, run the following command. Make sure to -specify the `$tp_size` supported by your Medusa checkpoint and the path to its stored location `$checkpoint_dir`. -Additionally, `$max_seq_len` should be set to the model's maximum position embedding. - -Using Llama-3.1 70B as an example, for a tensor parallel 8 and bfloat16 dtype: - -```shell -tp_size=8 -max_seq_len=131072 -trtllm-build --checkpoint_dir $checkpoint_dir \ - --speculative_decoding_mode medusa \ - --max_batch_size 1 \ - --gpt_attention_plugin bfloat16 \ - --max_seq_len $max_seq_len \ - --output_dir /tmp/meta-llama/Meta-Llama-3.1-70B/medusa/engine \ - --use_fused_mlp enable \ - --paged_kv_cache enable \ - --use_paged_context_fmha disable \ - --multiple_profiles enable \ - --reduce_fusion enable \ - --use_fp8_context_fmha enable \ - --workers $tp_size \ - --low_latency_gemm_plugin fp8 -``` - -After the engine is built, you need to define the Medusa choices. -The choices are specified with a YAML file like the following example (`medusa.yaml`): - -```yaml -- [0] -- [0, 0] -- [1] -- [0, 1] -- [2] -- [0, 0, 0] -- [1, 0] -- [0, 2] -- [3] -- [0, 3] -- [4] -- [0, 4] -- [2, 0] -- [0, 5] -- [0, 0, 1] -``` - -To run the Medusa-enabled engine, run the following command: - -```shell -env TRTLLM_ENABLE_PDL=1 \ - UB_ONESHOT=1 \ - UB_TP_SIZE=$tp_size \ - TRTLLM_ENABLE_PDL=1 \ - TRTLLM_PDL_OVERLAP_RATIO=0.15 \ - TRTLLM_PREFETCH_RATIO=-1 \ - trtllm-bench --model meta-llama/Meta-Llama-3-70B \ - latency \ - --dataset $DATASET_PATH \ - --engine_dir /tmp/meta-llama/Meta-Llama-3-70B/medusa/engine \ - --medusa_choices medusa.yml -``` - -## Summary - -The following table summarizes the commands needed for running benchmarks: - -| Scenario | Phase | Command | -| - | - | - | -| Dataset | Preparation | `python benchmarks/cpp/prepare_dataset.py --stdout --tokenizer $HF_MODEL token-norm-dist --input-mean $ISL --output-mean $OSL --input-stdev 0 --output-stdev 0 --num-requests $NUM_REQUESTS > $DATASET_PATH` | -| Throughput | Build | `trtllm-bench --model $HF_MODEL build --dataset $DATASET_PATH` | -| Throughput | Benchmark | `trtllm-bench --model $HF_MODEL throughput --dataset $DATASET_PATH --engine_dir $ENGINE_DIR` | -| Latency | Build | See [section about building low latency engines](#low-latency-tensorrt-llm-engine-for-llama-3-70b) | -| Non-Medusa Latency | Benchmark | `trtllm-bench --model $HF_MODEL latency --dataset $DATASET_PATH --engine_dir $ENGINE_DIR` | -| Medusa Latency | Benchmark | `trtllm-bench --model $HF_MODEL latency --dataset $DATASET_PATH --engine_dir $ENGINE_DIR --medusa_choices $MEDUSA_CHOICES` | - -where, - -`$HF_MODEL` -: The Hugging Face name of a model. - -`$NUM_REQUESTS` -: The number of requests to generate. - -`$DATASET_PATH` -: The path where the dataset was written when preparing the dataset. - -`$ENGINE_DIR` -: The engine directory as printed by `trtllm-bench build`. - -`$MEDUSA_CHOICES` -: A YAML config representing the Medusa tree for the benchmark. diff --git a/docs/source/legacy/performance/performance-tuning-guide/benchmarking-default-performance.md b/docs/source/legacy/performance/performance-tuning-guide/benchmarking-default-performance.md deleted file mode 100644 index 17cb9aef45ba..000000000000 --- a/docs/source/legacy/performance/performance-tuning-guide/benchmarking-default-performance.md +++ /dev/null @@ -1,243 +0,0 @@ -(benchmarking-default-performance)= - -# Benchmarking Default Performance - -This section discusses how to build an engine for the model using the LLM-API and benchmark it using TRTLLM-Bench. - -> Disclaimer: While performance numbers shown here are real, they are only for demonstration purposes. Differences in environment, SKU, interconnect, and workload can all significantly affect performance and lead to your results differing from what is shown here. - -## Before You Begin: TensorRT-LLM LLM-API - -TensorRT-LLM's LLM-API aims to make getting started with TensorRT-LLM quick and easy. For example, the following script instantiates `Llama-3.3-70B-Instruct` and runs inference on a small set of prompts. For those familiar with TensorRT-LLM's [CLI workflow](./benchmarking-default-performance.md#building-and-saving-engines-via-cli), the call to `LLM()` handles converting the model checkpoint and building the engine in one line. - -```python -#quickstart.py -from tensorrt_llm import LLM, SamplingParams - - -def main(): - prompts = [ - "Hello, I am", - "The president of the United States is", - "The capital of France is", - "The future of AI is", - ] - - sampling_params = SamplingParams(temperature=0.8, top_p=0.95) - - llm = LLM( - model="meta-llama/Llama-3.3-70B-Instruct", #HuggingFace model name, no need to download the checkpoint beforehand - tensor_parallel_size=4 - ) - - outputs = llm.generate(prompts, sampling_params) - - # Print the outputs. - for output in outputs: - prompt = output.prompt - generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - -if __name__ == '__main__': - main() -``` -### Troubleshooting Tips and Pitfalls To Avoid - -Since we are running on multiple GPUs, MPI is used to spawn processes for each GPU. This raises the following requirements - -1. The entrypoint to the script should be guarded via `if __name__ == '__main__'`. This requirement comes from mpi4py. -2. Depending on your environment, it might be required to wrap the `python` command with `mpirun`. For example the command to run the script above could be `mpirun -n 1 --oversubscribe --allow-run-as-root python quickstart.py`. For running on multiple GPUs on one node like the example is attempting to do it is usually not required to prefix with `mpirun` but if you are getting MPI errors then you should add it. Additionally, the `-n 1` which says just to launch one process is intentional as TensorRT-LLM handles spawning the processes for the remaining GPUs -3. If you get a HuggingFace access error when loading the Llama weights, this is likely because the model is gated. Request access on the HuggingFace page for the model. Then follow the instructions on [Huggingface's quickstart guide](https://huggingface.co/docs/huggingface_hub/en/quick-start#authentication) to authenticate in your environment. - - -## Building and Saving the Engine - -Save the engine using `.save()`. Just like the previous example, this script and all subsequent scripts might need to be run via `mpirun`. - -```python -from tensorrt_llm import LLM - -def main(): - llm = LLM( - model="/scratch/Llama-3.3-70B-Instruct", - tensor_parallel_size=4 - ) - - llm.save("baseline") - -if __name__ == '__main__': - main() -``` - -### Building and Saving Engines via CLI - -TensorRT-LLM also has a command line interface for building and saving engines. This workflow consists of two steps - -1. Convert model checkpoint (HuggingFace, Nemo) to TensorRT-LLM checkpoint via `convert_checkpoint.py`. Each supported model has a `convert_checkpoint.py` associated it with it and can be found in the examples folder. For example, the `convert_checkpoint.py` script for Llama models can be found [here](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/llama/convert_checkpoint.py) -2. Build engine by passing TensorRT-LLM checkpoint to `trtllm-build` command. The `trtllm-build` command is installed automatically when the `tensorrt_llm` package is installed. - -The README in the examples folder for supported models walks through building engines using this flow for a wide variety of situations. The examples folder for Llama models can be found at [https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/llama](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/llama). - -## Benchmarking with `trtllm-bench` - -`trtllm-bench` provides a command line interface for benchmarking the throughput and latency of saved engines. - -### Prepare Dataset - -`trtllm-bench` expects to be passed in a dataset of requests to run through the model. This guide creates a dummy dataset of 1000 requests with every request having input and output sequence length of 2048. TensorRT-LLM provides the `prepare_dataset.py` script to produce the dataset. To use it clone the TensorRT-LLM Repo and run the following command: - -`python benchmarks/cpp/prepare_dataset.py --stdout --tokenizer /path/to/hf/Llama-3.3-70B-Instruct/ token-norm-dist --input-mean 2048 --output-mean 2048 --input-stdev 0 --output-stdev 0 --num-requests 1000 > synthetic_2048_2048.txt` - -`trtllm-bench` can also take in real data, see [`trtllm-bench` documentation](../perf-benchmarking.md) for more details on the required format. - -### Running Throughput and Latency Benchmarks - - To benchmark the baseline engine built in the previous script, run the following commands. Again, due to the multi-gpu nature of the workload you may need prefix the `trtllm-bench` command with `mpirun -n 1 --oversubscribe --allow-run-as-root`. - -**Throughput** - -```bash -trtllm-bench \ ---model /path/to/hf/Llama-3.3-70B-Instruct/ \ -throughput \ ---dataset /path/to/dataset/synthetic_2048_2048_1000.txt \ ---engine_dir /path/to/engines/baseline #replace baseline with name used in llm.save() -``` - -This command will send all 1000 requests to the model immediately. Run `trtllm-bench throughput -h` to see a list of options that help you control the request rate and cap the total number of requests if the benchmark is taking too long. For reference, internal testing of the above command took around 20 minutes on a 4 NVLink connected H100-sxm-80GB. - -Running this command will provide a throughput overview like this: - -```bash -=========================================================== -= ENGINE DETAILS -=========================================================== -Model: /scratch/Llama-3.3-70B-Instruct/ -Engine Directory: /scratch/grid_search_engines/baseline -TensorRT-LLM Version: 0.16.0 -Dtype: bfloat16 -KV Cache Dtype: None -Quantization: None -Max Sequence Length: 131072 - -=========================================================== -= WORLD + RUNTIME INFORMATION -=========================================================== -TP Size: 4 -PP Size: 1 -Max Runtime Batch Size: 2048 -Max Runtime Tokens: 8192 -Scheduling Policy: Guaranteed No Evict -KV Memory Percentage: 90.00% -Issue Rate (req/sec): 7.9353E+13 - -=========================================================== -= PERFORMANCE OVERVIEW -=========================================================== -Number of requests: 1000 -Average Input Length (tokens): 2048.0000 -Average Output Length (tokens): 2048.0000 -Token Throughput (tokens/sec): 1585.7480 -Request Throughput (req/sec): 0.7743 -Total Latency (ms): 1291504.1051 - -=========================================================== -``` - -**Latency** - -```bash -trtllm-bench \ ---model /path/to/hf/Llama-3.3-70B-Instruct/ \ -latency \ ---dataset /path/to/dataset/synthetic_2048_2048_1000.txt \ ---num-requests 100 \ ---warmup 10 \ ---engine_dir /path/to/engines/baseline #replace baseline with name used in llm.save() -``` -The latency benchmark enforces a batch size of 1 to accurately measure latency, which can significantly increase testing duration. In the example above the total number of requests is limited to 100 via `--num-requests` to make the test duration more manageable. This example benchmark was designed to produce very stable numbers, but in real scenarios even 100 requests is likely more than you need and can take a long time to complete (in the case-study it took about an hour and a half). Reducing the number of requests to 10 would still provide accurate data and enable faster development iterations. In general you should adjust the number of requests per your needs. Run `trtllm-bench latency -h` to see other configurable options. - -Running this command will provide a latency overview like this: - -```bash -=========================================================== -= ENGINE DETAILS -=========================================================== -Model: /scratch/Llama-3.3-70B-Instruct/ -Engine Directory: /scratch/grid_search_engines/baseline -TensorRT-LLM Version: 0.16.0 -Dtype: bfloat16 -KV Cache Dtype: None -Quantization: None -Max Input Length: 1024 -Max Sequence Length: 131072 - -=========================================================== -= WORLD + RUNTIME INFORMATION -=========================================================== -TP Size: 4 -PP Size: 1 -Max Runtime Batch Size: 1 -Max Runtime Tokens: 8192 -Scheduling Policy: Guaranteed No Evict -KV Memory Percentage: 90.00% - -=========================================================== -= GENERAL OVERVIEW -=========================================================== -Number of requests: 100 -Average Input Length (tokens): 2048.0000 -Average Output Length (tokens): 2048.0000 -Average request latency (ms): 63456.0704 - -=========================================================== -= THROUGHPUT OVERVIEW -=========================================================== -Request Throughput (req/sec): 0.0158 -Total Token Throughput (tokens/sec): 32.2742 -Generation Token Throughput (tokens/sec): 32.3338 - -=========================================================== -= LATENCY OVERVIEW -=========================================================== -Total Latency (ms): 6345624.0554 -Average time-to-first-token (ms): 147.7502 -Average inter-token latency (ms): 30.9274 -Acceptance Rate (Speculative): 1.00 - -=========================================================== -= GENERATION LATENCY BREAKDOWN -=========================================================== -MIN (ms): 63266.8804 -MAX (ms): 63374.7770 -AVG (ms): 63308.3201 -P90 (ms): 63307.1885 -P95 (ms): 63331.7136 -P99 (ms): 63374.7770 - -=========================================================== -= ACCEPTANCE BREAKDOWN -=========================================================== -MIN: 1.00 -MAX: 1.00 -AVG: 1.00 -P90: 1.00 -P95: 1.00 -P99: 1.00 - -=========================================================== -``` - -## Results - -The baseline engine achieves the following performance for token throughput, request throughput, average time to first token, and average inter-token latency. These metrics will be analyzed throughout the guide. - - -| Metric | Value | -|-------------------------------|---------------| -| Token Throughput (tokens/sec) | 1564.3040 | -| Request Throughput (req/sec) | 0.7638 | -| Average Time To First Token (ms) | 147.6976 | -| Average Inter-Token Latency (ms) | 31.3276 | - -The following sections show ways you can improve these metrics using different configuration options. diff --git a/docs/source/legacy/performance/performance-tuning-guide/deciding-model-sharding-strategy.md b/docs/source/legacy/performance/performance-tuning-guide/deciding-model-sharding-strategy.md deleted file mode 100644 index 2970d9e7e444..000000000000 --- a/docs/source/legacy/performance/performance-tuning-guide/deciding-model-sharding-strategy.md +++ /dev/null @@ -1,50 +0,0 @@ -(deciding-model-sharding-strategy)= - -# Deciding Model Sharding Strategy - -Large models often can't fit on one GPU and need to be sharded across multiple GPUs. The sharding strategies used to accomplish this can have significant impacts on performance. This guide walks through how to determine if tensor parallelism, pipeline parallelism, or a mix of both are the best strategy for you. If you are not familiar with tensor parallelism and pipeline parallelism please refer to [Mastering LLM Techniques - Inference Optimization](https://developer.nvidia.com/blog/mastering-llm-techniques-inference-optimization/) - - -## How to Think about Model Sharding: Communication is Key - -Splitting your model weights between multiple GPUs requires them to communicate activations between each other, adding additional overhead. How expensive this overhead is on your system is the key factor in determining the best strategy for you. - -In pipeline parallelism, the model is split into sets of contiguous layers and each GPU houses one of these sets. In this setup, the only required communication is for each GPU to send the outputs of its set to the GPU with the next set. - -![Pipeline Parallel Visualization](../../media/Pipeline_Parallel_Vis.svg) - - - On the other hand, tensor parallelism takes each layer of the model and splits it between the GPUs. This means that every GPU houses a portion of every layer. However since each layer needs the full outputs of the previous layer as an input, each GPU has to perform the heavier All-Reduce communication operation to share its results with all other GPUs before it can begin processing the next layer. While this seems disadvantageous, because each GPU only holds partial layers, it also performs smaller matrix multiplications, allowing it to compute its outputs quicker. - - ![Tensor Parallel Visualization](../../media/Tensor_Parallelism_Vis.svg) - - - Ultimately deciding the best strategy comes down to whether the extra overhead from the All-Reduce operation overshadows the gains from the smaller matrix multiplications. If the interconnects between the GPUs are sufficiently fast, the gains from the reduced computation burden per layer can outweigh the additional communication cost. Consequently, a general rule of thumb is that if your GPUs have fast connections between them like NVLink then tensor parallel is likely a good choice. However if the communication will go over slow connections (across nodes for example) pipeline parallel is likely better. Overall we provide the following guidelines: - -**If your model fits on one gpu:** Unless you have a very specific reason, don't shard your model. The best communication overhead is no communication overhead. - -**If your model fits in one node:** Tensor parallel is likely the best option here, especially if you have fast connections between the GPUs like NVLink. If you don't, then pipeline parallel might be needed. Start with tensor parallel and sanity check if pipeline parallel is better. - -**If your model is sharded across multiple nodes:** Inter-node connections are typically significantly slower than intra-node connections, so if you have tensor parallelism across nodes it will be bottlenecked by the slow interconects. Consequently, a good starting point is having tensor parallelism within the node and pipeline parallelism between nodes. An exception is if you are running on NVL36 or NVL72 Blackwell systems. These have multinode NVLink so as long as you stay within the 36 or 72 GPUs, tensor parallel won't be bottlenecked by inter-node connections. - -## How to set Tensor Parallelism and Pipeline Parallelism - -The `LLM` class takes `tensor_parallel_size` and `pipeline_parallel_size` as parameters. `tensor_parallel_size * pipeline_parallel_size` should be equal to the total number of GPUs you are sharding the model over, referred to as the world size. For example, if you were sharding a model over 2 nodes, each with 16 GPUs, you might set tensor parallel to 8 (for tensor parallelism within the node) and pipeline parallel to 2 (pipeline parallel between nodes) like this: - -```python - llm = LLM( - model="/scratch/Llama-3.1-405B-Instruct", - tensor_parallel_size=8, - pipeline_parallel_size=2 - ) -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) you can specify tensor parallelism and pipeline parallelism by providing the `--tp_size` and `--tp_size` arguments to `convert_checkpoint.py` - -``` -python examples/models/core/llama/convert_checkpoint.py --model_dir ./tmp/llama/405B/ \ - --output_dir ./tllm_checkpoint_16gpu_tp8_pp2 \ - --dtype float16 \ - --tp_size 8 - --pp_size 2 -``` diff --git a/docs/source/legacy/performance/performance-tuning-guide/fp8-quantization.md b/docs/source/legacy/performance/performance-tuning-guide/fp8-quantization.md deleted file mode 100644 index 3a5121b7c173..000000000000 --- a/docs/source/legacy/performance/performance-tuning-guide/fp8-quantization.md +++ /dev/null @@ -1,217 +0,0 @@ -(fp8-quantization)= - -# FP8 Quantization - -Quantization is a technique that allows models to run in lower precisions like int8 and fp8 while maintaining acceptable output quality. Running in lower precisions can greatly boost performance, significantly increasing throughput and decreasing latency. The tradeoff is a drop in output quality, but in many cases the output quality is still acceptable and many real world deployments utilize quantization. If you want to learn more about quantization refer to [Mastering LLM Techniques - Inference Optimization](https://developer.nvidia.com/blog/mastering-llm-techniques-inference-optimization/) - -This section walks through enabling fp8 quantization and highlight some fp8 quantization specific configuration options for boosting performance. It also continues the case study of Llama-3.3-70B split across 4 H100-sxm-80GB GPUs via tensor parallelism and showcase the effects of enabling these configuration options on performance. - -> Disclaimer: While performance numbers shown here are real, they are only for demonstration purposes. Differences in environment, SKU, interconnect, and workload can all significantly affect performance and lead to your results differing from what is shown here. - -## Enabling Quantization - -To enable quantization you need to configure the `QuantConfig` class and pass it to the `quant_config` parameter of the LLM class. At a minimum the `quant_algo` parameter, which sets the quantization algorithm (fp8, fp8 per token, int8awq, etc.) must be specified. You can find all supported quantization algorithms and other configurable options for `QuantConfig` in the LLM-API->Reference section of the docs. While it is not required if you are using weights/checkpoints from that are already quantized, if you are using an fp16 checkpoint then you also need to specify the calibration dataset that will be used to determine the quantization scales via `CalibConfig`. `CalibConfig` provides several options for setting the calibration dataset that can also be referenced in the LLM-API->Reference section of the docs. Although TensorRT-LLM supports several other types of quantization, this guide focuses on fp8. - - -Here is an example of building and saving an fp8 engine from a bf16 checkpoint (Note that fp8 is supported only on devices with compute capability > 8.9 - Ada, Hopper, Blackwell, and beyond): -```python -from tensorrt_llm import LLM, BuildConfig -from tensorrt_llm.llmapi import QuantConfig, QuantAlgo, CalibConfig - -def main(): - - quant_config = QuantConfig(quant_algo=QuantAlgo.FP8) - - calib_config = CalibConfig( - calib_batches=512, - calib_batch_size=1, - calib_max_seq_length=2048, - tokenizer_max_seq_length=4096 - ) - - build_config = BuildConfig( - max_num_tokens=2048, - max_batch_size=512, - ) - - build_config.plugin_config.use_paged_context_fmha = True - build_config.plugin_config.multiple_profiles = True - - llm = LLM( - model="/path/to/Llama-3.3-70B", - tensor_parallel_size=4, - pipeline_parallel_size=1, - build_config=build_config, - quant_config=quant_config, - calib_config=calib_config - ) - - llm.save("baseline_fp8_engine") - -if __name__ == '__main__': - main() -``` - -For an example of how to build an fp8 engine using the [TensorRT-LLM CLI workflow](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) flow see [TensorRT-LLM LLaMA examples](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/llama). In short you first run [`examples/quantization/quantize.py`](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/quantization) to quantize and convert the model checkpoint to TensorRT-LLM format and then use `trtllm-build`. - -> ***Note: While quantization aims to preserve model accuracy this is not guaranteed and it is extremely important you check that the quality of outputs remains sufficient after quantization.*** - -## FP8 "Baseline" Performance - -Benchmarking the engine produced by the example above yielded the following performance results. Note that we enabled some of the build flags we mentioned [earlier](./useful-build-time-flags.md) (multiple profiles, paged_context_fmha) and also tuned max batch size and max num tokens. This is done to give a sense of what performance is achievable if you tune an fp8 engine but exclude options that have been tailored for quantization. We recommend disabling the gemm plugin for quantized engines which is why it is not included here (it is off by default). Reduce fusion has a quantization specific optimization that will be covered later. For the remainder of this page we will refer to this setup as the "baseline" numbers for fp8. - - -| Metric | Value | -| -------------------------------- | --------- | -| Token Throughput (tokens/sec) | 3389.5305 | -| Request Throughput (req/sec) | 1.6550 | -| Average Time To First Token (ms) | 96.1597 | -| Average Inter-Token Latency (ms) | 12.4248 | - - -## Quantized KV-Cache - -By default the KV-Cache is not quantized but TensorRT-LLM supports quantizing the KV-Cache to further improve performance. However, quantizing the model more aggressively also increases the risk of model output quality degrading so it is important to check that when using this feature. - -### Enabling Quantized KV Cache - -The LLM-API exposes the quantization algorithm to be used for kv cache via the `kv_cache_quant_algo` field in `QuantConfig`. To enable fp8 kv cache, you would modify `QuantConfig` as such: - -```python -quant_config = QuantConfig(quant_algo=QuantAlgo.FP8, - kv_cache_quant_algo=QuantAlgo.FP8) -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--kv_cache_dtype fp8` to [`examples/quantization/quantize.py`](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/quantization). - -### Performance with Quantized KV Cache - -| Metric | Baseline | FP8 KV-Cache ON | -| -------------------------------- | --------- | --------------- | -| Token Throughput (tokens/sec) | 3389.5305 | 5299.6372 | -| Request Throughput (req/sec) | 1.6550 | 2.5877 | -| Average Time To First Token (ms) | 96.1597 | 97.1287 | -| Average Inter-Token Latency (ms) | 12.4248 | 12.5496 | - -## Reduce Norm Fusion with User Buffers for Llama Models - -The [Reduce Norm Fusion](./useful-build-time-flags.md#reduce-norm-fusion-plugin-for-llama-models) feature is supported for fp8. An additional optimization called "User Buffers" is also supported for fp8 models. The user buffer feature aims to eliminate extra copies from the local buffer to the shared buffer in the communication kernel, leading to improved end-to-end performance. - - -### Enabling Reduce Norm Fusion with User Buffers - - -To enable reduce norm fusion with user buffers, add the following lines below `BuildConfig`'s initialization - -```python -build_config.plugin_config.reduce_fusion = True -build_config.plugin_config.user_buffer = True -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--reduce_fusion enable` and `--user_buffer enable` to `trtllm-build` to enable the feature. - -> Note: You must have enabled `reduce_fusion` in order to enable `user_buffer` - -### Performance with Reduce Norm Fusion + User Buffers: - -Reduce Norm Fusion + User Buffer ON: Same engine previously referred to as FP8 KV-Cache ON. - -Reduce Norm Fusion + User Buffer ON: Previous example with reduce fusion and user buffers enabled. Max-num tokens set to 16384 and max-batch size set to 512 after tuning. - - -| Metric | Reduce Norm Fusion + User Buffer OFF | Reduce Norm Fusion + User Buffer ON | -| -------------------------------- | ------------------------------------ | ----------------------------------- | -| Token Throughput (tokens/sec) | 5299.6372 | 5980.7842 | -| Request Throughput (req/sec) | 2.5877 | 2.9203 | -| Average Time To First Token (ms) | 97.1287 | 82.2679 | -| Average Inter-Token Latency (ms) | 12.5496 | 12.6975 | - -## GEMM + SwiGLU Fusion in Gated-MLP - -The GEMM + SwiGLU fusion in Gated-MLP combines two Matmul operations and one SwiGLU operation into a single kernel. Currently this is only supported for FP8 precision on Hopper. While this fusion improves performance, it can slightly reduce accuracy in FP8 PTQ because one quantization scaling factor is discarded. - -We recommend enabling this feature for large models running on Hopper with FP8 precision. We do not recommend enabling this feature for very small workloads or if the -accuracy loss is unacceptable. - -### Enabling GEMM + SwiGLU Fusion - -To enable the GEMM + SwiGLU fusion, add the following lines below `BuildConfig`'s initialization - -```python -build_config.plugin_config.gemm_swiglu_plugin = 'fp8' -``` -For small batch size cases where latency is important, you can replace the above line with - -```python -build_config.plugin_config.low_latency_gemm_swiglu_plugin = 'fp8' -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--gemm_swiglu_plugin=fp8` or `--low_latency_gemm_swiglu_plugin=fp8` for the low latency case (only include one or the other) to `trtllm-build`. - -### Performance with GEMM + SwiGLU Fusion - - -| Metric | GEMM + SwiGLU fusion OFF | GEMM + SwiGLU fusion ON | -| -------------------------------- | ------------------------ | ----------------------- | -| Token Throughput (tokens/sec) | 5980.7842 | 5976.7977 | -| Request Throughput (req/sec) | 2.9203 | 2.9184 | -| Average Time To First Token (ms) | 82.2679 | 81.8841 | -| Average Inter-Token Latency (ms) | 12.6975 | 11.7031 | - -In this case, the GEMM + SwiGLU plugin performs almost equivalently to when it was disabled. The throughput drop is within run to run variance and the TTFT and ITL improvements are slight. However, we found that when paired with the low latency gemm plugin discussed next, enabling this feature was necessary for getting the maximum throughput. - -## Low Latency GEMM Plugin - -Previously we mentioned the [GEMM Plugin](./useful-build-time-flags.md#gemm-plugin) feature. Although it has fp8 support we recommend disabling it (by default it is disabled). However for low-latency scenarios in fp8 we recommend trying the low latency GEMM plugin to see if it is effective for your workload. - -### Enabling Low Latency GEMM plugin - -To enable the low latency GEMM plugin, add the following lines below `BuildConfig`'s initialization - -```python -build_config.plugin_config.low_latency_gemm_plugin = 'fp8' -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--low_latency_gemm_plugin=fp8` to `trtllm-build` to enable the feature. Again, **we recommend disabling the gemm plugin for fp8** so if you are passing `--gemm_plugin=fp8` to `trtllm-build` we recommend removing that. - -### Performance with Low Latency GEMM plugin - -Low Latency GEMM ON: Same configuration as previous example but with low latency GEMM plugin enabled. Max num tokens was set to 16384 and max-batch size was set to 512 after tuning. - -| Metric | Low Latency GEMM OFF | Low Latency GEMM ON | -| -------------------------------- | -------------------- | ------------------- | -| Token Throughput (tokens/sec) | 5976.7977 | 6049.1625 | -| Request Throughput (req/sec) | 2.9184 | 2.9537 | -| Average Time To First Token (ms) | 81.8841 | 88.0162 | -| Average Inter-Token Latency (ms) | 11.7031 | 10.8225 | - -In this case, enabling the low-latency gemm plugin actually provided a meaningful boost to throughput. Additionally it also improved ITL but at the expense of TTFT. Furthermore, when used without the gemm+swiglu fusion, performance was actually worse than with out the plugin turned on. This suggests that for this workload the low-latency gemm plugin was choosing a worse kernel for the gemm right before the swiglu, but once that was handled by the gemm+swiglu fusion custom kernel, the rest of the kernels the low-latency gemm plugin was choosing was better than the baseline, resulting in improved performance. This underscores the importance of benchmarking different settings as the impact of this plugin is highly workload dependent. If possible some grid searching can be useful for extremely performance sensitive workloads - -## Conclusion - -Overall leveraging quantization can provide significant uplifts in performance. Here are the performance uplifts from our tuned fp8 model as compared to the tuned fp16 numbers we reached in the [previous page of guide](./tuning-max-batch-size-and-max-num-tokens.md) - -| Metric | Tuned FP16 Model | Tuned FP8 Model | % Improvement | -| -------------------------------- | ---------------- | --------------- | ------------- | -| Token Throughput (tokens/sec) | 2474.2581 | 6049.1625 | 144.48 | -| Request Throughput (req/sec) | 1.2081 | 2.9537 | 144.49 | -| Average Time To First Token (ms) | 147.5742 | 88.0162 | 40.36 | -| Average Inter-Token Latency (ms) | 14.6852 | 10.8225 | 26.30 | - -Additionally, compared to the fp8 baseline numbers (the baseline numbers had some degree of tuning, see [Baseline Performance](./fp8-quantization.md#fp8-baseline-performance) for details), we received the following performance uplifts from enabling the flags discussed above: - -| Metric | Baseline FP8 Model | Tuned FP8 Model | % Improvement | -| -------------------------------- | ------------------ | --------------- | ------------- | -| Token Throughput (tokens/sec) | 3389.5305 | 6049.1625 | 78.47 | -| Request Throughput (req/sec) | 1.6550 | 2.9537 | 78.47 | -| Average Time To First Token (ms) | 96.1597 | 88.0162 | 8.47 | -| Average Inter-Token Latency (ms) | 12.4248 | 10.8225 | 12.90 | - -As mentioned previously, the caveat with leveraging quantization are potential drops in accuracy, and we strongly recommend having a way to test whether model output quality is acceptable before attempting to use quantization. That said, many real world cases successfully use quantization and the significant performance boosts it enables are often worth the effort to see if it is a fit. - -### Summary of Configuration Option Recommendations: - -1. Quantized KV-cache: Typically provides significant throughput boost. We recommend turning it on as long as output quality is still acceptable with the feature enabled. -2. Reduce fusion + user buffers: This feature is only supported on fp8 Llama and Mistral/Mixtral models. Effectiveness is workload dependent so we recommend turning it on and benchmarking to check. -3. Gemm + Swiglu Plugin: This feature is only supported on fp8 models with Swiglu operators like Llama, Mixtral etc. Like reduce fusion effectiveness is workload dependent and we recommend sanity checking effectiveness. Has increased risk of affecting accuracy since it drops a quantization scale. -4. Low-Latency GEMM plugin: Effectiveness is workload dependent so we recommend turning it on and benchmarking. Effectiveness can be affected by other flags as we saw in our case study, so if possible benchmarking various combinations of configuration options is ideal. diff --git a/docs/source/legacy/performance/performance-tuning-guide/index.rst b/docs/source/legacy/performance/performance-tuning-guide/index.rst deleted file mode 100644 index f642a4ba2a88..000000000000 --- a/docs/source/legacy/performance/performance-tuning-guide/index.rst +++ /dev/null @@ -1,15 +0,0 @@ -Performance Tuning Guide -======================= - -.. include:: introduction.md - :parser: myst_parser.sphinx_ - -.. toctree:: - :maxdepth: 1 - - benchmarking-default-performance - useful-build-time-flags - tuning-max-batch-size-and-max-num-tokens - deciding-model-sharding-strategy - fp8-quantization - useful-runtime-flags diff --git a/docs/source/legacy/performance/performance-tuning-guide/introduction.md b/docs/source/legacy/performance/performance-tuning-guide/introduction.md deleted file mode 100644 index 800ae65f37bf..000000000000 --- a/docs/source/legacy/performance/performance-tuning-guide/introduction.md +++ /dev/null @@ -1,17 +0,0 @@ -While defaults are expected to provide solid performance, TensorRT-LLM has several configurable options that can improve performance for your particular workload. This guide is meant to help you tune TensorRT-LLM to extract the best performance for your use case. It covers several of the most helpful tunable parameters and provides intuition for thinking about them. This guide also doubles as an example of how to work with TensorRT-LLM's LLM-API and its TRTLLM-Bench benchmarking workflow. - -This guide uses Llama-3.3-70b on 4 H100-sxm-80GB connected via NVLink as a case study and focuses on optimizing performance on input sequence length/output sequence length of 2048/2048. Case study sections throughout this guide reference internal performance testing and results to help reinforce the conclusions and recommendations given. - -## Prerequisite Knowledge - -This guide expects you have some familiarity with the following concepts - -- Phases of Inference: Context (Prefill) Phase and Generation Phase -- Inflight Batching -- Tensor Parallelism and Pipeline Parallelism -- Quantization - - Please refer to [Mastering LLM Techniques - Inference Optimization](https://developer.nvidia.com/blog/mastering-llm-techniques-inference-optimization/) for an introduction to these concepts. - -## Table of Contents - diff --git a/docs/source/legacy/performance/performance-tuning-guide/tuning-max-batch-size-and-max-num-tokens.md b/docs/source/legacy/performance/performance-tuning-guide/tuning-max-batch-size-and-max-num-tokens.md deleted file mode 100644 index 6ed3a5c5c205..000000000000 --- a/docs/source/legacy/performance/performance-tuning-guide/tuning-max-batch-size-and-max-num-tokens.md +++ /dev/null @@ -1,126 +0,0 @@ -(tuning-max-batch-size-and-max-num-tokens)= - -# Tuning Max Batch Size and Max Num Tokens - -One of TensorRT-LLM's key features is its inflight batching scheduler and runtime, which simultaneously schedules and executes requests in both context and generation phases. Max num tokens, together with max batch size, dictate how and when the scheduler admits new requests and continues current ones, and understanding how to tune them can provide significant performance benefits. - -> Disclaimer: While performance numbers shown here are real, they are only for demonstration purposes. Differences in environment, SKU, interconnect, and workload can all significantly affect performance and lead to your results differing from what is shown here. - -## Understanding the TensorRT-LLM scheduler - -This section visualizes how TensorRT-LLM schedules requests based on max-batch size and max-num tokens. The example starts out with a newly initialized engine as well as a few unscheduled requests that have come in. For the sake of this example, toy values are set to `max batch size = 4` and `max num tokens = 12`. Each square block represents a token, and its color represents which request it belongs to. - -![TRT-LLM Scheduler Visualization 1](../../media/TRTLLM_Scheduler_Vis_1.svg) - - -Now the scheduler takes the first two requests, Request 1 and Request 2, and schedules them to execute the context phase. However, it cannot schedule any more requests because the prompts of the first two requests had 5 tokens each, leaving a budget of 2 tokens due to the max num tokens limit. Since all remaining requests have more than 2 prompt tokens none of them can be scheduled (context chunking can help in this situation, see the paged context attention section below). The tokens are marked with a "C" on them to represent that they are prompt tokens that were processed in the context phase. - -> Note: The tokens for different requests are shown on different rows simply for visualization purposes and are not representative of actual memory layouts - -![TRT-LLM Scheduler Visualization 2](../../media/TRTLLM_Scheduler_Vis_2.svg) - -Now the engine runs an iteration of execution, completing the context phases for both of the scheduled requests. After it is done, the kv-cache of the prompts for both requests have been created and the first token has been generated. Tokens that were generated are marked with "G(n)" - for example a token marked "G1" represents that it is the first token generated for its request. - -TRT-LLM prioritizes scheduling requests in generation phase first so the two generated tokens are queued to be processed in the next iteration. Now, since the two previously scheduled requests have entered generation phase and only take up two tokens out of the max num token budget of 12, the scheduler is able to schedule two additional requests, Request 3 and Request 4. It cannot schedule the last request, Request 5, even though there is space for it in the max num tokens budget because of the max batch size limit of 4. - -![TRT-LLM Scheduler Visualization 3](../../media/TRTLLM_Scheduler_Vis_3.svg) - -After the next iteration of execution, the second tokens for Requests 1 and 2 have been generated, and the first tokens for Request 3 and 4 have been generated. Lets say that G2 that was generated for Request 1 is the stop token, signifying that Request 1 is completed. In this case the scheduler would evict Request 1 before performing another execution iteration and prepare to return it to the user. This eviction puts the state of the engine below the max batch size limit and allows Request 5 to be scheduled. - -Another thing to note is that G1 that was generated for Request 2 has been added to the kv-cache for request 2, representing how kv-cache for a request grows as more and more tokens are generated. - -![TRT-LLM Scheduler Visualization 4](../../media/TRTLLM_Scheduler_Vis_4.svg) - -Overall, the max batch size and max num tokens limits play a key part in deciding when requests are actually executed, and tuning them can have significant impacts on throughput numbers as well as how the engine balances previously scheduled requests in generation phase with context phase on new requests - -> Note: This presents a simplified visualization of the scheduler to highlight how max batch size and max num tokens affect it. The scheduler also considers things like amount of free memory available to be used for kv-cache and has other configurable options that can affect its behavior. See the Runtime Flags of the Additional Options page for more on this. - - - ## Tuning Max Batch Size - - It is important to set the max batch size large enough that it doesn't bottleneck the scheduling of new requests. Due to this it's recommended to test additional values of max batch size for your workload. The default value is 2048. Powers of 2 are good initial values to sweep through. - -### How to change Max Batch Size - -You can specify the max batch size in the build config. - -```python -build_config = BuildConfig( - max_batch_size=512 -) -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--max_batch_size ` to `trtllm-build` to tune max batch size. - -### Tuning Case Study - -Continuing with our case study of Llama-3.3-70B on 4 H100s, we pick up where left off in the previous section (building an engine with multiple profiles, gemm plugin, paged context attention, and reduce fusion enabled). Sweeping across max batch sizes of 64, 512, and the default 2048 produced the following results - -| Metric | Max Batch Size 64 | Max Batch Size 512 | Max Batch Size 2048 | -| -------------------------------- | ----------------- | ------------------ | ------------------- | -| Token Throughput (tokens/sec) | 1944.3031 | 2466.7933 | 2044.2628 | -| Request Throughput (req/sec) | 0.9494 | 1.2045 | 0.9982 | -| Average Time To First Token (ms) | 145.7607 | 147.7876 | 146.6628 | -| Average Inter-Token Latency (ms) | 14.6475 | 14.6554 | 14.4493 | - -From this its clear a max batch size of 64 results in bottlenecking whereas a max batch size of 512 is the sweet spot, boosting throughput by almost 20% with negligible effect on latency. - -## Tuning Max Num Tokens - -If max num tokens is too small it can bottleneck request scheduling. However, if it is too large it can result in prompt tokens taking up too much memory (especially in long-context workloads), not leaving enough for kv-cache and resulting in performance hits or even out-of-memory errors. The default value of max num tokens is 8192. It's recommended that you sweep across several values of max num tokens to identify the best number for your workload. Good values to try are powers of 2 >= 1024. Since max num tokens and max batch size both affect scheduling, grid searching across combinations of them if possible is ideal. - -### How to change Max Num Tokens - -Like Max Batch size, max num tokens is also specified in the build config. - -```python -build_config = BuildConfig( - max_batch_size=512 - max_num_tokens=2048 -) -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--max_num_tokens ` to `trtllm-build` to tune max_num_tokens. - -### Tuning Case Study - -Sweeping across max num tokens values of 2048, 8192, and 16384 yielded the following performance numbers with max batch size set to 512. - -| Metric | Max Num Tokens 2048 | Max Num Tokens 8192 | Max Num Tokens 16384 | -| -------------------------------- | ------------------- | ------------------- | -------------------- | -| Token Throughput (tokens/sec) | 2474.2581 | 2466.7933 | 2461.0165 | -| Request Throughput (req/sec) | 1.2081 | 1.2045 | 1.2017 | -| Average Time To First Token (ms) | 147.5742 | 147.7876 | 147.9623 | -| Average Inter-Token Latency (ms) | 14.6852 | 14.6554 | 14.6769 | - -For this particular workload max num tokens of 2048 provides the best performance, but not by an extremely large margin. This reflects the reality that for any given workload tuning various flags can have varied impact. However it is important to check different values to ensure that you are not giving up large gains due to scheduling imbalances. - -## Revisiting Paged Context Attention and Context Chunking - -[Previously](./useful-build-time-flags.md#paged-context-attention) we recommended enabling paged context attention even though in our case study it didn't affect performance significantly. Having now understood the TensorRT-LLM scheduler we can now explain why this is beneficial. In short, we recommend enabling it because it enables context chunking, which allows the context phase of a request to be broken up into pieces and processed over several execution iterations, allowing the engine to provide a more stable balance of context and generation phase execution. - -The [visualization](#understanding-the-trt-llm-scheduler) of the TensorRT-LLM scheduler showed that initially Request 3 couldn't be scheduled because it would put the scheduler over the max-num tokens limit. However with context chunking, this is no longer the case, and the first chunk of Request 3 would be able to be scheduled. - -![TRT-LLM Scheduler Visualization Chunked Context 1](../../media/TRTLLM_Scheduler_Vis_Chunked_Context_1.svg) - -This is extremely beneficial for several reasons. Firstly it eliminates the possibility of requests with large prompts relative to max num tokens being unable to be scheduled due to other requests that are already in-flight. In production workloads, this can help improve worst case TTFT numbers. Secondly it allows for setting smaller values of max num tokens since you no longer need max num tokens to be at least as large as the longest prompt you want to support. For long-context cases this is extremely important, because setting extremely large values of max-num tokens takes away from memory available to be used as kv-cache. Given that in the worst case scenario chunked context has minimal impact on performance but can significantly benefit it in many scenarios, it's recommended that you always enable it. - -## Conclusion - -The TensorRT-LLM Scheduler plays a large role in performance, and properly tuning it can provide significant performance uplifts. In the case-study example, tuning max batch size and max num tokens provided the following boosts to performance when compared to the [results from the previous page](./useful-build-time-flags.md#conclusion): - -| Metric | Build-Time Flags ON | Tuned Max Batch Size and Max Num Tokens | % Improvement | -| -------------------------------- | ------------------- | --------------------------------------- | ------------- | -| Token Throughput (tokens/sec) | 2044.2628 | 2474.2581 | 21.03 | -| Request Throughput (req/sec) | 0.9982 | 1.2081 | 21.03 | -| Average Time To First Token (ms) | 146.6628 | 147.5742 | -0.62 | -| Average Inter-Token Latency (ms) | 14.4493 | 14.6852 | -1.63 | - -Interpreting these results, tuning max batch size and max num tokens significantly improved throughput and remained at par on latency (the slight drops are within run to run variance). Compared to the [initial baseline](./benchmarking-default-performance.md#results) the tuned engine achieves the following uplifts: - -| Metric | Baseline | Build-Time Flags ON and Tuned Max Batch Size and Max Num Tokens | % Improvement | -| -------------------------------- | --------- | --------------------------------------- | ------------- | -| Token Throughput (tokens/sec) | 1564.3040 | 2474.2581 | 58.17 | -| Request Throughput (req/sec) | 0.7638 | 1.2081 | 58.17 | -| Average Time To First Token (ms) | 147.6976 | 147.5742 | 0.08 | -| Average Inter-Token Latency (ms) | 31.3276 | 14.6852 | 53.12 | diff --git a/docs/source/legacy/performance/performance-tuning-guide/useful-build-time-flags.md b/docs/source/legacy/performance/performance-tuning-guide/useful-build-time-flags.md deleted file mode 100644 index 54ecd97fa6bd..000000000000 --- a/docs/source/legacy/performance/performance-tuning-guide/useful-build-time-flags.md +++ /dev/null @@ -1,171 +0,0 @@ -(useful-build-time-flags)= - -# Useful Build-Time Flags - -This page presents several build-time flags, set via the LLM-API's `BuildConfig` class that you can enable to improve upon the baseline performance. Build-time refers to the fact that these flags affect how the TensorRT-LLM engine is built and cannot be changed without rebuilding the engine. For each flag there is an explanation of what it does, a description of how to enable it, and then an example of running it through the benchmarking flow described in [Benchmarking Default Performance](./benchmarking-default-performance.md) to showcase its impact on performance. All options compatible with `trtllm-build` can be found in the Command Line Reference section of the docs. - -> Disclaimer: While performance numbers shown here are real, they are only for demonstration purposes. Differences in environment, SKU, interconnect, and workload can all significantly affect performance and lead to your results differing from what is shown here. - -## Multiple Profiles - -TensorRT-LLM is built on TensorRT, which handles engine building through "optimization profiles" defining min, optimal, and max input tensor shapes. TensorRT optimizes for the optimal shape while supporting the range between min and max. - -TensorRT-LLM abstracts away the need to create optimization profiles although flags like max_batch_size and max_num_tokens (covered later) influence how they are created. By default, only one profile is created. - -During inference serving, varying request loads can pose different tensor shapes to the engine. TensorRT addresses this by allowing multiple profiles, which TensorRT-LLM supports via the BuildConfig option in the LLM-API. Enabling multiple profiles increases build times but has no performance downsides, so it is recommended for production builds. - -The only thing to watch out for is that enabling this can lead to slightly different outputs when the same prompt is run multiple times as different profiles and consequently kernels might be used depending on the request load. However this variance should not affect output quality so it is safe to enable this flag as long as you don't need completely deterministic outputs. - -### Enabling building with multiple profiles - -Below is an example of how you can modify the baseline example to enable multiple profiles. - -```python -from tensorrt_llm import LLM, BuildConfig - -def main(): - build_config = BuildConfig() - build_config.plugin_config.multiple_profiles = True - - llm = LLM( - model="/scratch/Llama-3.3-70B-Instruct", - tensor_parallel_size=4, - build_config=build_config - ) - - llm.save("build_flags_multiple_profiles") - -if __name__ == '__main__': - main() -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--multiple_profiles` to `trtllm-build` to enable the feature. - - -### Performance with multiple profiles - -Baseline refers to the engine that was benchmarked in the previous Benchmarking Default Performance page. - -| Metric | Baseline | Multiple Profiles ON | -| -------------------------------- | --------- | -------------------- | -| Token Throughput (tokens/sec) | 1564.3040 | 1861.0881 | -| Request Throughput (req/sec) | 0.7638 | 0.9087 | -| Average Time To First Token (ms) | 147.6976 | 145.8958 | -| Average Inter-Token Latency (ms) | 31.3276 | 19.6452 | - -As you can see, enabling multiple profiles significantly improves the metrics across the board. - -## Paged Context Attention - -By default all the tokens of the prompt of a new request are processed in one iteration as the context phase. Enabling paged context attention allows TensorRT-LLM to break the context phase into chunks and handle the prompt over several iterations. This is particularly useful for workloads with large input length. In the worst case, this feature can provide a small performance hit in benchmarking runs (<2%) so it can be safely enabled. This feature is discussed further in the [next page](./tuning-max-batch-size-and-max-num-tokens.md#revisiting-paged-context-attention-and-context-chunking) of the guide. - - -### Enabling Paged Context Attention - -Add the following line to our multiple profiles example from above to enable paged context attention - -```python -build_config.plugin_config.use_paged_context_fmha=True -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--use_paged_context_fmha` to `trtllm-build` to enable the feature. - -### Performance - -Paged Context OFF refers to the same engine shown as Multiple Profiles ON in the previous example. - -| Metric | Paged Context OFF | Paged Context ON | -| -------------------------------- | ----------------- | ---------------- | -| Token Throughput (tokens/sec) | 1861.0881 | 1866.6684 | -| Request Throughput (req/sec) | 0.9087 | 0.9115 | -| Average Time To First Token (ms) | 145.8958 | 145.4089 | -| Average Inter-Token Latency (ms) | 19.6452 | 19.6523 | - -In this case enabling paged context attention provides a small boost to performance, but a rerun of our tests found this to be within run to run variance of around 10 tok/s for token throughput and 2ms for average time to first token (ITL was stable with <1ms and request throughput corresponded directly to token throughput). In other cases naively enabling it might actually provide a small hit to performance. However, further guidance on how to reason about this flag and why we recommend enabling it is discussed in the [next page](./tuning-max-batch-size-and-max-num-tokens.md#revisiting-paged-context-attention-and-context-chunking) as it is closely intertwined with how TensorRT-LLM schedules requests as well as the max-num tokens flag. - -## GEMM Plugin - -TensorRT allows you to add "plugins" or custom kernels that can be used instead of the kernels that TensorRT selects for particular operations. TensorRT-LLM has a host of custom plugins that are specifically tailored to speed up supported modules. The GEMM plugin utilizes NVIDIA cuBLASLt and some custom kernels to perform GEMM operations. On FP16 and BF16, it’s recommended to be enabled for better performance and smaller GPU memory usage. On FP8, it’s recommended to be disabled. - -### Enabling GEMM Plugin - -Add the following line to the multiple profiles example from above to enable the GEMM plugin. - -```python -build_config.plugin_config.gemm_plugin = 'auto' -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--gemm_plugin auto` to `trtllm-build` to enable the feature. `'auto'` tells the GEMM plugin to have the same type as the model (fp16, bf16, etc). It is fine to leave it on auto unless you are trying to do mixed precision. - -### Performance with GEMM Plugin - -GEMM Plugin OFF refers to the same engine shown as Paged Context ON in the previous example. - -| Metric | GEMM Plugin OFF | GEMM Plugin ON | -| -------------------------------- | --------------- | -------------- | -| Token Throughput (tokens/sec) | 1866.6684 | 2033.2640 | -| Request Throughput (req/sec) | 0.9115 | 0.9928 | -| Average Time To First Token (ms) | 145.4089 | 147.8307 | -| Average Inter-Token Latency (ms) | 19.6523 | 15.4133 | - -In this case the GEMM plugin greatly improves throughput as well as ITL, with a slight hit to TTFT. - -## Reduce Norm Fusion Plugin for Llama models: - -TensorRT-LLM has custom kernels for AllReduce operations that are enabled by default. This feature extends this functionality by fusing the ResidualAdd and LayerNorm kernels that run after AllReduce into the AllReduce kernel, resulting in a single kernel that handles those operations and improves end-to-end performance. This feature is currently only available for Llama models. It is most beneficial in workloads that are generation-phase heavy. For extremely context-phase heavy workloads its worth checking performance with and without this. Additionally, since this is an optimization for AllReduce, it is only beneficial for cases with tensor-parallelism. For scenarios only using pipeline parallelism this should stay disabled since pipeline parallelism doesn't require any AllReduce operations. - -### Enabling Reduce Norm Fusion Plugin - -Add the following line to the multiple profiles example from above to enable reduce norm fusion. - -```python -build_config.plugin_config.reduce_fusion = True -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--reduce_fusion enable` to `trtllm-build` to enable the feature. - -### Performance with Reduce Norm Fusion - -Reduce Fusion OFF refers to the same engine shown as GEMM Plugin ON in the previous example. - -| Metric | REDUCE FUSION OFF | REDUCE FUSION ON | -| -------------------------------- | ----------------- | ---------------- | -| Token Throughput (tokens/sec) | 2033.2640 | 2044.2628 | -| Request Throughput (req/sec) | 0.9928 | 0.9982 | -| Average Time To First Token (ms) | 147.8307 | 146.6628 | -| Average Inter-Token Latency (ms) | 15.4133 | 14.4493 | - -For the ISL/OSL pair of 2048/2048 enabling the reduce norm fusion plugin slightly improves performance all around. However, test reruns found that with run to run variance, in the worst case, they performed at par. Again this flag's effectiveness is dependent on the workload so users should check whether it provides meaningful performance boosts in their case. - - -## Pipeline Parallel Reduce Scatter Optimization - -This feature adds a pipeline parallelism optimization with ReduceScatter + AllGather targeting large mixture of experts models. -This can be enabled via the LLM-API as such -```python - build_config.plugin_config.pp_reduce_scatter = True -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) flow you can enable this feature by adding `--pp_reduce_scatter` to `trtllm-build`. - -As the Llama model is not a MoE model this flag was not included as part of the case study. - -## Conclusion - -Overall, enabling these flags can greatly boost performance. However, the degree to which they are effective can vary from workload to workload, and it's recommended that you run sanity checks on your workloads to verify performance. - -The case-study example showed that enabling these flags provided the following performance uplifts from the baseline numbers. This included significant boosts in Token Throughput, Request Throughput, and Average Inter-Token Latency. TTFT remained largely unchanged. - -| Metric | Baseline | Build-Time Flags ON | % Improvement | -| -------------------------------- | --------- | ------------------- | ------------- | -| Token Throughput (tokens/sec) | 1564.3040 | 2044.2628 | 30.68 | -| Request Throughput (req/sec) | 0.7638 | 0.9982 | 30.69 | -| Average Time To First Token (ms) | 147.6976 | 146.6628 | 0.70 | -| Average Inter-Token Latency (ms) | 31.3276 | 14.4493 | 53.88 | - -### Summary of Configuration Option Recommendations: - -1. Multiple profiles: Always enable. It may increase build times a little but will only ever help performance. Enabling might cause engine to produce slightly different outputs when the same prompt is run multiple times depending on request load but it should not affect output quality, see [Multiple Profiles section](./useful-build-time-flags.md#multiple-profiles) for explanation. -2. Paged Context Attention: In the worst case it may hurt performance a little initially but typically helps with request scheduling and boosts performance after further tuning of max batch size and max num tokens. More on this topic is discussed in the next page. -3. GEMM Plugin: It's recommended to enable it for FP16 and BF16 models as it usually helps. However, it is a good idea to benchmark your workload and double check that it is helping. -4. Reduce Fusion: This feature is only supported on Llama and Mistral/Mixtral models. Effectiveness is workload dependent and it's recommend that you benchmark your workload with and without it and compare the results. diff --git a/docs/source/legacy/performance/performance-tuning-guide/useful-runtime-flags.md b/docs/source/legacy/performance/performance-tuning-guide/useful-runtime-flags.md deleted file mode 100644 index 8c5579431b96..000000000000 --- a/docs/source/legacy/performance/performance-tuning-guide/useful-runtime-flags.md +++ /dev/null @@ -1,224 +0,0 @@ -(useful-runtime-flags)= - -# Useful Runtime Options - -This part summarizes the runtime configuration knobs that can be tweaked to -enhance the performance of already built engines. As compared to previous examples where - the LLM-API was used to build and save an engine but not to process any requests, -runtime knobs would be specified when you are using the LLM-API to actually run inference -like in the [LLM-API end-to-end example](./benchmarking-default-performance.md#before-you-begin-tensorrt-llm-llm-api) - - -## Capacity Scheduler Policy - -TensorRT-LLM currently supports three batch scheduler policies: `GUARANTEED_NO_EVICT` (default), -`MAX_UTILIZATION` and `STATIC_BATCH`. - -The scheduling policy can be set to `MAX_UTILIZATION` to pack as many -requests as possible at each iteration of the forward loop, when in-flight -sequence batching is enabled. It maximizes the utilization of the GPUs by -aggressively scheduling requests at the risk of having to pause requests if the -KV cache size limit is reached. - -For a more conservative approach with respect to the KV cache limitations in -terms of memory allocation, `CapacitySchedulerPolicy` should be set to -`GUARANTEED_NO_EVICT` to guarantee that a started request is never paused. - -If the goal is to maximize the throughput, users should try `MAX_UTILIZATION`. -However, they need to keep in mind that it may have a negative impact on -latency if requests have to be paused. - -`STATIC_BATCH` is a legacy mode and is not recommended for production usage. - -To switch the capacity scheduler policy from the default of `GUARANTEED_NO_EVICT` to `MAX_UTILIZATION` -you would modify the [LLM-API end-to-end example](./benchmarking-default-performance.md#before-you-begin-tensorrt-llm-llm-api) to be: - -```python -from tensorrt_llm import LLM, SamplingParams -from tensorrt_llm.bindings.executor import SchedulerConfig, CapacitySchedulerPolicy - - -def main(): - prompts = [ - "Hello, I am", - "The president of the United States is", - "The capital of France is", - "The future of AI is", - ] - - sampling_params = SamplingParams(temperature=0.8, top_p=0.95) - - scheduler_config = SchedulerConfig( - capacity_scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION - ) - - llm = LLM( - model="meta-llama/Llama-3.3-70B-Instruct", - tensor_parallel_size=4, - scheduler_config=scheduler_config - ) - - outputs = llm.generate(prompts, sampling_params) - - # Print the outputs. - for output in outputs: - prompt = output.prompt - generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - -if __name__ == '__main__': - main() -``` - -## Context Chunking Policy - -As discussed [previously](tuning-max-batch-size-and-max-num-tokens.md#revisiting-paged-context-attention-and-context-chunking) context chunking will increase the chance of batch processing between -the context and the generation phase, thereby balancing the calculation amount -of each iteration and typically increasing throughput. - -TensorRT-LLM currently supports two context chunking policies: `FIRST_COME_FIRST_SERVED` (default) which would prioritize scheduling all the context chunks of a request that comes in first, - and `EQUAL_PROGRESS` which schedules context chunks from all requests before scheduling the next chunk of any request. - -`FIRST_COME_FIRST_SERVED` should achieve overall better performance, while -`EQUAL_PROGRESS` can be helpful in theory to make sure time to first token (TTFT) -for most requests are relatively similar. - -To switch the context chunking policy from the default of `FIRST_COME_FIRST_SERVED` to `EQUAL_PROGRESS` -you would modify the [LLM-API end-to-end example](./benchmarking-default-performance.md#before-you-begin-tensorrt-llm-llm-api) to be: - -```python -from tensorrt_llm import LLM, SamplingParams -from tensorrt_llm.bindings.executor import SchedulerConfig, ContextChunkingPolicy - - -def main(): - prompts = [ - "Hello, I am", - "The president of the United States is", - "The capital of France is", - "The future of AI is", - ] - - sampling_params = SamplingParams(temperature=0.8, top_p=0.95) - - scheduler_config = SchedulerConfig( - context_chunking_policy=ContextChunkingPolicy.EQUAL_PROGRESS - ) - - llm = LLM( - model="meta-llama/Llama-3.3-70B-Instruct", - tensor_parallel_size=4, - scheduler_config=scheduler_config - ) - - outputs = llm.generate(prompts, sampling_params) - - # Print the outputs. - for output in outputs: - prompt = output.prompt - generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - -if __name__ == '__main__': - main() -``` - -## Max Tokens in Paged KV Cache and KV Cache Free GPU Memory Fraction - -The `max_tokens_in_paged_kv_cache` and `kv_cache_free_gpu_mem_fraction` -parameters can be used to control the maximum number of tokens handled by the -KV cache manager. Setting them properly helps better control the amount of -available memory for the KV cache manager during inference. Keeping in mind -that increasing the amount of memory available to the KV cache manager tends to -translate to a higher achievable throughput. - -The `max_tokens_in_paged_kv_cache` flag directly sets the maximum number of -tokens in the KV cache manager. When left unset, that value will be computed -based on the `kv_cache_free_gpu_mem_fraction` setting. - -The `kv_cache_free_gpu_mem_fraction` is a floating-point number between `0.0` -and `1.0` that indicates the maximum fraction of GPU memory (after loading the -model) that will be used for the KV cache. The default value is `0.90` and -means that 90% of the free GPU memory will be used to save tokens in the KV -cache. Based on that value, TensorRT-LLM can determine the maximum number of -tokens in the KV cache manager. - -When both parameters are set, the maximum number of tokens in the KV cache -manager will be set to the smaller value between `max_tokens_in_paged_kv_cache` -and the value computed from the amount of memory available for the KV cache. - -Unless users clearly know the maximum number of tokens in the KV cache needed -by the model, it is recommended to leave `max_tokens_in_paged_kv_cache` unset. -For `kv_cache_free_gpu_mem_fraction`, if no other programs are executed on the -same GPU, it is recommended to test with a as high value as `0.95` to target a -high throughput. Note that the `kv_cache_free_gpu_mem_fraction` parameter -cannot be set to `1.0` because some amount of memory has to be reserved for -inputs and outputs. - -To set `kv_cache_free_gpu_mem_fraction` you would modify the [LLM-API end-to-end example](./benchmarking-default-performance.md#before-you-begin-tensorrt-llm-llm-api) to be: - -```python -from tensorrt_llm import LLM, SamplingParams -from tensorrt_llm.bindings.executor import KvCacheConfig - - -def main(): - prompts = [ - "Hello, I am", - "The president of the United States is", - "The capital of France is", - "The future of AI is", - ] - - sampling_params = SamplingParams(temperature=0.8, top_p=0.95) - - kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.95) - - llm = LLM( - model="meta-llama/Llama-3.3-70B-Instruct", - tensor_parallel_size=8, - kv_cache_config=kv_cache_config - ) - - outputs = llm.generate(prompts, sampling_params) - - # Print the outputs. - for output in outputs: - prompt = output.prompt - generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - -if __name__ == '__main__': - main() -``` -If you wanted to set `max_tokens_in_paged_kv_cache` instead, you would replace `free_gpu_memory_fraction` with `max_tokens` and specify the number. - -```python - kv_cache_config = KvCacheConfig(max_tokens=) -``` - - -## Maximum Attention Window Size - -The `max_attention_window_size` flag sets the maximum number of tokens that are -attended to in order to generate one token when using techniques like sliding window -attention. See this -[Document](../../advanced/gpt-attention.md#sliding-window-attention-cyclic-rolling-buffer-kv-cache) -for more details. It defaults to the maximum sequence length -(`max_seq_len` when building the engine), which means -that the feature is disabled by default. - -When set to a smaller value than `max_seq_len` (during -engine build), only the KV cache of the last `max_attention_window_size` tokens -will be stored. If the input sequence length at runtime exceeds the -`max_attention_window_size` value, the accuracy may start dropping, but the -runtime performance will be better (due to the reduction in terms of -computations and GPU memory allocation). Users can modify that value to -increase runtime performance at the expense of reduced accuracy. - -Just like [`kv_cache_free_gpu_mem_fraction`](./useful-runtime-flags.md#max-tokens-in-paged-kv-cache-and-kv-cache-free-gpu-memory-fraction), `max_attention_window_size` can be specified in the LLM-API -via `KVCacheConfig`. To specify `max_attention_window_size` you would instantiate `KVCacheConfig` like so - -```python - kv_cache_config = KvCacheConfig(max_attention_window=) -``` diff --git a/docs/source/legacy/python-api/tensorrt_llm.functional.rst b/docs/source/legacy/python-api/tensorrt_llm.functional.rst deleted file mode 100644 index bfa0368f1793..000000000000 --- a/docs/source/legacy/python-api/tensorrt_llm.functional.rst +++ /dev/null @@ -1,11 +0,0 @@ -Functionals -=========================== - -.. automodule:: tensorrt_llm - -.. currentmodule:: tensorrt_llm - -.. automodule:: tensorrt_llm.functional - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/legacy/python-api/tensorrt_llm.layers.rst b/docs/source/legacy/python-api/tensorrt_llm.layers.rst deleted file mode 100644 index 7baa17ab0d67..000000000000 --- a/docs/source/legacy/python-api/tensorrt_llm.layers.rst +++ /dev/null @@ -1,69 +0,0 @@ -Layers -=========================== - -.. automodule:: tensorrt_llm - -.. currentmodule:: tensorrt_llm - -Activation ------------- -.. automodule:: tensorrt_llm.layers.activation - :members: - :undoc-members: - :show-inheritance: - -Attention ------------- -.. automodule:: tensorrt_llm.layers.attention - :members: - :undoc-members: - :show-inheritance: - -Cast ------------- -.. automodule:: tensorrt_llm.layers.cast - :members: - :undoc-members: - :show-inheritance: - -Conv ------------- -.. automodule:: tensorrt_llm.layers.conv - :members: - :undoc-members: - :show-inheritance: - -Embedding ------------- -.. automodule:: tensorrt_llm.layers.embedding - :members: - :undoc-members: - :show-inheritance: - -Linear ------------- -.. automodule:: tensorrt_llm.layers.linear - :members: - :undoc-members: - :show-inheritance: - -MLP ------------- -.. automodule:: tensorrt_llm.layers.mlp - :members: - :undoc-members: - :show-inheritance: - -Normalization ---------------- -.. automodule:: tensorrt_llm.layers.normalization - :members: - :undoc-members: - :show-inheritance: - -Pooling ------------- -.. automodule:: tensorrt_llm.layers.pooling - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/legacy/python-api/tensorrt_llm.models.rst b/docs/source/legacy/python-api/tensorrt_llm.models.rst deleted file mode 100644 index ca0a5c06f50c..000000000000 --- a/docs/source/legacy/python-api/tensorrt_llm.models.rst +++ /dev/null @@ -1,11 +0,0 @@ -Models -=========================== - -.. automodule:: tensorrt_llm - -.. currentmodule:: tensorrt_llm - -.. automodule:: tensorrt_llm.models - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/legacy/python-api/tensorrt_llm.plugin.rst b/docs/source/legacy/python-api/tensorrt_llm.plugin.rst deleted file mode 100644 index b4459f98e2ce..000000000000 --- a/docs/source/legacy/python-api/tensorrt_llm.plugin.rst +++ /dev/null @@ -1,10 +0,0 @@ -Plugin -=========================== - -.. automodule:: tensorrt_llm - -.. currentmodule:: tensorrt_llm - -.. automodule:: tensorrt_llm.plugin - :members: - :show-inheritance: diff --git a/docs/source/legacy/python-api/tensorrt_llm.quantization.rst b/docs/source/legacy/python-api/tensorrt_llm.quantization.rst deleted file mode 100644 index 1da9ec81fb6e..000000000000 --- a/docs/source/legacy/python-api/tensorrt_llm.quantization.rst +++ /dev/null @@ -1,10 +0,0 @@ -Quantization -=========================== - -.. automodule:: tensorrt_llm - -.. currentmodule:: tensorrt_llm - -.. automodule:: tensorrt_llm.quantization - :members: - :show-inheritance: diff --git a/docs/source/legacy/python-api/tensorrt_llm.runtime.rst b/docs/source/legacy/python-api/tensorrt_llm.runtime.rst deleted file mode 100644 index 2bd02e3dca5c..000000000000 --- a/docs/source/legacy/python-api/tensorrt_llm.runtime.rst +++ /dev/null @@ -1,11 +0,0 @@ -Runtime -=========================== - -.. automodule:: tensorrt_llm - -.. currentmodule:: tensorrt_llm - -.. automodule:: tensorrt_llm.runtime - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/legacy/reference/memory.md b/docs/source/legacy/reference/memory.md deleted file mode 100644 index b20429e061c2..000000000000 --- a/docs/source/legacy/reference/memory.md +++ /dev/null @@ -1,140 +0,0 @@ -(memory)= - -# 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 I/O tensors. For I/O tensors, the major memory footprint comes from the KV cache tensor. - - -### 1. 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. - - -### 2. 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 profile’s 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 **cannot be changed**, and can be queried by the API `trt.ICudaEngine.device_memory_size_v2`. - -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 number of input tokens (`max_num_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 `PretrainedModel.prepare_inputs` to affect the memory usage, or one can change the command line options of the `trtllm-build` command used in the examples. - - When using the [packed tensors](../advanced/gpt-attention.md#padded-and-packed-tensors) format and `max_num_tokens` is specified, reducing its value will also reduce activation memory size. - - When using the [padded tensors](../advanced/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. - - 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](../advanced/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. - - -### 3. I/O tensors - -#### 3.1 Runtime and decoder buffers except KV cache tensor - -##### C++ runtime - -Before KV cache blocks are allocated, some amount of GPU memory are pre-allocated by C++ runtime for storing I/O tensors of TensorRT engine and the decoupled dynamic decoder, it's allocated based on runtime max_batch_size and max_seq_len so that OOM can be avoided when there are indeed that amount of requests scheduled. - -#### 3.2 KV cache tensor - -##### C++ runtime - - TensorRT-LLM runtime pre-allocates paged KV cache pools 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 the `Executor`. If neither `maxTokens` nor `freeGpuMemoryFraction` is specified, KV cache will by default allocate 90% 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). - -##### Python runtime (Not recommended to be used) - -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.** - -## Memory pool - -TensorRT-LLM C++ runtime is using stream-ordered memory allocator to allocate and free buffers, see [BufferManager::initMemoryPool](source:cpp/tensorrt_llm/runtime/bufferManager.cpp), which uses the default memory pool managed by the CUDA driver. When a `TrtGptModel` object is destroyed, memory is returned to the memory pool and can be reused by the next instance of a `TrtGptModel` 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())](source:cpp/tensorrt_llm/runtime/bufferManager.cpp) and [BufferManager::memoryPoolFree())](source:cpp/tensorrt_llm/runtime/bufferManager.cpp), respectively. - -## Known Issues - - -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. How to debug the memory usage of TensorRT-LLM? - - When the `info` logging level is used, TensorRT and TensorRT-LLM will print messages about memory usage details. Here is part of a log example with `info` logging level at runtime: - ``` - [TensorRT-LLM][INFO] Loaded engine size: 6695 MiB - [TensorRT-LLM][INFO] [MemUsageChange] Allocated 1134.01 MiB for execution context memory. - [TensorRT-LLM][INFO] [MS] Running engine with multi stream info - [TensorRT-LLM][INFO] [MS] Number of aux streams is 1 - [TensorRT-LLM][INFO] [MS] Number of total worker streams is 2 - [TensorRT-LLM][INFO] [MS] The main stream provided by execute/enqueue calls is the first worker stream - [TensorRT-LLM][INFO] [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +0, GPU +0, now: CPU 0, GPU 6678 (MiB) - [TensorRT-LLM][INFO] [MemUsageChange] Allocated 43.29 MB GPU memory for runtime buffers. - [TensorRT-LLM][INFO] [MemUsageChange] Allocated 180.30 MB GPU memory for decoder. - [TensorRT-LLM][INFO] Memory usage when calculating max tokens in paged kv cache: total: 79.10 GiB, available: 70.48 GiB - [TensorRT-LLM][INFO] Number of blocks in KV cache primary pool: 4060 - [TensorRT-LLM][INFO] Number of blocks in KV cache secondary pool: 0, onboard blocks to primary memory before reuse: true - [TensorRT-LLM][INFO] Max KV cache pages per sequence: 32 - [TensorRT-LLM][INFO] Number of tokens per block: 64. - [TensorRT-LLM][INFO] [MemUsageChange] Allocated 63.44 GiB for max tokens in paged KV cache (259840). - ``` - You can see that there are several GPU memory allocation started with `[MemUsageChange]` keyword happened at runtime. - - 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. - -2. 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`, see [Activation size](#activation-size) for details. - - -3. 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. - -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/legacy/reference/multimodal-feature-support-matrix.md b/docs/source/legacy/reference/multimodal-feature-support-matrix.md deleted file mode 100644 index b6d99e24ca69..000000000000 --- a/docs/source/legacy/reference/multimodal-feature-support-matrix.md +++ /dev/null @@ -1,13 +0,0 @@ -# Multimodal Feature Support Matrix (PyTorch Backend) - -| Model | CUDA Graph | Encoder IFB | KV Cache Reuse | Chunked Prefill | -| :----------------- | :--------- | :------------------ | :------------- | :-------------- | -| Gemma 3 | Yes | Yes | N/A | N/A | -| HyperCLOVA | Yes | Yes | No | No | -| VILA | Yes | No | No | No | -| LLaVA-NeXT | Yes | Yes | Yes | Yes | -| Llama 4 | Yes | Yes | No | No | -| Mistral-Small-3.1 | Yes | Yes | Yes | Yes | -| Phi-4-multimodal | Yes | Yes | Yes | Yes | -| Qwen2-VL | Yes | Yes | Yes | Yes | -| Qwen2.5-VL | Yes | Yes | Yes | Yes | diff --git a/docs/source/legacy/reference/precision.md b/docs/source/legacy/reference/precision.md deleted file mode 100644 index b31eff6d6238..000000000000 --- a/docs/source/legacy/reference/precision.md +++ /dev/null @@ -1,184 +0,0 @@ -(precision)= - -# Numerical Precision - -This document describes the different quantization recipes implemented in TensorRT-LLM and contains a support matrix -for the different models. - -## FP32, FP16 and BF16 - -The different models implemented in TensorRT-LLM work with 32-bit IEEE -floating-point (FP32) numbers. When checkpoints are available, the models also -support 16-bit IEEE floating-point numbers (FP16) and 16-bit Bfloat16 (BF16) as -described [here](https://en.wikipedia.org/wiki/Bfloat16_floating-point_format). - -## Quantization and Dequantization (Q/DQ) - -Given a floating-point number `x` and a floating-point scaling factor `s`, -TensorRT-LLM implements INT8 quantization as: - -``` -q = int8.satfinite(x * s) -``` - -Given an INT8 number `q` and a floating-point scaling factor `s`, TensorRT-LLM -implements INT8 dequantization to the floating-point (FP) type as: - -``` -x = static_cast(q) * s -``` - -Given a matrix (2D tensor) of shape `M x N` (`M` rows and `N` columns) where -`M` is the number of tokens and `N` is the number of channels. TensorRT-LLM has -the three following modes to quantize and dequantize the elements of the -tensor: - - * Per-tensor: It uses a single scaling factor for all the elements, - * Per-token: It uses a different scaling factor for each token. There are `M` - scaling factors in that case, - * Per-channel: It uses a different scaling factor for each channel. There are - `N` scaling factors in that case. - -Note that per-token and per-channel scaling modes can be used together (i.e. -they are _not_ mutually exclusive). - -In pseudo-code, the quantization can be implemented as follows for the three -different modes: - -```python -# Per-tensor scaling. -for mi in range(M): - for ni in range(N): - q[mi][ni] = int8.satfinite(x[mi][ni] * s) - -# Per-token scaling. -for mi in range(M): - for ni in range(N): - q[mi][ni] = int8.satfinite(x[mi][ni] * s[mi]) - -# Per-channel scaling. -for mi in range(M): - for ni in range(N): - q[mi][ni] = int8.satfinite(x[mi][ni] * s[ni]) -``` - -## INT8 SmoothQuant (W8A8) - -The SmoothQuant technique was introduced in -[https://arxiv.org/abs/2211.10438](https://arxiv.org/abs/2211.10438). It is a -method to run inference using INT8 for both activations and weights while -maintaining the accuracy of the network (on downstream tasks). - -As explained in the research paper, preprocessing must be applied to the -weights of the model. TensorRT-LLM includes scripts to prepare the model to -run using the SmoothQuant method. - -Examples of how to enable SmoothQuant for GPT, GPT-J and LLaMA can be found in -the [examples/quantization](source:examples/quantization) folder of that release. - -## INT4 and INT8 Weight-Only (W4A16 and W8A16) - -The INT4 and INT8 Weight-Only techniques consist in quantizing the weights of -a model and dequantizing those weights on-the-fly in linear layers (Matmuls). -The activations are encoded using floating-point values (FP16 or BF16). - -To use INT4/INT8 Weight-Only methods, the user must determine the scaling -factors to use to quantize and dequantize the weights of the model. - -This release includes examples for [GPT](source:examples/models/core/gpt) and -[LLaMA](source:examples/models/core/llama). - -## GPTQ and AWQ (W4A16) - -The GPTQ and AWQ techniques are presented in -[https://arxiv.org/abs/2210.17323](https://arxiv.org/abs/2210.17323) -and -[https://arxiv.org/abs/2306.00978](https://arxiv.org/abs/2306.00978), -respectively. TensorRT-LLM supports per-group scaling factors and -zero-offsetting in linear layers to implement GPTQ and AWQ methods. See the -[WeightOnlyGroupwiseQuantMatmulPlugin](source:cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin) -plugin and the corresponding -[`weight_only_groupwise_quant_matmul`](source:tensorrt_llm/quantization/functional.py) -Python function, for details. - -This release includes examples of applying GPTQ to [GPT-NeoX](source:examples/models/core/gpt) -and [LLaMA-v2](source:examples/models/core/llama), as well as an example of using AWQ with -[GPT-J](source:examples/models/contrib/gptj). - -## FP8 (Hopper) - -This release of TensorRT-LLM contains implementations of FP8 for GPT-NeMo, -GPT-J and LLaMA. Those examples can be found in -[examples/quantization](source:examples/quantization). - -## NVFP4 (Blackwell) - -LLama and Mixtral can run in NVFP4 datatype. Those examples can be found in Llama examples. - -## Support matrix - -This release of TensorRT-LLM contains the following examples: - -| Model | FP32 | FP16 | BF16 | FP8 | NVFP4 | W8A8 SQ | W8A16 | W4A16 | W4A16 AWQ | W4A16 GPTQ | -| :------------- | :---: | :---: | :---: | :---: | :---: | :-----: | :---: | :---: | :-------: | :--------: | -| Baichuan | Y | Y | Y | Y | . | Y | Y | Y | Y | Y | -| BERT | Y | Y | Y | . | . | . | . | . | . | . | -| BLIP-2 | Y | Y | Y | . | . | . | . | . | . | . | -| BLOOM | Y | Y | Y | Y | . | Y | Y | Y | . | . | -| ChatGLM | Y | Y | Y | . | . | . | . | . | . | . | -| ChatGLM-v2 | Y | Y | Y | . | . | . | . | . | . | . | -| ChatGLM-v3 | Y | Y | Y | . | . | . | . | . | . | . | -| DBRX | Y | Y | Y | . | . | . | Y | Y | . | . | -| Falcon | Y | Y | Y | Y | . | . | Y | Y | Y | . | -| Flan-T5 | Y | Y | Y | . | . | . | . | . | . | . | -| Gemma | Y | Y | Y | Y | . | Y | 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 | . | . | -| InternLM2 | Y | Y | Y | . | . | . | . | . | . | . | -| LLaMA | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | -| LLaMA-v2 | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | -| Mamba | Y | Y | Y | . | . | . | . | . | . | . | -| Mistral | Y | Y | Y | Y | . | Y | Y | Y | Y | . | -| Mixtral | Y | Y | Y | Y | Y | . | Y | Y | . | . | -| MPT | Y | Y | Y | Y | . | Y | Y | Y | Y | . | -| OPT | Y | Y | Y | . | . | . | . | . | . | . | -| Phi | Y | Y | Y | . | . | . | . | . | . | . | -| Qwen | Y | Y | Y | . | . | Y | Y | Y | Y | Y | -| RecurrentGemma | Y | Y | Y | Y | . | Y | . | . | Y | . | -| Replit Code | Y | Y | Y | . | . | . | . | . | . | . | -| SantaCoder | Y | Y | Y | . | . | . | Y | Y | . | . | -| Skywork | Y | Y | Y | . | . | . | . | . | . | . | -| StarCoder1 | Y | Y | Y | . | . | . | Y | Y | . | . | -| StarCoder2 | Y | Y | Y | Y | . | . | Y | Y | . | . | -| T5 | Y | Y | Y | . | . | . | . | . | . | . | -| Whisper | Y | Y | Y | . | . | . | Y | Y | . | . | -| BLIP2-OPT | Y | Y | Y | . | . | . | . | . | . | . | -| BLIP2-T5 | Y | Y | Y | . | . | . | . | . | . | . | -| LLaVA | Y | Y | Y | Y | . | Y | Y | Y | Y | Y | -| VILA | Y | Y | Y | Y | . | Y | Y | Y | Y | Y | -| Nougat | Y | Y | Y | . | . | . | . | . | . | . | - -Note: The vision component of multi-modal models(BLIP2-OPT/BLIP2-T5/LLaVA/VILA/Nougat) uses FP16 by default. -The language component decides which quantization methods are supported by a given multi-modal model. - -## Technical Detail: The `QuantMode` Flags - -The quantization method is controlled by the -[`QuantMode`](source:tensorrt_llm/quantization/mode.py) flags. The different fields -are: - - * `INT4_WEIGHTS`, the weights are quantized to 4 bits (W4A\*), - * `INT8_WEIGHTS`, the weights are quantized to 8 bits (W8A\*), - * `ACTIVATIONS`, the activations are quantized to 8 bits (W\*A8), - * `PER_CHANNEL`, the scaling factors are defined per channel, - * `PER_TOKEN`, the scaling factors are defined per token, - * `PER_GROUP`, the scaling factors are defined per group. - -There are three additional flags to control TensorRT-LLM: - - * `INT8_KV_CACHE`, the K/V cache stores K and V using 8-bit integers, - * `FP8_KV_CACHE`, the K/V cache stores K and V using 8-bit floating-point numbers, - * `FP8_QDQ`, TensorRT-LLM relies on automatic fusion of Q/DQ nodes in TensorRT. diff --git a/docs/source/legacy/reference/support-matrix.md b/docs/source/legacy/reference/support-matrix.md deleted file mode 100644 index e63f439b9446..000000000000 --- a/docs/source/legacy/reference/support-matrix.md +++ /dev/null @@ -1,185 +0,0 @@ -# Support Matrix - -```{deprecated} -This page is outdated and no longer maintained. Please refer to the up-to-date [Supported Models](../../models/supported-models.md) page instead. -``` - -TensorRT-LLM optimizes the performance of a range of well-known models on NVIDIA GPUs. The following sections provide a list of supported GPU architectures as well as important features implemented in TensorRT-LLM. - -## Models (PyTorch Backend) - -| Architecture | Model | HuggingFace Example | Modality | -|--------------|-------|---------------------|----------| -| `BertForSequenceClassification` | BERT-based | `textattack/bert-base-uncased-yelp-polarity` | L | -| `DeciLMForCausalLM` | Nemotron | `nvidia/Llama-3_1-Nemotron-51B-Instruct` | L | -| `DeepseekV3ForCausalLM` | DeepSeek-V3 | `deepseek-ai/DeepSeek-V3 `| L | -| `Exaone4ForCausalLM` | EXAONE 4.0 | `LGAI-EXAONE/EXAONE-4.0-32B` | L | -| `Gemma3ForCausalLM` | Gemma 3 | `google/gemma-3-1b-it` | L | -| `Gemma3ForConditionalGeneration` | Gemma 3 | `google/gemma-3-27b-it` | L + I | -| `HCXVisionForCausalLM` | HyperCLOVAX-SEED-Vision | `naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B` | L + I | -| `LlavaLlamaModel` | VILA | `Efficient-Large-Model/NVILA-8B` | L + I + V | -| `LlavaNextForConditionalGeneration` | LLaVA-NeXT | `llava-hf/llava-v1.6-mistral-7b-hf` | L + I | -| `LlamaForCausalLM` | Llama 3.1, Llama 3, Llama 2, LLaMA | `meta-llama/Meta-Llama-3.1-70B` | L | -| `Llama4ForConditionalGeneration` | Llama 4 | `meta-llama/Llama-4-Scout-17B-16E-Instruct` | L + I | -| `MistralForCausalLM` | Bielik | `speakleash/Bielik-11B-v2.2-Instruct` | L | -| `MistralForCausalLM` | Mistral | `mistralai/Mistral-7B-v0.1` | L | -| `Mistral3ForConditionalGeneration` | Mistral3 | `mistralai/Mistral-Small-3.1-24B-Instruct-2503` | L + I | -| `MixtralForCausalLM` | Mixtral | `mistralai/Mixtral-8x7B-v0.1` | L | -| `MllamaForConditionalGeneration` | Llama 3.2 | `meta-llama/Llama-3.2-11B-Vision` | L | -| `NemotronForCausalLM` | Nemotron-3, Nemotron-4, Minitron | `nvidia/Minitron-8B-Base` | L | -| `NemotronNASForCausalLM` | NemotronNAS | `nvidia/Llama-3_3-Nemotron-Super-49B-v1` | L | -| `Phi3ForCausalLM` | Phi-4 | `microsoft/Phi-4` | L | -| `Phi4MMForCausalLM` | Phi-4-multimodal | `microsoft/Phi-4-multimodal-instruct` | L + I + A | -| `Qwen2ForCausalLM` | QwQ, Qwen2 | `Qwen/Qwen2-7B-Instruct` | L | -| `Qwen2ForProcessRewardModel` | Qwen2-based | `Qwen/Qwen2.5-Math-PRM-7B` | L | -| `Qwen2ForRewardModel` | Qwen2-based | `Qwen/Qwen2.5-Math-RM-72B` | L | -| `Qwen2VLForConditionalGeneration` | Qwen2-VL | `Qwen/Qwen2-VL-7B-Instruct` | L + I + V | -| `Qwen2_5_VLForConditionalGeneration` | Qwen2.5-VL | `Qwen/Qwen2.5-VL-7B-Instruct` | L + I + V | -| `Qwen3ForCausalLM` | Qwen3 | `Qwen/Qwen3-8B` | L | -| `Qwen3MoeForCausalLM` | Qwen3MoE | `Qwen/Qwen3-30B-A3B` | L | - -Note: -- L: Language -- I: Image -- V: Video -- A: Audio - -## Models (TensorRT Backend) - -### LLM Models - -- [Arctic](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/arctic) -- [Baichuan/Baichuan2](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/baichuan) -- [BART](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/enc_dec) -- [BERT](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/bert) -- [BLOOM](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/bloom) -- [ByT5](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/enc_dec) -- [ChatGLM](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/chatglm-6b) -- [ChatGLM2](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/chatglm2-6b) -- [ChatGLM3](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/chatglm3-6b-32k) -- [Code LLaMA](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/llama) -- [DBRX](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/dbrx) -- [Exaone](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/exaone) -- [FairSeq NMT](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/enc_dec) -- [Falcon](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/falcon) -- [Flan-T5](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/enc_dec) [^encdec] -- [Gemma/Gemma2](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/gemma) -- [GLM-4](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/glm-4-9b) -- [GPT](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/gpt) -- [GPT-J](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/gptj) -- [GPT-Nemo](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/gpt) -- [GPT-NeoX](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/gptneox) -- [Granite-3.0](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/granite) -- [Grok-1](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/grok) -- [InternLM](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples//models/contrib/internlm) -- [InternLM2](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/internlm2) -- [LLaMA/LLaMA 2/LLaMA 3/LLaMA 3.1](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/llama) -- [Mamba](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/mamba) -- [mBART](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/enc_dec) -- [Minitron](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/nemotron) -- [Mistral](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/llama) -- [Mistral NeMo](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/llama) -- [Mixtral](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/mixtral) -- [MPT](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/mpt) -- [Nemotron](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/nemotron) -- [mT5](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/enc_dec) -- [OPT](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/opt) -- [Phi-1.5/Phi-2/Phi-3](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/phi) -- [Qwen/Qwen1.5/Qwen2/Qwen3](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/qwen) -- [Qwen-VL](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/qwenvl) -- [RecurrentGemma](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/recurrentgemma) -- [Replit Code](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/mpt) [^replitcode] -- [RoBERTa](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/bert) -- [SantaCoder](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/gpt) -- [Skywork](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/skywork) -- [Smaug](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/smaug) -- [StarCoder](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/gpt) -- [T5](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/enc_dec) -- [Whisper](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/whisper) - - -### Multi-Modal Models [^multimod] - -- [BLIP2 w/ OPT](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [BLIP2 w/ T5](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [CogVLM](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) [^bf16only] -- [Deplot](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [Fuyu](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [Kosmos](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [LLaVA-v1.5](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [LLaVa-Next](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [LLaVa-OneVision](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [NeVA](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [Nougat](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [Phi-3-vision](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [Video NeVA](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [VILA](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [MLLaMA](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [LLama 3.2 VLM](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) - - -(support-matrix-hardware)= -## Hardware - -The following table shows the supported hardware for TensorRT-LLM. - -If a GPU architecture is not listed, the TensorRT-LLM team does not develop or test the software on the architecture and support is limited to community support. -In addition, older architectures can have limitations for newer software releases. - -```{list-table} -:header-rows: 1 -:widths: 20 80 - -* - - - Hardware Compatibility -* - Operating System - - TensorRT-LLM requires Linux x86_64 or Linux aarch64. -* - GPU Model Architectures - - - - [NVIDIA GB300 NVL72](https://www.nvidia.com/en-us/data-center/gb300-nvl72/) - - [NVIDIA GB200 NVL72](https://www.nvidia.com/en-us/data-center/gb200-nvl72/) - - [NVIDIA GB300 NVL72](https://www.nvidia.com/en-us/data-center/gb300-nvl72/) - - [NVIDIA Blackwell Architecture](https://www.nvidia.com/en-us/data-center/technologies/blackwell-architecture/) - - [NVIDIA Grace Hopper Superchip](https://www.nvidia.com/en-us/data-center/grace-hopper-superchip/) - - [NVIDIA Hopper Architecture](https://www.nvidia.com/en-us/data-center/technologies/hopper-architecture/) - - [NVIDIA Ada Lovelace Architecture](https://www.nvidia.com/en-us/technologies/ada-architecture/) - - [NVIDIA Ampere Architecture](https://www.nvidia.com/en-us/data-center/ampere-architecture/) -``` - -(support-matrix-software)= -## Software - -The following table shows the supported software for TensorRT-LLM. - -```{list-table} -:header-rows: 1 -:widths: 20 80 - -* - - - Software Compatibility -* - Container - - [26.04](https://docs.nvidia.com/deeplearning/frameworks/support-matrix/index.html) -* - TensorRT - - [10.16](https://docs.nvidia.com/deeplearning/tensorrt/release-notes/index.html) -* - Precision - - - - Blackwell (SM100/SM103/SM120) - FP32, FP16, BF16, FP8, FP4, INT8, INT4 - - Hopper (SM90) - FP32, FP16, BF16, FP8, INT8, INT4 - - Ada Lovelace (SM89) - FP32, FP16, BF16, FP8, INT8, INT4 - - Ampere (SM80, SM86) - FP32, FP16, BF16, INT8, INT4[^smgte89] -``` - -[^replitcode]: Replit Code is not supported with the transformers 4.45+. - -[^smgte89]: INT4 AWQ and GPTQ with FP8 activations require SM >= 89. - -[^encdec]: Encoder-Decoder provides general encoder-decoder functionality that supports many encoder-decoder models such as T5 family, BART family, Whisper family, NMT family, and so on. - -[^multimod]: Multi-modal provides general multi-modal functionality that supports many multi-modal architectures such as BLIP2 family, LLaVA family, and so on. - -[^bf16only]: Only supports bfloat16 precision. - - -```{note} -Support for FP8 and quantized data types (INT8 or INT4) is not implemented for all the models. Refer to {ref}`precision` and [examples](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples) folder for additional information. -``` diff --git a/docs/source/legacy/reference/troubleshooting.md b/docs/source/legacy/reference/troubleshooting.md deleted file mode 100644 index 1b29849bbe64..000000000000 --- a/docs/source/legacy/reference/troubleshooting.md +++ /dev/null @@ -1,328 +0,0 @@ -(troubleshooting)= - -# Troubleshooting - -This document describes some of the frequently asked questions and their solutions in TensorRT-LLM, including problems of installation, model-building, model-execution, and input / output size. - -## Installation Errors - -During compilation and installation of TensorRT-LLM, many build errors can be resolved by simply deleting the build tree and rebuilding again. - -In most occasions, these problems are caused by the workflow like: an old compilation -> some code change (update of the repo or users' writing) -> a later compilation. - -Solution: try running build script with `--clean`, or try running `rm -r build cpp/build` before running build script again. - -## Debug on Unit Tests - -Here is an example to print the values of the MLP output tensor in a unit test ([full example](../../../../tests/unittest/others/test_debugging_api.py)). - -1. Register the intermediate tensors as the network outputs with `register_network_output` API. - -```python -class MLP(Module): - - def __init__(self, ...): - super().__init__() - # Do not modify the definition in `__init__` method - self.fc = ... - self.proj = ... - - def forward(self, hidden_states): - inter = self.fc(hidden_states) - inter = tensorrt_llm.functional.relu(inter) - # Here register the tensor `inter` as our debug output tensor - self.register_network_output('inter', inter) - output = self.proj(inter) - return output -``` - -2. Mark the intermediate tensors as network outputs. - -```python -for k, v in gm.named_network_outputs(): - net._mark_output(v, k, dtype) -``` - -3. Print the tensors at runtime. - -```python -print(outputs.keys()) -print(outputs['inter']) -``` - -## Debug on E2E Models - -Here is an example to print the values of the MLP output tensor in the GPT model. - -1. Register the MLP output tensor in `tensorrt_llm/models/gpt/model.py`. - -```python - hidden_states = residual + attention_output.data - - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states) - # Register as model output - # ------------------------------------------------------ - self.register_network_output('mlp_output', hidden_states) - # ------------------------------------------------------ - - hidden_states = residual + hidden_states -``` - -2. Build the TensorRT engine of the model. - -Enable the `--enable_debug_output` option when building engines with `trtllm-build` - -```bash -cd examples/models/core/gpt - -# Download hf gpt2 model -rm -rf gpt2 && git clone https://huggingface.co/gpt2-medium gpt2 -pushd gpt2 && rm pytorch_model.bin model.safetensors && wget -q https://huggingface.co/gpt2-medium/resolve/main/pytorch_model.bin && popd - -# Convert to TensorRT-LLM checkpoint -python3 convert_checkpoint.py \ - --model_dir gpt2 \ - --dtype float16 \ - --output_dir gpt2/trt_ckpt/fp16/1-gpu - -# Build TensorRT-LLM engines with --enable_debug_output -trtllm-build \ - --checkpoint_dir gpt2/trt_ckpt/fp16/1-gpu \ - --enable_debug_output \ - --output_dir gpt2/trt_engines/fp16/1-gpu -``` - -3. Print the intermediate output tensors. - -Add debug info in `tensorrt_llm/runtime/generation.py`. - -```python - stream = torch.cuda.current_stream().cuda_stream - instance_idx = step % 2 - if self.cuda_graph_mode and self.runtime.cuda_graph_instances[ - instance_idx] is not None: - # launch cuda graph - CUASSERT( - cudart.cudaGraphLaunch( - self.runtime.cuda_graph_instances[instance_idx], stream)) - ok = True - else: - ok = self.runtime._run(context, stream) - - if not ok: - raise RuntimeError(f"Executing TRT engine failed step={step}!") - if self.debug_mode: - torch.cuda.synchronize() - # ------------------------------------------- - if step == 0: - print(self.debug_buffer.keys()) - print(f"Step: {step}") - print(self.debug_buffer['transformer.layers.6.mlp_output']) - # ------------------------------------------- -``` - -4. Run `../run.py` with `--debug_mode` and `--use_py_session`. - -```bash -python3 ../run.py \ - --engine_dir gpt2/trt_engines/fp16/1-gpu \ - --tokenizer_dir gpt2 \ - --max_output_len 8 \ - --debug_mode \ - --use_py_session -``` - -5. See the value of the tensor. - -```txt -...... -dict_keys(['context_lengths', 'cache_indirection', 'position_ids', 'logits', 'last_token_ids', 'input_ids', 'kv_cache_block_pointers', 'host_kv_cache_block_pointers', 'sequence_length', 'host_past_key_value_lengths', 'host_sink_token_length', 'host_request_types', 'host_max_attention_window_sizes', 'host_context_lengths', 'transformer.layers.0.mlp_output', 'transformer.layers.1.mlp_output', 'transformer.layers.2.mlp_output', 'transformer.layers.3.mlp_output', 'transformer.layers.4.mlp_output', 'transformer.layers.5.mlp_output', 'transformer.layers.6.mlp_output', 'transformer.layers.7.mlp_output', 'transformer.layers.8.mlp_output', 'transformer.layers.9.mlp_output', 'transformer.layers.10.mlp_output', 'transformer.layers.11.mlp_output', 'transformer.layers.12.mlp_output', 'transformer.layers.13.mlp_output', 'transformer.layers.14.mlp_output', 'transformer.layers.15.mlp_output', 'transformer.layers.16.mlp_output', 'transformer.layers.17.mlp_output', 'transformer.layers.18.mlp_output', 'transformer.layers.19.mlp_output', 'transformer.layers.20.mlp_output', 'transformer.layers.21.mlp_output', 'transformer.layers.22.mlp_output', 'transformer.layers.23.mlp_output']) -Step: 0 -tensor([[ 0.0294, -0.0260, -0.0776, ..., -0.0560, -0.0235, 0.0273], - [-0.0071, 0.5879, 0.1993, ..., -1.0449, -0.6299, 0.5957], - [-0.8779, 0.1050, 0.7090, ..., 0.0910, 1.0713, -0.2939], - ..., - [ 0.1212, -0.0903, -0.5918, ..., -0.1045, -0.3445, 0.1082], - [-1.0723, -0.0732, 0.6157, ..., 0.3452, 0.2998, 0.2649], - [-0.7134, 0.9692, -0.1141, ..., -0.0096, 0.9521, 0.1437]], - device='cuda:0', dtype=torch.float16) -Step: 1 -tensor([[-0.2107, 0.5874, 0.8179, ..., 0.7900, -0.6890, 0.6064]], - device='cuda:0', dtype=torch.float16) -Step: 2 -tensor([[ 0.4192, -0.0047, 1.3887, ..., -0.9028, -0.0682, -0.2820]], - device='cuda:0', dtype=torch.float16) -Step: 3 -tensor([[-0.7949, -0.5073, -0.1721, ..., -0.5830, -0.1378, -0.0070]], - device='cuda:0', dtype=torch.float16) -Step: 4 -tensor([[-0.0804, 0.1272, -0.6255, ..., -0.1072, -0.0523, 0.7144]], - device='cuda:0', dtype=torch.float16) -Step: 5 -tensor([[-0.3328, -0.8828, 0.3442, ..., 0.8149, -0.0630, 1.2305]], - device='cuda:0', dtype=torch.float16) -Step: 6 -tensor([[-0.2225, -0.2079, -0.1459, ..., -0.3555, -0.1672, 0.1135]], - device='cuda:0', dtype=torch.float16) -Step: 7 -tensor([[ 0.1290, -0.1556, 0.3977, ..., -0.8218, -0.3291, -0.8672]], - device='cuda:0', dtype=torch.float16) -Input [Text 0]: "Born in north-east France, Soyer trained as a" -Output [Text 0 Beam 0]: " chef before moving to London in the early" -``` - -## Debug Execution Errors - -If problems come from plugins, try setting the environment variable `CUDA_LAUNCH_BLOCKING=1` to make kernels launch synchronously with their return status checked immediately. - -If problems come from runtime-shape of the input tensors, double-check the shape (rank and length of each rank) and location (CPU / GPU) of input tensors for the engine obey the build-time setting. - -For example, one possible reason of getting the error information like below is, we use mismatched configuration between engine building and running, including code change (update of repo or users' rewriting), too large or too small input shape, etc.. - -```txt -unexpected shape for input 'XXX' for model 'YYY'. Expected [-1,-1,-1], got [8,16]. NOTE: Setting a non-zero max_batch_size in the model config requires a batch dimension to be prepended to each input shape. If you want to specify the full shape including the batch dim in your input dims config, try setting max_batch_size to zero. See the model configuration docs for more info on max_batch_size. - -[TensorRT-LLM][ERROR] Assertion failed: Tensor 'input_ids' has invalid shape (8192), expected (-1) (/code/tensorrt_llm/cpp/tensorrt_llm/runtime/tllmRuntime.cpp:149) - -RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 8192 but got size 1024 for tensor number 1 in the list. -``` - -By setting environment variable `export TLLM_LOG_LEVEL=TRACE`, we can get more information about the TensorRT engine and context at runtime. - -Before the first forward computation, the shapes of all input / output tensors and their corresponding allowed ranges are provided in a table like: - -```txt -[TensorRT-LLM][TRACE] Information of engine input / output. -[TensorRT-LLM][TRACE] ===================================================================== -[TensorRT-LLM][TRACE] Name |I/O|Location|DataType| Shape | -[TensorRT-LLM][TRACE] --------------------------------------------------------------------- -[TensorRT-LLM][TRACE] input_ids | I | GPU | INT32 | (-1) | -[TensorRT-LLM][TRACE] position_ids | I | GPU | INT32 | (-1) | -[TensorRT-LLM][TRACE] last_token_ids | I | GPU | INT32 | (-1) | -[TensorRT-LLM][TRACE] kv_cache_block_offsets | I | GPU | INT32 |(1, -1, 2, -1)| -[TensorRT-LLM][TRACE] host_kv_cache_block_offsets | I | GPU | INT32 |(1, -1, 2, -1)| -[TensorRT-LLM][TRACE] host_kv_cache_pool_pointers | I | GPU | INT64 | (1, 2) | -[TensorRT-LLM][TRACE] host_kv_cache_pool_mapping | I | GPU | INT32 | (28) | -[TensorRT-LLM][TRACE] sequence_length | I | GPU | INT32 | (-1) | -[TensorRT-LLM][TRACE] host_request_types | I | GPU | INT32 | (-1) | -[TensorRT-LLM][TRACE] host_past_key_value_lengths | I | GPU | INT32 | (-1) | -[TensorRT-LLM][TRACE] context_lengths | I | GPU | INT32 | (-1) | -[TensorRT-LLM][TRACE] host_runtime_perf_knobs | I | GPU | INT64 | (16) | -[TensorRT-LLM][TRACE] host_context_lengths | I | GPU | INT32 | (-1) | -[TensorRT-LLM][TRACE] host_max_attention_window_sizes| I | GPU | INT32 | (28) | -[TensorRT-LLM][TRACE] host_sink_token_length | I | GPU | INT32 | (1) | -[TensorRT-LLM][TRACE] cache_indirection | I | GPU | INT32 | (-1, 1, -1) | -[TensorRT-LLM][TRACE] logits | O | GPU | FP32 | (-1, 65024) | -[TensorRT-LLM][TRACE] ===================================================================== -[TensorRT-LLM][TRACE] Information of optimization profile. -[TensorRT-LLM][TRACE] Optimization Profile 0: -[TensorRT-LLM][TRACE] ============================================================================= -[TensorRT-LLM][TRACE] Name | Min | Opt | Max | -[TensorRT-LLM][TRACE] ----------------------------------------------------------------------------- -[TensorRT-LLM][TRACE] input_ids | (1) | (8) | (8192) | -[TensorRT-LLM][TRACE] position_ids | (1) | (8) | (8192) | -[TensorRT-LLM][TRACE] last_token_ids | (1) | (4) | (8) | -[TensorRT-LLM][TRACE] kv_cache_block_offsets | (1, 1, 2, 1) |(1, 4, 2, 16) |(1, 8, 2, 32) | -[TensorRT-LLM][TRACE] host_kv_cache_block_offsets | (1, 1, 2, 1) |(1, 4, 2, 16) |(1, 8, 2, 32) | -[TensorRT-LLM][TRACE] host_kv_cache_pool_pointers | (1, 2) | (1, 2) | (1, 2) | -[TensorRT-LLM][TRACE] host_kv_cache_pool_mapping | (28) | (28) | (28) | -[TensorRT-LLM][TRACE] sequence_length | (1) | (4) | (8) | -[TensorRT-LLM][TRACE] host_request_types | (1) | (4) | (8) | -[TensorRT-LLM][TRACE] host_past_key_value_lengths | (1) | (4) | (8) | -[TensorRT-LLM][TRACE] context_lengths | (1) | (4) | (8) | -[TensorRT-LLM][TRACE] host_runtime_perf_knobs | (16) | (16) | (16) | -[TensorRT-LLM][TRACE] host_context_lengths | (1) | (4) | (8) | -[TensorRT-LLM][TRACE] host_max_attention_window_sizes| (28) | (28) | (28) | -[TensorRT-LLM][TRACE] host_sink_token_length | (1) | (1) | (1) | -[TensorRT-LLM][TRACE] cache_indirection | (1, 1, 1) | (4, 1, 1024) | (8, 1, 2048) | -[TensorRT-LLM][TRACE] logits | (1, 65024) | (4, 65024) | (8, 65024) | -[TensorRT-LLM][TRACE] ============================================================================= -``` - -Before each forward computation, the real shapes of all input / output tensors for TRT engine are provided by a table like: - -```txt -[TensorRT-LLM][TRACE] Information of context input / output. -[TensorRT-LLM][TRACE] Using Optimization Profile: 0 -[TensorRT-LLM][TRACE] ================================================= -[TensorRT-LLM][TRACE] Name |I/O| Shape | -[TensorRT-LLM][TRACE] ------------------------------------------------- -[TensorRT-LLM][TRACE] input_ids | I | (33) | -[TensorRT-LLM][TRACE] position_ids | I | (33) | -[TensorRT-LLM][TRACE] last_token_ids | I | (3) | -[TensorRT-LLM][TRACE] kv_cache_block_offsets | I |(1, 3, 2, 4)| -[TensorRT-LLM][TRACE] host_kv_cache_block_offsets | I |(1, 3, 2, 4)| -[TensorRT-LLM][TRACE] host_kv_cache_pool_pointers | I | (1, 2) | -[TensorRT-LLM][TRACE] host_kv_cache_pool_mapping | I | (28) | -[TensorRT-LLM][TRACE] sequence_length | I | (3) | -[TensorRT-LLM][TRACE] host_request_types | I | (3) | -[TensorRT-LLM][TRACE] host_past_key_value_lengths | I | (3) | -[TensorRT-LLM][TRACE] context_lengths | I | (3) | -[TensorRT-LLM][TRACE] host_runtime_perf_knobs | I | (16) | -[TensorRT-LLM][TRACE] host_context_progress | I | (1) | -[TensorRT-LLM][TRACE] host_context_lengths | I | (3) | -[TensorRT-LLM][TRACE] host_max_attention_window_sizes| I | (28) | -[TensorRT-LLM][TRACE] host_sink_token_length | I | (1) | -[TensorRT-LLM][TRACE] cache_indirection | I |(3, 2, 256) | -[TensorRT-LLM][TRACE] logits | O | (3, 65024) | -[TensorRT-LLM][TRACE] ================================================= -``` - -## Tips - -* It's recommended to add options `–shm-size=1g –ulimit memlock=-1` to the - docker or nvidia-docker run command. Otherwise you may see NCCL errors when - running multiple GPU inferences. See - https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/troubleshooting.html#errors - for details. - -* When building models, memory-related issues such as -``` -[09/23/2023-03:13:00] [TRT] [E] 9: GPTLMHeadModel/layers/0/attention/qkv/PLUGIN_V2_Gemm_0: could not find any supported formats consistent with input/output data types -[09/23/2023-03:13:00] [TRT] [E] 9: [pluginV2Builder.cpp::reportPluginError::24] Error Code 9: Internal Error (GPTLMHeadModel/layers/0/attention/qkv/PLUGIN_V2_Gemm_0: could not find any supported formats consistent with input/output data types) -``` -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: `--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. --------------------------------------------------------------------------- -``` - -You may experience other problems like hanging on the program startup. - -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/models/core/gpt/build.py ...` - -It's critical that it's always `-n 1` regardless of how many GPUs are being used. If you'd use `-n 2` for a 2 GPU program it will not work. `mpirun` here isn't being used to orchestrate multiple processes, but to invoke the right environment on SLURM. The internal MPI implementation deals with spawning the additional processes. diff --git a/docs/source/legacy/tensorrt-backend-removal.md b/docs/source/legacy/tensorrt-backend-removal.md new file mode 100644 index 000000000000..0f8bf6777bad --- /dev/null +++ b/docs/source/legacy/tensorrt-backend-removal.md @@ -0,0 +1,115 @@ +# Migration Guide: TensorRT Backend Removed + +```{note} +**Breaking change.** The TensorRT engine backend has been **removed** in this +release. PyTorch is now the sole execution backend for TensorRT LLM. AutoDeploy +(built on the PyTorch backend) remains available. +``` + +This guide explains what was removed, the errors you will encounter if you rely +on the old TensorRT path, and how to migrate to the PyTorch backend. + +## What Was Removed + +The following public APIs, command-line tools, and packaging behaviors are no +longer available: + +- **`LLM(backend="tensorrt")`** — passing `backend="tensorrt"` to the `LLM` + constructor now raises a `ValueError`. PyTorch is the only supported backend. +- **`TrtLlmArgs`** — the TensorRT-specific arguments class has been removed. + Use `TorchLlmArgs` (the default) instead. +- **`tensorrt_llm._tensorrt_engine.LLM`** — the TensorRT engine `LLM` entry + point has been removed. +- **CLI build/engine tooling** — `trtllm-build`, `trtllm-refit`, and + `trtllm-prune` have been removed. The PyTorch backend loads HuggingFace + checkpoints directly and does not require a separate engine-build step. +- **`--backend tensorrt`** — the `tensorrt` choice has been removed from CLI + tools (for example `trtllm-serve`, `trtllm-bench`, `trtllm-eval`). Omit the + flag (PyTorch is the default) or pass `--backend pytorch`. +- **`tensorrt` pip dependency** — the `tensorrt` Python package is no longer a + dependency of TensorRT LLM and is not installed. +- **Checkpoint conversion scripts** — the per-model `convert_checkpoint.py` + scripts and the associated TensorRT example directories have been removed. + +## Breaking-Change Contract + +| Removed surface | Previous behavior | New behavior | +| --- | --- | --- | +| `LLM(backend="tensorrt")` | Built and ran a TensorRT engine | Raises `ValueError` | +| `TrtLlmArgs` | TensorRT argument schema | Removed; use `TorchLlmArgs` | +| `tensorrt_llm._tensorrt_engine.LLM` | TensorRT engine entry point | Removed | +| `trtllm-build` / `trtllm-refit` / `trtllm-prune` | Engine build / refit / prune CLIs | Removed | +| `--backend tensorrt` | Selected the TensorRT backend | Removed choice; use `pytorch` | +| `tensorrt` pip dependency | Installed with TensorRT LLM | Dropped | +| `convert_checkpoint.py` | Converted HF weights to TRT checkpoints | Removed | + +## Examples and Documentation + +The per-model TensorRT example directories (the `convert_checkpoint.py` / +`trtllm-build` workflow under `examples/models/`) have been removed. Examples +and docs that cover the **PyTorch backend** are retained: + +- PyTorch model usage now lives with the LLM Python API examples under + [`examples/llm-api/`](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/llm-api) + (for example `quickstart_advanced.py` and `quickstart_multimodal.py`). +- Per-model READMEs that documented PyTorch-backend usage (for example + Qwen3/Qwen3-Next, Gemma, Phi-4-multimodal, Nemotron-NAS) have been + trimmed to keep only their PyTorch sections. Models that already have a + dedicated deployment walkthrough (for example Llama-3.3-70B) point to the + [deployment guides](../deployment-guide/) instead. +- Model-specific deployment walkthroughs remain under + [deployment guides](../deployment-guide/). +- The legacy TensorRT documentation tree under `docs/source/legacy/` has been + removed. Content that still applies to the PyTorch backend was relocated — + for example the + [Low-Precision AllReduce](../features/lowprecision-allreduce.md) guide now + lives under Features. + +## How to Migrate + +### Python API + +Before: + +```python +from tensorrt_llm import LLM + +llm = LLM(model="", backend="tensorrt") +``` + +After (PyTorch is the default backend — no `backend` argument needed): + +```python +from tensorrt_llm import LLM + +llm = LLM(model="") +``` + +The PyTorch backend loads HuggingFace checkpoints directly. There is no separate +checkpoint-conversion or engine-build step, so `convert_checkpoint.py` and +`trtllm-build` are no longer part of the workflow. + +### Command Line + +Before: + +```bash +trtllm-build --checkpoint_dir --output_dir +trtllm-serve --backend tensorrt +``` + +After: + +```bash +trtllm-serve +``` + +PyTorch is the default backend, so `--backend` may be omitted. To be explicit, +pass `--backend pytorch`. + +## Where to Go Next + +- [Quick Start Guide](../quick-start-guide.md) +- [LLM API Reference](../llm-api/index.md) +- [Supported Models](../models/supported-models.md) +- [PyTorch Backend Architecture](../torch/arch_overview.md) diff --git a/docs/source/legacy/tensorrt_quickstart.md b/docs/source/legacy/tensorrt_quickstart.md deleted file mode 100644 index e74a0f5e9e28..000000000000 --- a/docs/source/legacy/tensorrt_quickstart.md +++ /dev/null @@ -1,9 +0,0 @@ -# LLM API with TensorRT Engine -A simple inference example with TinyLlama using the LLM API: - -```{literalinclude} ../../../examples/llm-api/_tensorrt_engine/quickstart_example.py - :language: python - :linenos: -``` - -For more advanced usage including distributed inference, multimodal, and speculative decoding, please refer to this [README](../../../examples/llm-api/README.md). diff --git a/docs/source/legacy/torch.md b/docs/source/legacy/torch.md deleted file mode 100644 index e3fcf1dd8835..000000000000 --- a/docs/source/legacy/torch.md +++ /dev/null @@ -1,44 +0,0 @@ -# PyTorch Backend - -```{note} -Note: -This feature is currently in beta, and the related API is subject to change in future versions. -``` -To enhance the usability of the system and improve developer efficiency, TensorRT LLM launches a new backend based on PyTorch. - -The PyTorch backend of TensorRT LLM is available in version 0.17 and later. You can try it via importing `tensorrt_llm._torch`. - -## Quick Start - -Here is a simple example to show how to use `tensorrt_llm.LLM` API with Llama model. - -```{literalinclude} ../../examples/llm-api/quickstart_example.py - :language: python - :linenos: -``` - -## Features - -- [Sampling](../torch/features/sampling.md) -- [Quantization](../torch/features/quantization.md) -- [Overlap Scheduler](../torch/features/overlap_scheduler.md) -- [Feature Combination Matrix](../features/feature-combination-matrix.md) - -## Developer Guide - -- [Architecture Overview](../torch/arch_overview.md) -- [Adding a New Model](../torch/adding_new_model.md) - -## Key Components - -- [Attention](../torch/attention.md) -- [KV Cache Manager](../torch/kv_cache_manager.md) -- [Scheduler](../torch/scheduler.md) - -## Known Issues - -- The PyTorch backend on SBSA is incompatible with bare metal environments like Ubuntu 24.04. Please use the [PyTorch NGC Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch) for optimal support on SBSA platforms. - -## Prototype Features - -- [AutoDeploy: Seamless Model Deployment from PyTorch to TensorRT LLM](../features/auto_deploy/auto-deploy.md) diff --git a/docs/source/release-notes.md b/docs/source/release-notes.md index 9f0ed17d4f76..7b38dc720118 100644 --- a/docs/source/release-notes.md +++ b/docs/source/release-notes.md @@ -26,6 +26,7 @@ All published functionality in the Release Notes has been fully tested and verif ### API Changes +- **[BREAKING CHANGE] TensorRT backend removed.** PyTorch is now the sole execution backend. `LLM(backend="tensorrt")` now raises a `ValueError`; `TrtLlmArgs`, `tensorrt_llm._tensorrt_engine.LLM`, the `trtllm-build` / `trtllm-refit` / `trtllm-prune` CLIs, the `--backend tensorrt` CLI choice, and the per-model `convert_checkpoint.py` scripts have all been removed. The `tensorrt` pip dependency is no longer installed. See the [TensorRT Backend Removed migration guide](legacy/tensorrt-backend-removal.md) for details. - `trtllm-serve`, `trtllm-eval`, `trtllm-bench`: explicit CLI flags now take precedence over values in `--config` / `--extra_llm_api_options` YAML files (was: YAML overrode CLI). Un-set CLI flags continue to fall back to the YAML, then to model-specific and built-in defaults. ### Fixed Issues diff --git a/docs/source/torch/arch_overview.md b/docs/source/torch/arch_overview.md index d67a4668e259..5856dc92ee58 100644 --- a/docs/source/torch/arch_overview.md +++ b/docs/source/torch/arch_overview.md @@ -17,7 +17,7 @@ The `LLM` also manages the tokenization and detokenization processes of the inpu ## PyExecutor -Similar to the TensorRT backend, which uses [Executor API](../legacy/advanced/executor.md), the PyTorch backend employs a `PyExecutor` class. +The PyTorch backend employs a `PyExecutor` class. This class has a similar interface to Executor, allowing it to be integrated into LLM as an alternative backend. Key components of the `PyExecutor` include: diff --git a/docs/source/torch/kv_cache_manager.md b/docs/source/torch/kv_cache_manager.md index 57d02c8ab59f..25196a1b3164 100644 --- a/docs/source/torch/kv_cache_manager.md +++ b/docs/source/torch/kv_cache_manager.md @@ -4,7 +4,7 @@ In Transformer-based models, the KV (Key-Value) Cache is a mechanism used to opt Since KV Cache requires memory to store, it is also an important resource. In TensorRT LLM, KV Cache is managed by the `KVCacheManager`. -For details of the TensorRT LLM `KVCacheManager` implementation see [KV Cache Management](../legacy/advanced/kv-cache-management.md). +For details on KV cache behavior in TensorRT LLM, see [KV Cache Management](../features/kvcache.md). ## KV Cache Manager Introduction diff --git a/examples/draft_target_model/README.md b/examples/draft_target_model/README.md deleted file mode 100644 index 1703b3def7ec..000000000000 --- a/examples/draft_target_model/README.md +++ /dev/null @@ -1,364 +0,0 @@ -# Draft-Target-Model Speculative Decoding (DTM) - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a model using DTM speculative decoding (also known as `Speculative-Sampling`, [`Paper`](https://arxiv.org/abs/2302.01318)) in TensorRT LLM on single GPU, or single node multiple GPU. - -## Overview - -We provide two styles of running DTM now: using TensorRT-LLM-BLS in Triton Inference Server, or using TensorRT LLM directly. Here we introduce the detailed steps of running DTM in both workflows. - -## Support Matrix - * GPU Compute Capability >= 8.0 (Ampere or newer) - * FP16 / BF16 / FP8 (both draft and target model) - * Paged KV Cache - * Tensor Parallel - -## Usage - -### Build draft and target engines (the same in two workflows) - -+ We use open-source `llama-7B/13B` as draft and target models in this example, assuming the paths to the models' repository are `DRAFT_MODEL_PATH` and `TARGET_MODEL_PATH`. -+ `--use_paged_context_fmha=enable` must be specified since we need KV-Cache reuse in this approach. -+ `--speculative_decoding_mode=draft_tokens_external` and `--max_draft_len` must be specified for target model. -+ `--gather_generation_logits` is necessary if using generation logits for selecting tokens in target model. -+ `--tp_size` can be modified set if using TP mode for draft / target model. - -```bash -cd examples/models/core/llama -export DRAFT_CKPT_PATH=/workspace/ckpt-draft -export TARGET_CKPT_PATH=/workspace/ckpt-target -export DRAFT_ENGINE_PATH=/workspace/engine-draft -export TARGET_ENGINE_PATH=/workspace/engine-target -export MAX_BATCH_SIZE=4 -export MAX_DRAFT_LEN=10 -export MAX_INPUT_LEN=3200 -export MAX_SEQ_LEN=4800 - -python3 convert_checkpoint.py \ - --model_dir=${DRAFT_MODEL_PATH} \ - --output_dir=${DRAFT_CKPT_PATH} \ - --dtype=float16 - -python3 convert_checkpoint.py \ - --model_dir=${TARGET_MODEL_PATH} \ - --output_dir=${TARGET_CKPT_PATH} \ - --dtype=float16 - -trtllm-build \ - --checkpoint_dir=${DRAFT_CKPT_PATH} \ - --output_dir=${DRAFT_ENGINE_PATH} \ - --gemm_plugin=float16 \ - --use_paged_context_fmha=enable \ - --max_batch_size=${MAX_BATCH_SIZE} \ - --max_input_len=${MAX_INPUT_LEN} \ - --max_seq_len=${MAX_SEQ_LEN} - -trtllm-build \ - --checkpoint_dir=${TARGET_CKPT_PATH} \ - --output_dir=${TARGET_ENGINE_PATH} \ - --gemm_plugin=float16 \ - --use_paged_context_fmha=enable \ - --speculative_decoding_mode=draft_tokens_external \ - --max_batch_size=${MAX_BATCH_SIZE} \ - --max_draft_len=${MAX_DRAFT_LEN} \ - --max_input_len=${MAX_INPUT_LEN} \ - --max_seq_len=${MAX_SEQ_LEN} -``` - -### TensorRT LLM workflow - -+ `--draft_engine_dir` and `--engine_dir` must be specified for the draft and target engines respectively. -+ `--draft_target_model_config` is corresponding configuration of DTM, which has 4 hyperparameters that you need to specify to control the process of generation: - - `draft_len`: the number of tokens the draft model generated in one iteration, which the range is from 4 to 10 in common usage. Empirically, the larger the value is, the higher acceptance ratio but higher overhead is expected at the same time, so the right balance based on the models and application scenarios needs to be found. - - `draft_model_device_list`: the index list of device(s) to run the draft model. The length of it must be the same as the TP size of the draft model engine. For instances, `draft_model_device_list=[1]` means using tp_size=1 and GPU 1 for draft model, `draft_model_device_list=[4,5,6,7]` means using tp=4 and GPU from 4 to 7 for draft model. - - `target_model_device_list`: the index list of device(s) to run the target model. The length of it must be the same as the TP size of the target model engine. For instances, `draft_model_device_list=[0]` means using tp_size=1 and GPU 0 for target model, `draft_model_device_list=[2,3]` means using tp=2 and GPU from 2 to 3 for target model. - - `use_logits`: there are two methods to accept tokens proposed by draft model. When `use_logits=True`, the draft tokens are accepted based on the ratio of the logits from draft and target model (modified rejection sampling method in the original paper); When `use_logits=False`, the draft tokens are accepted based on per-token comparison with target predictions regardless of the logits. - - As an example, `[4,[0],[1],False]` means `draft_len=4`, device of draft model is `GPU0`, device of target model is `GPU1`, and use tokens rather than logits to accept. -+ `--kv_cache_enable_block_reuse` must be specified for this approach. -+ Only CPP session is supported, so `--use_py_session` must not be specified. -+ `--kv_cache_free_gpu_memory_fraction` should be specified if we want to place two models on one GPU, or one of the models would use out of the GPU memory. -+ `--num_beams` can not be specified as larger than 1 since beam search is not supported in this approach yet. -+ `--output_generation_logits` is optional. In original paper, we accept the tokens by comparing logits of draft and target models, so this parameter is needed. But for simplification, we can accept the tokens by comparing the output token directly, in this occasion, we can skip this parameter. - -```bash -python3 examples/run.py \ - --tokenizer_dir=${TARGET_MODEL_PATH} \ - --draft_engine_dir=/workspace/engine-draft \ - --engine_dir=/workspace/engine-target \ - --draft_target_model_config="[4,[0],[1],False]" \ - --max_output_len=256 \ - --kv_cache_enable_block_reuse \ - --kv_cache_free_gpu_memory_fraction=0.5 \ - --input_text="How does Draft-Sampling work?" -``` - -### Triton Inference Server workflow - -+ This example is based on TensorRT-LLM-0.18.0 and TRTLLM-backend-0.18.0 with docker image `nvcr.io/nvidia/tritonserver:25.03-trtllm-python-py3`. -+ DTM model approach is supported since TensorRT-LLM-0.7.0 (using two separate Tritonserver to maintain draft and target model respectively), but has significant optimization in TensorRT-LLM-0.10.0 (using one Tritonserver with [Business Logic Scripting](https://github.com/triton-inference-server/python_backend?tab=readme-ov-file#business-logic-scripting), BLS). - -#### Get related repository inside the container - -```bash -git clone https://github.com/triton-inference-server/tensorrtllm_backend.git -cd tensorrtllm_backend -git checkout rel -git lfs pull -git submodule update --init --recursive -pip install -r requirements.txt -pip install SentencePiece tritonclient - -export DRAFT_MODEL_NAME="tensorrt_llm_draft" -export TARGET_MODEL_NAME="tensorrt_llm" -export TRITON_MODEL_REPO=llama_dtm -``` - -#### Simple deploy - -+ Edit model configuration. - -```bash -export DRAFT_DEVICE_IDS="0" -export TARGET_DEVICE_IDS="1" - -rm -rf ${TRITON_MODEL_REPO} -cp -r all_models/inflight_batcher_llm/ ${TRITON_MODEL_REPO} -cp -r ${TRITON_MODEL_REPO}/tensorrt_llm ${TRITON_MODEL_REPO}/tensorrt_llm_draft -sed -i 's/name: "tensorrt_llm"/name: "tensorrt_llm_draft"/g' ${TRITON_MODEL_REPO}/tensorrt_llm_draft/config.pbtxt - -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/ensemble/config.pbtxt triton_max_batch_size:4,logits_datatype:TYPE_FP32 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/preprocessing/config.pbtxt triton_max_batch_size:4,tokenizer_dir:${HF_MODEL},preprocessing_instance_count:1 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/postprocessing/config.pbtxt triton_max_batch_size:4,tokenizer_dir:${HF_MODEL},postprocessing_instance_count:1 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/tensorrt_llm_bls/config.pbtxt triton_max_batch_size:4,decoupled_mode:False,logits_datatype:TYPE_FP32,bls_instance_count:1,accumulate_tokens:False,tensorrt_llm_model_name:${TARGET_MODEL_NAME},tensorrt_llm_draft_model_name:${DRAFT_MODEL_NAME} - -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/tensorrt_llm/config.pbtxt triton_max_batch_size:4,decoupled_mode:False,logits_datatype:TYPE_FP32,triton_backend:tensorrtllm,max_tokens_in_paged_kv_cache:2560,max_attention_window_size:2560,kv_cache_free_gpu_mem_fraction:0.5,exclude_input_in_output:False,enable_kv_cache_reuse:True,batching_strategy:inflight_fused_batching,max_queue_delay_microseconds:0,encoder_input_features_data_type:TYPE_FP16,engine_dir:${TARGET_ENGINE_PATH},gpu_device_ids:${TARGET_DEVICE_IDS} -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/tensorrt_llm_draft/config.pbtxt triton_max_batch_size:4,decoupled_mode:False,logits_datatype:TYPE_FP32,triton_backend:tensorrtllm,max_tokens_in_paged_kv_cache:2560,max_attention_window_size:2560,kv_cache_free_gpu_mem_fraction:0.5,exclude_input_in_output:False,enable_kv_cache_reuse:True,batching_strategy:inflight_fused_batching,max_queue_delay_microseconds:0,encoder_input_features_data_type:TYPE_FP16,engine_dir:${DRAFT_ENGINE_PATH},gpu_device_ids:${DRAFT_DEVICE_IDS} -``` - -+ Start the triton inference server. - + Verbose log will be written in to file `triton_log.txt` if specifying `--log`. - -```bash -python3 scripts/launch_triton_server.py \ - --model_repo=${TRITON_MODEL_REPO} \ - --multi-model \ - --log -``` - -+ You can see the output below in the file if Triton server launches successfully: - -```txt -Started HTTPService at 0.0.0.0:8000 -Started GRPCInferenceService at 0.0.0.0:8001 -Started Metrics Service at 0.0.0.0:8002 -``` - -+ Send a request for inference. - -```bash -python3 inflight_batcher_llm/client/e2e_grpc_speculative_decoding_client.py \ - --url-target=localhost:8001 \ - --draft-tensorrt-llm-model-name=${DRAFT_MODEL_NAME} \ - --target-tensorrt-llm-model-name=${TARGET_MODEL_NAME} \ - --output-len=100 \ - --num-draft-tokens=4 \ - --end-id=2 \ - --pad-id=2 \ - --prompt "What is Ubuntu operation system?" -``` - -+ You can receive the following results if everything goes smoothly. - -```txt -Final text: - What is Ubuntu operation system? -Ubuntu is a free and open source operating system that runs from the desktop, to the cloud, to all your internet connected things. Ubuntu is used by millions of people around the world who want to explore new ideas and discover new opportunities. -Ubuntu is a community developed operating system that is perfect for laptops, desktops, servers, and cloud. It is used by millions of people around the world who want to explore new ideas and discover new opportunities. -``` - -+ Test DTM with a script. - + Prepare a JSON file `input_data.json` containing input data as below (more requests are acceptable). - -```json -[ - { - "input": "What is Ubuntu operation system?", - "instruction": "Answer the question shortly.", - "output": " " - } -] -``` - -+ Use command below to launch test. - -```bash -### Use BLS speculative decoding -python3 tools/inflight_batcher_llm/speculative_decoding_test.py \ - --max-input-len 2500 \ - --dataset input_data.json \ - --url-control=localhost:8001 \ - --url-target=localhost:8001 \ - --url-draft=localhost:8001 \ - --draft-tensorrt-llm-model-name="${DRAFT_MODEL_NAME}" \ - --target-tensorrt-llm-model-name="${TARGET_MODEL_NAME}" \ - --bls-speculative-tensorrt-llm-model-name="tensorrt_llm_bls" \ - --execute-bls-speculative-decoding \ - --disable-output-comparison \ - --num-draft-tokens=4 \ - --use-draft-logits - -### Use client-side speculative decoding -python3 tools/inflight_batcher_llm/speculative_decoding_test.py \ - --max-input-len 2500 \ - --dataset input_data.json \ - --url-control=localhost:8001 \ - --url-target=localhost:8001 \ - --url-draft=localhost:8001 \ - --draft-tensorrt-llm-model-name="${DRAFT_MODEL_NAME}" \ - --target-tensorrt-llm-model-name="${TARGET_MODEL_NAME}" \ - --bls-speculative-tensorrt-llm-model-name="tensorrt_llm_bls" \ - --disable-output-comparison \ - --num-draft-tokens=4 \ - --use-draft-logits -``` - -+ You can receive the following results if everything goes smoothly. - -```txt -Ubuntu is a free and open source operating system. It is a Linux based operating system. ... -``` - -+ Stop triton inference server after all work is done. - -```bash -pkill tritonserver -``` - -+ In addition, it appears better performance can be achieved with both draft and target engines deployed on a single GPU (llama-7B-FP8 + llama-30B-FP8, for a total of 40GiB on one H100-80GiB GPU for example). - -#### Usage of Tensor-Parallelization mode. - -+ In this example, we use draft engine with TP=1 and target engine with TP=2 (both symmetrical or asymmetrical TP size are acceptable), and want to place the draft engine on GPU0, target engine on GPU1 and GPU2. -+ Edit model configuration. - -```bash -export DRAFT_DEVICE_IDS="0" -export TARGET_DEVICE_IDS="1,2" - -rm -rf ${TRITON_MODEL_REPO} -cp -r all_models/inflight_batcher_llm/ ${TRITON_MODEL_REPO} -cp -r ${TRITON_MODEL_REPO}/tensorrt_llm ${TRITON_MODEL_REPO}/tensorrt_llm_draft -sed -i 's/name: "tensorrt_llm"/name: "tensorrt_llm_draft"/g' ${TRITON_MODEL_REPO}/tensorrt_llm_draft/config.pbtxt - -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/ensemble/config.pbtxt triton_max_batch_size:4,logits_datatype:TYPE_FP32 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/preprocessing/config.pbtxt triton_max_batch_size:4,tokenizer_dir:${HF_MODEL},preprocessing_instance_count:1 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/postprocessing/config.pbtxt triton_max_batch_size:4,tokenizer_dir:${HF_MODEL},postprocessing_instance_count:1 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/tensorrt_llm_bls/config.pbtxt triton_max_batch_size:4,decoupled_mode:False,logits_datatype:TYPE_FP32,bls_instance_count:1,accumulate_tokens:False,tensorrt_llm_model_name:${TARGET_MODEL_NAME},tensorrt_llm_draft_model_name:${DRAFT_MODEL_NAME} - -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/tensorrt_llm/config.pbtxt triton_max_batch_size:4,decoupled_mode:False,logits_datatype:TYPE_FP32,triton_backend:tensorrtllm,max_tokens_in_paged_kv_cache:2560,max_attention_window_size:2560,kv_cache_free_gpu_mem_fraction:0.5,exclude_input_in_output:False,enable_kv_cache_reuse:True,batching_strategy:inflight_fused_batching,max_queue_delay_microseconds:0,encoder_input_features_data_type:TYPE_FP16,engine_dir:${TARGET_ENGINE_PATH} -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/tensorrt_llm_draft/config.pbtxt triton_max_batch_size:4,decoupled_mode:False,logits_datatype:TYPE_FP32,triton_backend:tensorrtllm,max_tokens_in_paged_kv_cache:2560,max_attention_window_size:2560,kv_cache_free_gpu_mem_fraction:0.5,exclude_input_in_output:False,enable_kv_cache_reuse:True,batching_strategy:inflight_fused_batching,max_queue_delay_microseconds:0,encoder_input_features_data_type:TYPE_FP16,engine_dir:${DRAFT_ENGINE_PATH} - -sed -i 's/\${gpu_device_ids}/'"${DRAFT_DEVICE_IDS}"'/g' ${TRITON_MODEL_REPO}/tensorrt_llm_draft/config.pbtxt -sed -i 's/\${gpu_device_ids}/'"${TARGET_DEVICE_IDS}"'/g' ${TRITON_MODEL_REPO}/tensorrt_llm/config.pbtxt -``` - -+ As you see, the only difference is `gpu_device_ids`, which needs fix manually since comma is not supported in script `python3 tools/fill_template.py`. - -+ Start the triton inference server - + Use `--multi-model` to enable orchestrator mode in TP>1 scenario. See [model config](https://github.com/triton-inference-server/tensorrtllm_backend/blob/main/docs/model_config.md) for more information. - -```bash -python3 scripts/launch_triton_server.py \ - --model_repo=${TRITON_MODEL_REPO} \ - --tensorrt_llm_model_name "tensorrt_llm,tensorrt_llm_draft" \ - --multi-model -``` - -+ All other operations are the same as `Simple deploy` part. - -#### Usage of Fast logits D2D transfer - -+ Fast logits boosts the performance (TPS) by hiding the latency of logits transfer from draft engine to target engine supported since TensorRT-LLM-0.15.0. -+ In this example, we use draft engine with TP=1 and target engine with TP=2 (both symmetrical or asymmetrical TP size are acceptable), and want to place the draft engine on GPU0, target engine on GPU1 and GPU2. -+ For `participant_ids`, rank 0 is reserved for the orchestrator; rank (`1` ~ `tp_size_draft`) are for draft engine; rank (`tp_size_draft+1` ~ `tp_size_draft+tp_size_target`) are for target engine. -+ Edit model configuration. - -```bash -cd examples/models/core/llama -export TARGET_CKPT_PATH=/workspace/ckpt-target -export TARGET_ENGINE_PATH=/workspace/engine-target -export MAX_BATCH_SIZE=4 -export MAX_DRAFT_LEN=10 -export MAX_INPUT_LEN=3200 -export MAX_SEQ_LEN=4800 - -python3 convert_checkpoint.py \ - --model_dir=${TARGET_MODEL_PATH} \ - --output_dir=${TARGET_CKPT_PATH} \ - --dtype=float16 \ - --tp_size=2 - -trtllm-build \ - --checkpoint_dir=${TARGET_CKPT_PATH} \ - --output_dir=${TARGET_ENGINE_PATH} \ - --gemm_plugin=float16 \ - --use_paged_context_fmha=enable \ - --speculative_decoding_mode=draft_tokens_external \ - --max_batch_size=${MAX_BATCH_SIZE} \ - --max_draft_len=${MAX_DRAFT_LEN} \ - --max_input_len=${MAX_INPUT_LEN} \ - --max_seq_len=${MAX_SEQ_LEN} -``` - -```bash -export DRAFT_DEVICE_IDS="0" -export TARGET_DEVICE_IDS="1,2" -export DRAFT_PARTICIPANT_IDS="1" -export TARGET_PARTICIPANT_IDS="2,3" - -cd /work/tensorrtllm_backend -rm -rf ${TRITON_MODEL_REPO} -cp -r all_models/inflight_batcher_llm/ ${TRITON_MODEL_REPO} -cp -r ${TRITON_MODEL_REPO}/tensorrt_llm ${TRITON_MODEL_REPO}/tensorrt_llm_draft -sed -i 's/name: "tensorrt_llm"/name: "tensorrt_llm_draft"/g' ${TRITON_MODEL_REPO}/tensorrt_llm_draft/config.pbtxt - -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/ensemble/config.pbtxt triton_max_batch_size:4,logits_datatype:TYPE_FP32 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/preprocessing/config.pbtxt triton_max_batch_size:4,tokenizer_dir:${HF_MODEL},preprocessing_instance_count:1 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/postprocessing/config.pbtxt triton_max_batch_size:4,tokenizer_dir:${HF_MODEL},postprocessing_instance_count:1 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/tensorrt_llm_bls/config.pbtxt triton_max_batch_size:4,decoupled_mode:False,bls_instance_count:1,accumulate_tokens:False,tensorrt_llm_model_name:${TARGET_MODEL_NAME},logits_datatype:TYPE_FP32,tensorrt_llm_draft_model_name:${DRAFT_MODEL_NAME} - -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/tensorrt_llm/config.pbtxt triton_max_batch_size:4,triton_backend:tensorrtllm,decoupled_mode:False,max_beam_width:1,engine_dir:${TARGET_ENGINE_PATH},max_tokens_in_paged_kv_cache:2560,max_attention_window_size:2560,kv_cache_free_gpu_mem_fraction:0.5,exclude_input_in_output:False,enable_kv_cache_reuse:True,batching_strategy:inflight_fused_batching,max_queue_delay_microseconds:0,encoder_input_features_data_type:TYPE_FP16,logits_datatype:TYPE_FP32,gpu_device_ids:${TARGET_DEVICE_IDS},participant_ids:2,3,speculative_decoding_fast_logits:1 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/tensorrt_llm_draft/config.pbtxt triton_max_batch_size:4,triton_backend:tensorrtllm,decoupled_mode:False,max_beam_width:1,engine_dir:${DRAFT_ENGINE_PATH},max_tokens_in_paged_kv_cache:2560,max_attention_window_size:2560,kv_cache_free_gpu_mem_fraction:0.5,exclude_input_in_output:False,enable_kv_cache_reuse:True,batching_strategy:inflight_fused_batching,max_queue_delay_microseconds:0,encoder_input_features_data_type:TYPE_FP16,logits_datatype:TYPE_FP32,gpu_device_ids:${DRAFT_DEVICE_IDS},participant_ids:1,speculative_decoding_fast_logits:1 - -sed -i 's/\${gpu_device_ids}/'"${DRAFT_DEVICE_IDS}"'/g' ${TRITON_MODEL_REPO}/tensorrt_llm_draft/config.pbtxt -sed -i 's/\${participant_ids}/'"${DRAFT_PARTICIPANT_IDS}"'/g' ${TRITON_MODEL_REPO}/tensorrt_llm_draft/config.pbtxt -sed -i 's/\${gpu_device_ids}/'"${TARGET_DEVICE_IDS}"'/g' ${TRITON_MODEL_REPO}/tensorrt_llm/config.pbtxt -sed -i 's/\${participant_ids}/'"${TARGET_PARTICIPANT_IDS}"'/g' ${TRITON_MODEL_REPO}/tensorrt_llm/config.pbtxt -``` - -+ As you see, the differences are `participant_ids` and `speculative_decoding_fast_logits`. - -+ Start the triton inference server. - + Use `--disable-spawn-process` to enable pre-spawn variant in orchestrator mode. - + `--world_size` must be equal to `1 + tp_size_draft + tp_size_target`, which is 4 in this example. - -```bash -python3 scripts/launch_triton_server.py \ - --model_repo ${TRITON_MODEL_REPO} \ - --tensorrt_llm_model_name tensorrt_llm,tensorrt_llm_draft \ - --multi-model \ - --world_size 4 \ - --disable-spawn-processes -``` - -+ All other operations are the same as the `Simple deploy` part. - -### Additional information - -+ With the fast logits enabled and following optimization tips in [model configuration](https://github.com/triton-inference-server/tensorrtllm_backend/blob/main/docs/model_config.md#some-tips-for-model-configuration), speculative decoding with draft logits achieves 2.x throughput in BS1, 1.x throughput in BS16 comparing to auto-regressive decoding using Llama 3.2 1B draft and Llama 3.1 70B target. -+ Streaming mode or batched-request mode are not supported in DTM yet. diff --git a/examples/draft_target_model/requirements.txt b/examples/draft_target_model/requirements.txt deleted file mode 100644 index 423ba2d4b0e0..000000000000 --- a/examples/draft_target_model/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ --c ../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -rouge_score -sentencepiece>=0.1.99 -evaluate diff --git a/examples/eagle/README.md b/examples/eagle/README.md deleted file mode 100644 index ac95e21516cd..000000000000 --- a/examples/eagle/README.md +++ /dev/null @@ -1,195 +0,0 @@ -# EAGLE speculative Decoding - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a model using EAGLE decoding ([`GitHub`](https://github.com/SafeAILab/EAGLE/tree/main), [`BLOG`](https://sites.google.com/view/eagle-llm)) in TensorRT LLM on a single node with one or multiple GPUs. - -## Overview -Different from other models, EAGLE decoding needs a base model and an EAGLE model. - -The TensorRT LLM EAGLE decoding implementation can be found in [tensorrt_llm/models/eagle/model.py](../../tensorrt_llm/models/eagle/model.py). -The implementation adds an EAGLE drafter network to a base model. - -For more info about EAGLE, refer to [speculative decoding documentation](https://nvidia.github.io/TensorRT-LLM/features/speculative-decoding.html). - -## Limitations - * EAGLE-2 is not supported. - * All EAGLE choices have to have exactly the same depth as `num_eagle_layers` of the engine. - * Pipeline parallelism is not supported. - -## Support Matrix - * GPU Compute Capability >= 8.0 (Ampere or newer) - * FP16/BF16 - * Paged KV cache - * Inflight-fused-batching - * C++ runtime - * Tensor Parallel - -This example is based on the Vicuna-7b v1.3 model, a fine-tuned Llama. With some modifications, you can add EAGLE to other base models as well. Some TensorRT LLM models might not work with EAGLE due to the missing head size in the speculative decoding XQA attention kernels. - -## Usage -The TensorRT LLM EAGLE example code is located in [`examples/eagle`](./). There is one [`convert_checkpoint.py`](./convert_checkpoint.py) file to convert and build the [TensorRT](https://developer.nvidia.com/tensorrt) engine(s) needed to run models with EAGLE decoding support. -In this example, we use the model from HuggingFace [`yuhuili/EAGLE-Vicuna-7B-v1.3`](https://huggingface.co/yuhuili/EAGLE-Vicuna-7B-v1.3), which is a LLAMA-based model. - -### Build TensorRT engine(s) -Get the weights by downloading the base model [`vicuna-7b-v1.3`](https://huggingface.co/lmsys/vicuna-7b-v1.3) and the EAGLE draft model [`EAGLE-Vicuna-7B-v1.3`](https://huggingface.co/yuhuili/EAGLE-Vicuna-7B-v1.3) from HF. - -``` -pip install -r requirements.txt - -git lfs install -git clone https://huggingface.co/lmsys/vicuna-7b-v1.3 -https://huggingface.co/yuhuili/EAGLE-Vicuna-7B-v1.3 -``` - -Here is the example: -```bash -# Convert and Build EAGLE decoding support for vicuna-7b-v1.3 -python convert_checkpoint.py --model_dir ./vicuna-7b-v1.3 \ - --eagle_model_dir EAGLE-Vicuna-7B-v1.3 \ - --output_dir ./tllm_checkpoint_1gpu_eagle \ - --dtype float16 \ - --max_draft_len 63 \ - --num_eagle_layers 4 \ - --max_non_leaves_per_layer 10 - -# Note: Increasing the batch size may have a negative impact on performance -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_eagle \ - --output_dir ./tmp/eagle/7B/trt_engines/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --use_paged_context_fmha enable \ - --speculative_decoding_mode eagle \ - --max_batch_size 4 - -# Convert and Build EAGLE decoding support for vicuna-7b-v1.3 with 4-way tensor parallelism. -python convert_checkpoint.py --model_dir ./vicuna-7b-v1.3 \ - --eagle_model_dir EAGLE-Vicuna-7B-v1.3 \ - --output_dir ./tllm_checkpoint_4gpu_eagle \ - --dtype float16 \ - --max_draft_len 63 \ - --num_eagle_layers 4 \ - --max_non_leaves_per_layer 10 \ - --tp_size 4 \ - --workers 4 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_4gpu_eagle \ - --output_dir ./tmp/eagle/7B/trt_engines/fp16/4-gpu/ \ - --gemm_plugin float16 \ - --use_paged_context_fmha enable \ - --speculative_decoding_mode eagle \ - --max_batch_size 4 -``` - -### Run - -To run a TensorRT LLM model with EAGLE-1 decoding support, you can use `../run.py` script, with an additional argument -`--eagle_choices`. -The `--eagle_choices` argument is of type `list[list[int]]`. If you do not specify any choices, the -default, [mc_sim_7b_63](https://github.com/FasterDecoding/Medusa/blob/main/medusa/model/medusa_choices.py#L1) choices -are used. For more information regarding choices tree, refer -to [Medusa Tree](https://nvidia.github.io/TensorRT-LLM/legacy/advanced/speculative-decoding.html#medusa-tree). - -The number of non-leaf nodes at each level can not exceed `max_non_leaves_per_layer` set -to `convert_checkpoint`. For example, in the tree below (`mc_sim_7b_63`) the minimum number of -`max_non_leaves_per_layer` is 10. There are exactly 10 non leaf nodes at depth 0, check `[0, 0], [1, 0], ..., [9, 0]`. - -The maximum depth, meaning the maximum length of inner `list[int]` specified in the `--eagle_choices` argument, should -be equal to `num_eagle_layers`. - -To run non-greedy sampling and use typical acceptance, set `--eagle_posterior_threshold` to `run.py`. `eagle_posterior_threshold` corresponds to epsilon in typical acceptance criteria from [Medusa paper](https://arxiv.org/pdf/2401.10774). -`--temperature` can be specified as well. When no `--eagle_posterior_threshold` is specified or `--temperature=0.0` is set, greedy sampling is used. - -#### Run EAGLE-2 - -EAGLE-2 can be enabled with 2 runtime flags (`--eagle_use_dynamic_tree` and `--eagle_dynamic_tree_max_top_k=N`). The same engine can be used for EAGLE-1 and EAGLE-2. Eagle choices must not be set in case of EAGLE-2. EAGLE-2 will generate the tree corresponding to choices dynamically in the runtime. For more details, please refer to [EAGLE-2 paper](https://arxiv.org/pdf/2406.16858). - -When using EAGLE-2, please enable `--eagle_use_dynamic_tree`, which indicates whether to use a dynamic tree (default is `False`, i.e., use EAGLE-1 by default). Then set `--eagle_dynamic_tree_max_top_k=N`, which indicates how many new child nodes are expanded for the nodes in the dynamic tree. -- In EagleNet0, `N` draft tokens are generated. -- In EagleNet1, each draft token expands `N` new draft tokens. Therefore, this layer has `N * N` draft tokens. We select the top `N` as the output of this layer. -- In EagleNet2, the `N` output nodes of EagleNet1 are expanded, and each node expands `N` new draft tokens. Therefore, this layer also has a total of `N * N` draft tokens. And select the top `N` as the output of this layer. -- Etc. - -Finally, after `num_eagle_layer` EagleNets, `N + N * N * (num_eagle_layer - 1)` draft tokens are generated. We will rebuild the final tree based on all draft tokens and their scores. The final generated tree will have `min(N + N * N * (num_eagle_layer - 1), max_draft_len)` nodes. - - - -```bash -# Eagle greedy decoding using vicuna-7b-v1.3 model with 1 GPU -python ../run.py --engine_dir ./tmp/eagle/7B/trt_engines/fp16/1-gpu/ \ - --tokenizer_dir ./vicuna-7b-v1.3/ \ - --max_output_len=100 \ - --eagle_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --input_text "Once upon" - -# Eagle typical acceptance decoding using vicuna-7b-v1.3 model with 1 GPU -python ../run.py --engine_dir ./tmp/eagle/7B/trt_engines/fp16/1-gpu/ \ - --tokenizer_dir ./vicuna-7b-v1.3/ \ - --max_output_len=100 \ - --eagle_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --input_text "Once upon" \ - --temperature 0.7 \ - --eagle_posterior_threshold 0.09 - -# Eagle decoding using vicuna-7b-v1.3 model with 4 GPUs -mpirun -np 4 --allow-run-as-root --oversubscribe \ - python ../run.py --engine_dir ./tmp/eagle/7B/trt_engines/fp16/4-gpu/ \ - --tokenizer_dir ./vicuna-7b-v1.3/ \ - --max_output_len=100 \ - --eagle_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --input_text "Once upon" - -# Run EAGLE-2 -mpirun -np 1 --allow-run-as-root --oversubscribe \ - python ../run.py --engine_dir ./tmp/eagle/7B/trt_engines/fp16/1-gpu/ \ - --tokenizer_dir ./vicuna-7b-v1.3/ \ - --max_output_len=100 \ - --eagle_use_dynamic_tree \ - --eagle_dynamic_tree_max_top_k 10 \ - --input_text "Once upon" -``` - -For greedy decoding, refer to the following example output: -```text -...... -Input [Text 0]: " Once upon" -Output [Text 0 Beam 0]: "a time, there was a young girl who loved to read. She would spend hours in the library, devouring books of all genres. She had a special love for fairy tales, and would often dream of living in a magical world where she could meet princes and princesses, and have adventures with talking animals. -One day, while she was reading a book, she came across a passage that spoke to her heart. It said, "You are the author of" -``` - -### Summarization using EAGLE decoding -```bash -# EAGLE decoding using vicuna-7b-v1.3 model with 1 GPU -python ../summarize.py --engine_dir ./tmp/eagle/7B/trt_engines/fp16/1-gpu/ \ - --hf_model_dir ./vicuna-7b-v1.3/ \ - --tokenizer_dir ./vicuna-7b-v1.3/ \ - --test_trt_llm \ - --data_type fp16 \ - --eagle_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --batch_size 1 - -# EAGLE decoding using vicuna-7b-v1.3 with 4 GPUs -mpirun -np 4 --allow-run-as-root --oversubscribe \ - python ../summarize.py --engine_dir ./tmp/eagle/7B/trt_engines/fp16/4-gpu/ \ - --hf_model_dir ./vicuna-7b-v1.3/ \ - --tokenizer_dir ./vicuna-7b-v1.3/ \ - --test_trt_llm \ - --data_type fp16 \ - --eagle_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --batch_size 1 - -# Run EAGLE-2 -mpirun -np 1 --allow-run-as-root --oversubscribe \ - python ../summarize.py --engine_dir ./tmp/eagle/7B/trt_engines/fp16/1-gpu/ \ - --hf_model_dir ./vicuna-7b-v1.3/ \ - --tokenizer_dir ./vicuna-7b-v1.3/ \ - --test_trt_llm \ - --data_type fp16 \ - --eagle_use_dynamic_tree \ - --eagle_dynamic_tree_max_top_k 10 \ - --batch_size 1 - -``` diff --git a/examples/eagle/convert_checkpoint.py b/examples/eagle/convert_checkpoint.py deleted file mode 100644 index 130faee4453c..000000000000 --- a/examples/eagle/convert_checkpoint.py +++ /dev/null @@ -1,493 +0,0 @@ -import argparse -import json -import os -import time -from pathlib import Path - -from tqdm import tqdm -from transformers import LlamaConfig - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import get_hf_rope_theta -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.eagle.config import EagleConfig -from tensorrt_llm.models.eagle.model import EagleForCausalLM -from tensorrt_llm.models.model_weights_loader import ModelWeightsLoader -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument('--meta_ckpt_dir', type=str, default=None) - 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('--dtype', - type=str, - default='auto', - choices=['auto', 'float16', 'bfloat16', 'float32']) - parser.add_argument('--vocab_size', type=int, default=32000) - parser.add_argument('--n_positions', type=int, default=2048) - parser.add_argument('--n_layer', type=int, default=32) - - 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', 'int4_gptq'], - 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( - '--calib_dataset', - type=str, - default='ccdv/cnn_dailymail', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - 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( - '--per_channel', - action="store_true", - default=False, - 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', - action="store_true", - default=False, - 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( - '--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( - '--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('--load_by_shard', - action='store_true', - help='Load a pretrained model shard-by-shard.') - parser.add_argument('--hidden_act', type=str, default='silu') - - parser.add_argument('--rotary_base', type=float, default=10000.0) - parser.add_argument('--rotary_scaling', nargs=2, type=str, default=None) - - parser.add_argument('--group_size', - type=int, - default=128, - help='Group size used in GPTQ/AWQ quantization.') - - parser.add_argument("--storage-type", - "-t", - 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") - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - - parser.add_argument('--eagle_model_dir', type=str, default=None) - parser.add_argument('--max_draft_len', type=int, default=63) - parser.add_argument( - '--num_eagle_layers', - type=int, - default=4, - help= - 'Maximum depth of the EAGLE choices tree, i.e. maximum number of accepted draft tokens.' - ) - parser.add_argument( - '--max_non_leaves_per_layer', - type=int, - default=10, - help='Maximum number of non-leaf nodes in the EAGLE choice tree.') - args = parser.parse_args() - return args - - -def convert_and_save_hf(config, args): - world_size = args.tp_size * args.pp_size - tllm_config = EagleConfig.from_dict(config) - for rank in range(world_size): - tllm_config.mapping = Mapping(world_size=world_size, - rank=rank, - cp_size=1, - tp_size=args.tp_size, - pp_size=args.pp_size) - - model = EagleForCausalLM(tllm_config) - - def check_and_update(module, dict): - if hasattr(module, 'tllm_to_externel_key_dict'): - module.tllm_to_externel_key_dict.update(dict) - else: - module.tllm_to_externel_key_dict = dict - - def copy(tensors): - if isinstance(tensors, list): - if None in tensors: - return tensors - else: - return [tensor.clone() for tensor in tensors] - elif tensors is None: - return tensors - else: - return tensors.clone() - - shared_weight_prefixs = [] - tllm_weights = {} - customized_dict = {"drafter": ""} - if args.eagle_model_dir is None: - # Single checkpoint for ModelOpt - for idx, eagle_net in enumerate(model.eagle_nets): - check_and_update(eagle_net.drafter.fc, {"fc": "fc"}) - check_and_update(eagle_net.drafter.vocab_embedding, - {f"eagle_nets.{idx}": "model"}) - check_and_update(eagle_net.lm_head, {f"eagle_nets.{idx}": ""}) - shared_weight_prefixs.append(f"eagle_nets.{idx}") - customized_dict[f'eagle_nets.{idx}'] = 'eagle_module' - loader = ModelWeightsLoader(eagle_model_dir, customized_dict) - loader.update_key_mapping(model) - for tllm_key, _ in tqdm(model.named_parameters()): - if any([ - tllm_key.startswith(prefix) - for prefix in shared_weight_prefixs - ]): - tllm_weights.update(loader.load(tllm_key, preprocess=copy)) - else: - tllm_weights.update(loader.load(tllm_key)) - loader.fill(tllm_weights) - else: - # Double checkpoint for HF - for idx, eagle_net in enumerate(model.eagle_nets): - check_and_update(eagle_net.drafter.fc, {"fc": "fc"}) - check_and_update(eagle_net.drafter.vocab_embedding, - {f"eagle_nets.{idx}": ""}) - check_and_update(eagle_net.lm_head, {f"eagle_nets.{idx}": ""}) - shared_weight_prefixs.append(f"eagle_nets.{idx}") - customized_dict[f'eagle_nets.{idx}'] = '' - - # Load base model - base_loader = ModelWeightsLoader(args.model_dir) - base_loader.update_key_mapping(model) - for tllm_key, _ in tqdm(model.transformer.named_parameters()): - tllm_weights.update(base_loader.load("transformer." + tllm_key)) - tllm_weights.update(base_loader.load("lm_head.weight")) - for idx in range(args.num_eagle_layers): - tllm_weights.update( - base_loader.load(f"eagle_nets.{idx}.lm_head.weight", - preprocess=copy)) - - # Load eagle model - eagle_loader = ModelWeightsLoader(eagle_model_dir, customized_dict) - eagle_loader.update_key_mapping(model) - for tllm_key, _ in tqdm(model.eagle_nets.named_parameters()): - if not tllm_key.endswith("lm_head.weight"): - if any([ - tllm_key.startswith(prefix) - for prefix in shared_weight_prefixs - ]): - tllm_weights.update( - eagle_loader.load("eagle_nets." + tllm_key, - preprocess=copy)) - else: - tllm_weights.update( - eagle_loader.load("eagle_nets." + tllm_key)) - base_loader.fill(tllm_weights) - model.save_checkpoint(args.output_dir, save_config=(rank == 0)) - - -if __name__ == '__main__': - emit_engine_arch_deprecation("convert_checkpoint.py") - # TODO(qijun): Currently, the convert script depends on a torch op: - # torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix, - # which is included in tensorrt_llm Python package. Otherwise, the convert - # script does not need to import tensorrt_llm. Will remove it after reimplementing - # the op with PyTorch. - print(tensorrt_llm.__version__) - args = parse_arguments() - world_size = args.tp_size * args.pp_size - - assert args.pp_size == 1, "Pipeline parallelism is not supported in EAGLE yet." - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - hf_config = None - eagle_model_dir = args.model_dir if args.eagle_model_dir is None else args.eagle_model_dir - if args.model_dir is not None: - hf_config = LlamaConfig.from_pretrained(args.model_dir) - - args.model_type = hf_config.model_type - args.n_head = hf_config.num_attention_heads - args.inter_size = hf_config.intermediate_size - args.n_layer = hf_config.num_hidden_layers - args.n_embd = hf_config.hidden_size - args.n_kv_head = hf_config.num_key_value_heads - args.rms_norm_eps = hf_config.rms_norm_eps - args.vocab_size = hf_config.vocab_size - args.rotary_scaling = hf_config.rope_scaling - args.rotary_base = get_hf_rope_theta(hf_config, 10000.0) - args.n_positions = hf_config.max_position_embeddings - args.dtype = str( - hf_config.torch_dtype)[6:] if args.dtype == 'auto' else args.dtype - if 'head_dim' in hf_config: - args.head_dim = hf_config.head_dim - else: - args.head_dim = args.n_embd // args.n_head - if 'head_size' in hf_config: - args.head_size = hf_config.head_size - else: - args.head_size = args.head_dim - - if args.eagle_model_dir is None: - hf_config_eagle = hf_config.eagle - args.n_head_eagle = hf_config_eagle['num_attention_heads'] - args.inter_size_eagle = hf_config_eagle['intermediate_size'] - args.n_layer_eagle = hf_config_eagle['num_hidden_layers'] - args.n_embd_eagle = hf_config_eagle['hidden_size'] - args.n_kv_head_eagle = hf_config_eagle['num_key_value_heads'] - args.rms_norm_eps_eagle = hf_config_eagle['rms_norm_eps'] - args.n_positions_eagle = hf_config_eagle['max_position_embeddings'] - if 'head_dim' in hf_config_eagle: - args.head_dim_eagle = hf_config_eagle['head_dim'] - else: - args.head_dim_eagle = args.n_embd_eagle // args.n_head_eagle - if 'head_size' in hf_config_eagle: - args.head_size_eagle = hf_config_eagle['head_size'] - else: - args.head_size_eagle = args.head_dim_eagle - else: - hf_config_eagle = LlamaConfig.from_pretrained(args.eagle_model_dir) - args.n_head_eagle = hf_config_eagle.num_attention_heads - args.inter_size_eagle = hf_config_eagle.intermediate_size - args.n_layer_eagle = hf_config_eagle.num_hidden_layers - args.n_embd_eagle = hf_config_eagle.hidden_size - args.n_kv_head_eagle = hf_config_eagle.num_key_value_heads - args.rms_norm_eps_eagle = hf_config_eagle.rms_norm_eps - args.n_positions_eagle = hf_config_eagle.max_position_embeddings - if 'head_dim' in hf_config_eagle: - args.head_dim_eagle = hf_config_eagle.head_dim - else: - args.head_dim_eagle = args.n_embd_eagle // args.n_head_eagle - if 'head_size' in hf_config_eagle: - args.head_size_eagle = hf_config_eagle.head_size - else: - args.head_size_eagle = args.head_dim_eagle - - elif args.meta_ckpt_dir is not None: - assert False, "meta ckpt is not supported yet" - - with open(Path(args.meta_ckpt_dir, "params.json")) as fp: - meta_config: dict = json.load(fp) - args.n_embd = meta_config["dim"] - args.n_head = meta_config["n_heads"] - args.n_layer = meta_config["n_layers"] - args.n_kv_head = meta_config.get("n_kv_heads", args.n_head) - - if "hidden_dim" in meta_config: - args.inter_size = meta_config["hidden_dim"] - else: - args.multiple_of = meta_config.get("multiple_of", 1) - n_embd = int(4 * args.n_embd * 2 / 3) - args.ffn_dim_multiplier = meta_config.get("ffn_dim_multiplier", 1) - args.inter_size = args.multiple_of * ( - (int(n_embd * args.ffn_dim_multiplier) + args.multiple_of - 1) - // args.multiple_of) - args.rms_norm_eps = meta_config["norm_eps"] - - 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["rope_type"], - } - args.rotary_scaling = rotary_scaling - - eagle_net_config = { - 'architecture': "LlamaForCausalLM", - 'dtype': args.dtype, - 'logits_dtype': 'float32', - 'num_hidden_layers': args.n_layer_eagle, - 'num_attention_heads': args.n_head_eagle, - 'hidden_size': args.n_embd_eagle, - 'intermediate_size': args.inter_size_eagle, - 'num_key_value_heads': args.n_kv_head_eagle, - 'vocab_size': args.vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'max_position_embeddings': args.n_positions_eagle, - 'hidden_act': args.hidden_act, - 'rotary_base': args.rotary_base, - 'rotary_scaling': args.rotary_scaling, - 'norm_epsilon': args.rms_norm_eps_eagle, - 'quantization': { - 'quant_algo': None, - 'kv_cache_quant_algo': None, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'head_dim': args.head_dim_eagle, - 'head_size': args.head_size_eagle - } - - config = { - 'architecture': 'EagleForCausalLM', - 'dtype': args.dtype, - 'logits_dtype': 'float32', - 'num_hidden_layers': args.n_layer, - 'num_attention_heads': args.n_head, - 'hidden_size': args.n_embd, - 'intermediate_size': args.inter_size, - 'num_key_value_heads': args.n_kv_head, - 'vocab_size': args.vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'max_position_embeddings': args.n_positions, - 'hidden_act': args.hidden_act, - 'rotary_base': args.rotary_base, - 'rotary_scaling': args.rotary_scaling, - 'norm_epsilon': args.rms_norm_eps, - 'quantization': { - 'quant_algo': None, - 'kv_cache_quant_algo': None, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'max_draft_len': args.max_draft_len, - 'num_eagle_layers': args.num_eagle_layers, - 'max_non_leaves_per_layer': args.max_non_leaves_per_layer, - 'eagle_net_config': eagle_net_config, - 'head_dim': args.head_dim, - 'head_size': args.head_size - } - - assert args.max_draft_len <= 256, "args.max_draft_len > 256 is not supported" - - if args.use_weight_only: - if args.weight_only_precision == 'int8': - config['quantization']['quant_algo'] = QuantAlgo.W8A16 - elif args.weight_only_precision == 'int4': - config['quantization']['quant_algo'] = QuantAlgo.W4A16 - elif args.smoothquant: - if args.per_channel: - if args.per_token: - config['quantization'][ - 'quant_algo'] = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN - else: - config['quantization'][ - 'quant_algo'] = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN - else: - if args.per_token: - config['quantization'][ - 'quant_algo'] = QuantAlgo.W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN - else: - config['quantization'][ - 'quant_algo'] = QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN - - if args.int8_kv_cache: - config['quantization']['kv_cache_quant_algo'] = QuantAlgo.INT8 - - if args.weight_only_precision == 'int4_gptq': - config['quantization'].update({ - "group_size": args.group_size, - "has_zero_point": True, - "pre_quant_scale": False, - 'quant_algo': QuantAlgo.W4A16_GPTQ - }) - - # Update quant config if hf_quant_config.json exists - quant_config = {} - try: - with open(eagle_model_dir + '/' + 'hf_quant_config.json') as f: - quant_config = json.load(f) - if "lm_head" in quant_config['quantization']['exclude_modules']: - quant_config['quantization']['exclude_modules'] += [ - f"eagle_nets.{i}.lm_head" - for i in range(args.num_eagle_layers) - ] - config['quantization'].update(quant_config['quantization']) - config['eagle_net_config']['quantization'].update( - quant_config['quantization']) - except IOError: - pass - - convert_and_save_hf(config, args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') diff --git a/examples/eagle/requirements.txt b/examples/eagle/requirements.txt deleted file mode 100644 index 32bd8ce04018..000000000000 --- a/examples/eagle/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ --c ../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -rouge_score -SentencePiece~=0.1.99 -evaluate diff --git a/examples/llm-api/_tensorrt_engine/llm_eagle2_decoding.py b/examples/llm-api/_tensorrt_engine/llm_eagle2_decoding.py deleted file mode 100755 index 86b5ca28af4a..000000000000 --- a/examples/llm-api/_tensorrt_engine/llm_eagle2_decoding.py +++ /dev/null @@ -1,55 +0,0 @@ -### Generate Text Using Eagle2 Decoding - -from tensorrt_llm._tensorrt_engine import LLM -from tensorrt_llm.llmapi import (EagleDecodingConfig, KvCacheConfig, - SamplingParams) - - -def main(): - # Sample prompts. - prompts = [ - "Hello, my name is", - "The capital of France is", - "The future of AI is", - ] - # The end user can customize the sampling configuration with the SamplingParams class - sampling_params = SamplingParams(temperature=0.8, top_p=0.95) - - # The end user can customize the kv cache configuration with the KVCache class - kv_cache_config = KvCacheConfig(enable_block_reuse=True) - - llm_kwargs = {} - - model = "lmsys/vicuna-7b-v1.3" - - # The end user can customize the eagle decoding configuration by specifying the - # speculative_model, max_draft_len, num_eagle_layers, max_non_leaves_per_layer, eagle_choices - # greedy_sampling,posterior_threshold, use_dynamic_tree and dynamic_tree_max_topK - # with the EagleDecodingConfig class - - speculative_config = EagleDecodingConfig( - speculative_model="yuhuili/EAGLE-Vicuna-7B-v1.3", - max_draft_len=63, - num_eagle_layers=4, - max_non_leaves_per_layer=10, - use_dynamic_tree=True, - dynamic_tree_max_topK=10) - - llm = LLM(model=model, - kv_cache_config=kv_cache_config, - speculative_config=speculative_config, - max_batch_size=1, - max_seq_len=1024, - **llm_kwargs) - - outputs = llm.generate(prompts, sampling_params) - - # Print the outputs. - for output in outputs: - prompt = output.prompt - generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - - -if __name__ == '__main__': - main() diff --git a/examples/llm-api/_tensorrt_engine/llm_eagle_decoding.py b/examples/llm-api/_tensorrt_engine/llm_eagle_decoding.py deleted file mode 100644 index e6e89a622eeb..000000000000 --- a/examples/llm-api/_tensorrt_engine/llm_eagle_decoding.py +++ /dev/null @@ -1,60 +0,0 @@ -### Generate Text Using Eagle Decoding - -from tensorrt_llm import SamplingParams -from tensorrt_llm._tensorrt_engine import LLM -from tensorrt_llm.llmapi import EagleDecodingConfig, KvCacheConfig - - -def main(): - # Sample prompts. - prompts = [ - "Hello, my name is", - "The capital of France is", - "The future of AI is", - ] - # The end user can customize the sampling configuration with the SamplingParams class - sampling_params = SamplingParams(temperature=0.8, top_p=0.95) - - # The end user can customize the kv cache configuration with the KVCache class - kv_cache_config = KvCacheConfig(enable_block_reuse=True) - - llm_kwargs = {} - - model = "lmsys/vicuna-7b-v1.3" - - # The end user can customize the eagle decoding configuration by specifying the - # speculative_model, max_draft_len, num_eagle_layers, max_non_leaves_per_layer, eagle_choices - # greedy_sampling,posterior_threshold, use_dynamic_tree and dynamic_tree_max_topK - # with the EagleDecodingConfig class - - speculative_config = EagleDecodingConfig( - speculative_model="yuhuili/EAGLE-Vicuna-7B-v1.3", - max_draft_len=63, - num_eagle_layers=4, - max_non_leaves_per_layer=10, - eagle_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], \ - [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], \ - [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], \ - [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], \ - [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], \ - [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]] - ) - - llm = LLM(model=model, - kv_cache_config=kv_cache_config, - speculative_config=speculative_config, - max_batch_size=1, - max_seq_len=1024, - **llm_kwargs) - - outputs = llm.generate(prompts, sampling_params) - - # Print the outputs. - for output in outputs: - prompt = output.prompt - generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - - -if __name__ == '__main__': - main() diff --git a/examples/llm-api/_tensorrt_engine/llm_inference_customize.py b/examples/llm-api/_tensorrt_engine/llm_inference_customize.py deleted file mode 100644 index 523501752459..000000000000 --- a/examples/llm-api/_tensorrt_engine/llm_inference_customize.py +++ /dev/null @@ -1,55 +0,0 @@ -### Generate text with customization -import tempfile - -from tensorrt_llm._tensorrt_engine import LLM -from tensorrt_llm.llmapi import BuildConfig, KvCacheConfig, SamplingParams - - -def main(): - # The end user can customize the build configuration with the build_config class and other arguments borrowed from the lower-level APIs - build_config = BuildConfig() - build_config.max_batch_size = 128 - build_config.max_num_tokens = 2048 - - build_config.max_beam_width = 4 - - # Model could accept HF model name or a path to local HF model. - - llm = LLM( - model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", - build_config=build_config, - kv_cache_config=KvCacheConfig( - free_gpu_memory_fraction=0.8 - ), # Similar to `build_config`, you can also customize the runtime configuration with the `kv_cache_config`, `runtime_config`, `peft_cache_config` or \ - # other arguments borrowed from the lower-level APIs. - ) - - # You can save the engine to disk and load it back later, the LLM class can accept either a HF model or a TRT-LLM engine. - llm.save(tempfile.mkdtemp()) - - # Sample prompts. - prompts = [ - "Hello, my name is", - "The capital of France is", - "The future of AI is", - ] - - # With SamplingParams, you can customize the sampling strategy, such as beam search, temperature, and so on. - sampling_params = SamplingParams(temperature=0.8, - top_p=0.95, - n=4, - use_beam_search=True) - - for output in llm.generate(prompts, sampling_params): - print( - f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}" - ) - - # Got output like - # Prompt: 'Hello, my name is', Generated text: '\n\nJane Smith. I am a student pursuing my degree in Computer Science at [university]. I enjoy learning new things, especially technology and programming' - # Prompt: 'The capital of France is', Generated text: 'Paris.' - # Prompt: 'The future of AI is', Generated text: 'an exciting time for us. We are constantly researching, developing, and improving our platform to create the most advanced and efficient model available. We are' - - -if __name__ == '__main__': - main() diff --git a/examples/llm-api/_tensorrt_engine/llm_inference_kv_events.py b/examples/llm-api/_tensorrt_engine/llm_inference_kv_events.py deleted file mode 100644 index 9a1d11782d58..000000000000 --- a/examples/llm-api/_tensorrt_engine/llm_inference_kv_events.py +++ /dev/null @@ -1,49 +0,0 @@ -### Get KV Cache Events - -from tensorrt_llm import SamplingParams -from tensorrt_llm._tensorrt_engine import LLM -from tensorrt_llm.llmapi import KvCacheConfig - - -def main(): - - llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", - tensor_parallel_size=2, - enable_autotuner=False, - kv_cache_dtype='auto', - kv_cache_config=KvCacheConfig(enable_block_reuse=True, - event_buffer_max_size=1024), - backend="pytorch") - - # Sample prompts having a common prefix. - common_prefix = ( - "After the ghost's departure, Barnardo notes Horatio's pale appearance and asks if he's okay. " - "Horatio concedes that he's shaken and confesses that, without witnessing the ghost himself, he wouldn't have believed it existed. " - "He's also disturbed by the ghost's striking resemblance to the king. It even seems to be wearing the former king's armor. " - "Horatio thinks the ghost's presence foretells that something is about to go wrong in Denmark. " - "Marcellus concurs with Horatio, as he and the other guards have observed that their schedules have become more rigorous and have also noticed the preparations taking place within Elsinore, including the building of cannons, the storing of weapons, and the preparation of ships." - ) - prompts = [ - common_prefix, common_prefix + " Marcellus also notes that the king's" - ] - - # Create a sampling params. - sampling_params = SamplingParams(temperature=0.001, - top_p=0.001, - max_tokens=5) - - for output in llm.generate(prompts, sampling_params=sampling_params): - print( - f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}" - ) - - kv_events = llm.get_kv_cache_events(10) - print(kv_events) - - # Got output like follows: - # [{'event_id': 0, 'data': {'type': 'created', 'num_blocks_per_cache_level': [101230, 0]}}, - # {'event_id': 1, 'data': {'type': 'stored', 'parent_hash': None, 'blocks': [{'type': 'stored_block', 'block_hash': 4203099703668305365, 'tokens': [{'type': 'unique_token', 'token_id': 1, 'token_extra_id': 0}, ... - - -if __name__ == '__main__': - main() diff --git a/examples/llm-api/_tensorrt_engine/llm_lookahead_decoding.py b/examples/llm-api/_tensorrt_engine/llm_lookahead_decoding.py deleted file mode 100644 index ed2c94450dd4..000000000000 --- a/examples/llm-api/_tensorrt_engine/llm_lookahead_decoding.py +++ /dev/null @@ -1,38 +0,0 @@ -### Generate Text Using Lookahead Decoding -from tensorrt_llm._tensorrt_engine import LLM -from tensorrt_llm.llmapi import (BuildConfig, KvCacheConfig, - LookaheadDecodingConfig, SamplingParams) - - -def main(): - - # The end user can customize the build configuration with the build_config class - build_config = BuildConfig() - build_config.max_batch_size = 32 - - # The configuration for lookahead decoding - lookahead_config = LookaheadDecodingConfig(max_window_size=4, - max_ngram_size=4, - max_verification_set_size=4) - - kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.4) - llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", - kv_cache_config=kv_cache_config, - build_config=build_config, - speculative_config=lookahead_config) - - prompt = "NVIDIA is a great company because" - print(f"Prompt: {prompt!r}") - - sampling_params = SamplingParams(lookahead_config=lookahead_config) - - output = llm.generate(prompt, sampling_params=sampling_params) - print(output) - - #Output should be similar to: - # Prompt: 'NVIDIA is a great company because' - #RequestOutput(request_id=2, prompt='NVIDIA is a great company because', prompt_token_ids=[1, 405, 13044, 10764, 338, 263, 2107, 5001, 1363], outputs=[CompletionOutput(index=0, text='they are always pushing the envelope. They are always trying to make the best graphics cards and the best processors. They are always trying to make the best', token_ids=[896, 526, 2337, 27556, 278, 427, 21367, 29889, 2688, 526, 2337, 1811, 304, 1207, 278, 1900, 18533, 15889, 322, 278, 1900, 1889, 943, 29889, 2688, 526, 2337, 1811, 304, 1207, 278, 1900], cumulative_logprob=None, logprobs=[], finish_reason='length', stop_reason=None, generation_logits=None)], finished=True) - - -if __name__ == '__main__': - main() diff --git a/examples/llm-api/_tensorrt_engine/llm_medusa_decoding.py b/examples/llm-api/_tensorrt_engine/llm_medusa_decoding.py deleted file mode 100644 index d371600d00fc..000000000000 --- a/examples/llm-api/_tensorrt_engine/llm_medusa_decoding.py +++ /dev/null @@ -1,93 +0,0 @@ -### Generate Text Using Medusa Decoding -import argparse -from pathlib import Path - -from tensorrt_llm._tensorrt_engine import LLM -from tensorrt_llm.llmapi import (BuildConfig, KvCacheConfig, - MedusaDecodingConfig, SamplingParams) - - -def run_medusa_decoding(use_modelopt_ckpt=False, model_dir=None): - # Sample prompts. - prompts = [ - "Hello, my name is", - "The capital of France is", - "The future of AI is", - ] - # The end user can customize the sampling configuration with the SamplingParams class - sampling_params = SamplingParams(temperature=0.8, top_p=0.95) - - # The end user can customize the build configuration with the BuildConfig class - build_config = BuildConfig( - max_batch_size=1, - max_seq_len=1024, - ) - - # The end user can customize the kv cache configuration with the KVCache class - kv_cache_config = KvCacheConfig(enable_block_reuse=True) - - llm_kwargs = {} - - if use_modelopt_ckpt: - # This is a Llama-3.1-8B combined with Medusa heads provided by Model Optimizer. - # Both the base model (except lm_head) and Medusa heads have been quantized in FP8. - model = model_dir or "nvidia/Llama-3.1-8B-Medusa-FP8" - - # ModelOpt ckpt uses 3 Medusa heads - speculative_config = MedusaDecodingConfig( - max_draft_len=63, - num_medusa_heads=3, - medusa_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], \ - [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], \ - [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], \ - [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], \ - [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [1, 6], [0, 7, 0]] - ) - else: - # In this path, base model and Medusa heads are stored and loaded separately. - model = "lmsys/vicuna-7b-v1.3" - - # The end user can customize the medusa decoding configuration by specifying the - # speculative_model, max_draft_len, medusa heads num and medusa choices - # with the MedusaDecodingConfig class - speculative_config = MedusaDecodingConfig( - speculative_model="FasterDecoding/medusa-vicuna-7b-v1.3", - max_draft_len=63, - num_medusa_heads=4, - medusa_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], \ - [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], \ - [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], \ - [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], \ - [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], \ - [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]] - ) - - # Add 'tensor_parallel_size=2' if using ckpt for - # a larger model like nvidia/Llama-3.1-70B-Medusa. - llm = LLM(model=model, - build_config=build_config, - kv_cache_config=kv_cache_config, - speculative_config=speculative_config, - **llm_kwargs) - - outputs = llm.generate(prompts, sampling_params) - - # Print the outputs. - for output in outputs: - prompt = output.prompt - generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - - -if __name__ == '__main__': - parser = argparse.ArgumentParser( - description="Generate text using Medusa decoding.") - parser.add_argument( - '--use_modelopt_ckpt', - action='store_true', - help="Use FP8-quantized checkpoint from Model Optimizer.") - # TODO: remove this arg after ModelOpt ckpt is public on HF - parser.add_argument('--model_dir', type=Path, default=None) - args = parser.parse_args() - - run_medusa_decoding(args.use_modelopt_ckpt, args.model_dir) diff --git a/examples/llm-api/_tensorrt_engine/llm_quantization.py b/examples/llm-api/_tensorrt_engine/llm_quantization.py deleted file mode 100644 index 9db421a7573a..000000000000 --- a/examples/llm-api/_tensorrt_engine/llm_quantization.py +++ /dev/null @@ -1,80 +0,0 @@ -### Generation with Quantization -import logging - -import torch - -from tensorrt_llm import SamplingParams -from tensorrt_llm._tensorrt_engine import LLM -from tensorrt_llm.llmapi import CalibConfig, QuantAlgo, QuantConfig - -major, minor = torch.cuda.get_device_capability() -enable_fp8 = major > 8 or (major == 8 and minor >= 9) -enable_nvfp4 = major >= 10 - -quant_and_calib_configs = [] - -if not enable_nvfp4: - # Example 1: Specify int4 AWQ quantization to QuantConfig. - # We can skip specifying CalibConfig or leave a None as the default value. - quant_and_calib_configs.append( - (QuantConfig(quant_algo=QuantAlgo.W4A16_AWQ), None)) - -if enable_fp8: - # Example 2: Specify FP8 quantization to QuantConfig. - # We can create a CalibConfig to specify the calibration dataset and other details. - # Note that the calibration dataset could be either HF dataset name or a path to local HF dataset. - quant_and_calib_configs.append( - (QuantConfig(quant_algo=QuantAlgo.FP8, - kv_cache_quant_algo=QuantAlgo.FP8), - CalibConfig(calib_dataset='cnn_dailymail', - calib_batches=256, - calib_max_seq_length=256))) -else: - logging.error( - "FP8 quantization only works on post-ada GPUs. Skipped in the example.") - -if enable_nvfp4: - # Example 3: Specify NVFP4 quantization to QuantConfig. - quant_and_calib_configs.append( - (QuantConfig(quant_algo=QuantAlgo.NVFP4, - kv_cache_quant_algo=QuantAlgo.FP8), - CalibConfig(calib_dataset='cnn_dailymail', - calib_batches=256, - calib_max_seq_length=256))) -else: - logging.error( - "NVFP4 quantization only works on Blackwell. Skipped in the example.") - - -def main(): - - for quant_config, calib_config in quant_and_calib_configs: - # The built-in end-to-end quantization is triggered according to the passed quant_config. - llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", - quant_config=quant_config, - calib_config=calib_config) - - # Sample prompts. - prompts = [ - "Hello, my name is", - "The capital of France is", - "The future of AI is", - ] - - # Create a sampling params. - sampling_params = SamplingParams(temperature=0.8, top_p=0.95) - - for output in llm.generate(prompts, sampling_params): - print( - f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}" - ) - llm.shutdown() - - # Got output like - # Prompt: 'Hello, my name is', Generated text: 'Jane Smith. I am a resident of the city. Can you tell me more about the public services provided in the area?' - # Prompt: 'The capital of France is', Generated text: 'located in Paris, France. The population of Paris, France, is estimated to be 2 million. France is home to many famous artists, including Picasso' - # Prompt: 'The future of AI is', Generated text: 'an open and collaborative project. The project is an ongoing effort, and we invite participation from members of the community.\n\nOur community is' - - -if __name__ == '__main__': - main() diff --git a/examples/llm-api/_tensorrt_engine/quickstart_example.py b/examples/llm-api/_tensorrt_engine/quickstart_example.py deleted file mode 100644 index d02f55c46b35..000000000000 --- a/examples/llm-api/_tensorrt_engine/quickstart_example.py +++ /dev/null @@ -1,39 +0,0 @@ -from tensorrt_llm import BuildConfig, SamplingParams -from tensorrt_llm._tensorrt_engine import LLM # NOTE the change - - -def main(): - - build_config = BuildConfig() - build_config.max_batch_size = 256 - build_config.max_num_tokens = 1024 - - # Model could accept HF model name, a path to local HF model, - # or Model Optimizer's quantized checkpoints like nvidia/Llama-3.1-8B-Instruct-FP8 on HF. - llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", - build_config=build_config) - - # Sample prompts. - prompts = [ - "Hello, my name is", - "The capital of France is", - "The future of AI is", - ] - - # Create a sampling params. - sampling_params = SamplingParams(temperature=0.8, top_p=0.95) - - for output in llm.generate(prompts, sampling_params): - print( - f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}" - ) - - # Got output like - # Prompt: 'Hello, my name is', Generated text: '\n\nJane Smith. I am a student pursuing my degree in Computer Science at [university]. I enjoy learning new things, especially technology and programming' - # Prompt: 'The president of the United States is', Generated text: 'likely to nominate a new Supreme Court justice to fill the seat vacated by the death of Antonin Scalia. The Senate should vote to confirm the' - # Prompt: 'The capital of France is', Generated text: 'Paris.' - # Prompt: 'The future of AI is', Generated text: 'an exciting time for us. We are constantly researching, developing, and improving our platform to create the most advanced and efficient model available. We are' - - -if __name__ == '__main__': - main() diff --git a/examples/lookahead/README.md b/examples/lookahead/README.md deleted file mode 100644 index b7ac61fa3bcd..000000000000 --- a/examples/lookahead/README.md +++ /dev/null @@ -1,108 +0,0 @@ -# Lookahead Speculative Decoding - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a model using Lookahead speculative decoding ([Break the Sequential Dependency of LLM Inference Using Lookahead Decoidng](https://arxiv.org/pdf/2402.02057)) in TensorRT-LLM. - -## Overview - -Lookahead decoding algorithm operates through two parallel computation branches within the same LLM - a lookahead branch that generates n-grams using a fixed-sized 2D window, and a verification branch that validates promising n-gram candidates. - -Lookahead algorithm is configured with a tuple of `(windows_size, ngram_size, verification_set_size)` or, shortly, `(W, N, G)`. -+ `windows_size` is the Jacobi window size, meaning number of n-grams in lookahead branch that explores future draft tokens. -+ `ngram_size` is the n-gram size, meaning the maximum number of draft tokens accepted per iteration. -+ `verification_set_size` is the maximum number of n-grams considered for verification, meaning the number of draft token beam hypotheses. - -You can enable Lookahead decoding for any of decoder-only autoregressive LLM models without any fine-tuning. Some TensorRT LLM models might not work with Lookahead due to the missing head size in the speculative decoding XQA attention kernels. Lookahead performance greatly depends on the base model, hardware, batch size, sequence length, and the dataset. It is recommended to profile various configurations to find the best `(W, N, G)` configuration given the setup. - -Specify the Lookahead related flags in three places: - -1. *Build the engine* - -To build an engine with Lookahead support, specify the `--speculative_decoding_mode lookahead_decoding` and `--max_draft_len` arguments. -For Lookahead, the `max_draft_len` is defined as: -```python -def max_draft_len(windows_size, ngram_size, verification_set_size): - return (0 if (ngram_size == 1) else ngram_size - 2) - + (windows_size - 1 + verification_set_size) * (ngram_size - 1) -``` - -2. *Setup TensorRT LLM runtime* -When TensorRT LLM server starts, the server reserves resources according to the `executor_lookahead_config`. `executor_lookahead_config` is noted as `(W, N, G)`. Ensure the `max_draft_len` derived from `executor_lookahead_config` equals to the `max_draft_len` specified in the engine-building phase -- `--max_draft_len == max_draft_len(W, N, G)`. - -3. *Setup the request* -Each request can specify a Lookahead configuration, noted as `(w, n, g)`. If none are specified, the `executor_lookahead_config` is used. The minimum Lookahead config `(1, 1, 0)` forces non speculative, autoregressive mode. The meaningful minimum configuration is `(2, 2, 1)`. Ensure the Lookahead configuration for each request satisfies `w <= W, n <= N, g <= G`. - -## Support Matrix - * GPU Compute Capability >= 8.0 (Ampere or newer) - * FP16 / BF16 / FP8 - * Paged KV Cache - * Inflight-fused-batching - * C++ runtime - * Tensor Parallel - -## Usage -### Convert Checkpoint - -This example is based on the Vicuna-7b v1.3 model, a fine-tuned Llama model. -Checkpoint conversion is similar to any standard autoregressive model, such as the models located in the [examples/models/core/llama](../../examples/models/core/llama) directory. - -```bash -MODEL_DIR=/path/to/vicuna-7b-v1.3 -ENGINE_DIR=tmp/engine -CKPT_DIR=tmp/engine/ckpt - -python3 examples/models/core/llama/convert_checkpoint.py \ - --model_dir=$MODEL_DIR \ - --output_dir=$CKPT_DIR \ - --dtype=float16 \ - --tp_size=1 \ - --pp_size=1 -``` - -### Build engine - -```bash -trtllm-build \ - --checkpoint_dir=$CKPT_DIR \ - --output_dir=$ENGINE_DIR \ - --gpt_attention_plugin=float16 \ - --gemm_plugin=float16 \ - --max_batch_size=32 \ - --max_input_len=1024 \ - --max_seq_len=2048 \ - --max_beam_width=1 \ - --log_level=error \ - --max_draft_len=83 \ - --speculative_decoding_mode=lookahead_decoding -``` - -### Run decoding - -+ `--lookahead_config` is a server-level configuration of Lookahead decoding in the form of a triplet `[W, N, G]`. Note that `run.py` and `summarize.py` interfaces allow to set only per server but not per-request config. - -Run `examples/run.py` to generate sequences. -```bash -python examples/run.py \ - --tokenizer_dir=$MODEL_DIR \ - --engine_dir=$ENGINE_DIR \ - --max_output_len=32 \ - --lookahead_config=[7,7,7] \ - --log_level=verbose \ - --input_text 'Once upon' 'To be, or not' 'Be not afraid of greatness' -``` - -Run `examples/summarize.py` to summarize the CNN daily dataset. -```bash -python examples/summarize.py \ - --test_hf \ - --test_trt_llm \ - --hf_model_dir=$MODEL_DIR \ - --engine_dir=$ENGINE_DIR \ - --data_type=fp16 \ - --lookahead_config=[7,7,7] -``` diff --git a/examples/lookahead/requirements.txt b/examples/lookahead/requirements.txt deleted file mode 100644 index 423ba2d4b0e0..000000000000 --- a/examples/lookahead/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ --c ../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -rouge_score -sentencepiece>=0.1.99 -evaluate diff --git a/examples/medusa/README.md b/examples/medusa/README.md deleted file mode 100644 index a596608625ea..000000000000 --- a/examples/medusa/README.md +++ /dev/null @@ -1,232 +0,0 @@ -# Medusa Decoding - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a model using Medusa decoding([`Github`](https://github.com/FasterDecoding/Medusa), [`BLOG`](https://sites.google.com/view/medusa-llm)) in TensorRT LLM on single GPU, single node multiple GPU. - -## Overview -Different from other models, Medusa decoding needs a base model and Medusa heads. The TensorRT LLM Medusa Decoding implementation can be found in [tensorrt_llm/models/medusa/model.py](../../tensorrt_llm/models/medusa/model.py). The implementation adds Medusa heads to a base model. - -For more info about Medusa visit [speculative decoding documentation](https://nvidia.github.io/TensorRT-LLM/features/speculative-decoding.html). - -## Support Matrix - * GPU Compute Capability >= 8.0 (Ampere or newer) - * FP16 - * BF16 - * FP8 (base model) - * PAGED_KV_CACHE - * Tensor Parallel - -## Usage -The TensorRT LLM Medusa example code is located in [`examples/medusa`](./). There is one [`convert_checkpoint.py`](./convert_checkpoint.py) file to convert and build the [TensorRT](https://developer.nvidia.com/tensorrt) engine(s) needed to run models with Medusa decoding support. -In this example, we demonstrate the usage of two models: -1. The Vucuna 7B model from Hugging Face [`FasterDecoding/medusa-vicuna-7b-v1.3`](https://huggingface.co/FasterDecoding/medusa-vicuna-7b-v1.3) with its Medusa heads [`medusa-vicuna-7b-v1.3`](https://huggingface.co/FasterDecoding/medusa-vicuna-7b-v1.3). -2. The quantized checkpoint [`nvidia/Llama-3.1-8B-Medusa-FP8`](https://huggingface.co/nvidia/Llama-3.1-8B-Medusa-FP8) on Hugging Face by [Model Optimizer](https://github.com/NVIDIA/Model-Optimizer) (ModelOpt). This model is based on [Llama-3.1 8B](https://huggingface.co/meta-llama/Llama-3.1-8B) and enhanced with Medusa heads, with both the base model (except lm_head) and Medusa heads already quantized in FP8. - -### Build TensorRT engine(s) -Get the weights by downloading base model [`vicuna-7b-v1.3`](https://huggingface.co/lmsys/vicuna-7b-v1.3) and Medusa Heads [`medusa-vicuna-7b-v1.3`](https://huggingface.co/FasterDecoding/medusa-vicuna-7b-v1.3) from HF. - -``` -pip install -r requirements.txt - -git lfs install -git clone https://huggingface.co/lmsys/vicuna-7b-v1.3 -https://huggingface.co/FasterDecoding/medusa-vicuna-7b-v1.3 -``` - -We use `convert_checkpoint.py` script to convert the model for Medusa decoding into TensorRT LLM checkpoint format. -We could use `--num_medusa_heads` to set the number of medusa heads that we want to use. If not, `num_medusa_heads` will be set according to the `medusa_num_heads` from medusa weights' `config.json`. - -Here is the example: -```bash -# Convert and Build Medusa decoding support for vicuna-7b-v1.3 -python convert_checkpoint.py --model_dir ./vicuna-7b-v1.3 \ - --medusa_model_dir medusa-vicuna-7b-v1.3 \ - --output_dir ./tllm_checkpoint_1gpu_medusa \ - --dtype float16 \ - --num_medusa_heads 4 - -# Note: Increasing the batch size may have a negative impact on performance -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_medusa \ - --output_dir ./tmp/medusa/7B/trt_engines/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --speculative_decoding_mode medusa \ - --max_batch_size 4 - -# Convert and Build Medusa decoding support for vicuna-13b-v1.3 with 4-way tensor parallelism. -python convert_checkpoint.py --model_dir ./vicuna-7b-v1.3 \ - --medusa_model_dir medusa-vicuna-7b-v1.3 \ - --output_dir ./tllm_checkpoint_1gpu_medusa \ - --dtype float16 \ - --num_medusa_heads 4 \ - --tp_size 4 \ - --workers 4 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_medusa \ - --output_dir ./tmp/medusa/7B/trt_engines/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --speculative_decoding_mode medusa \ - --max_batch_size 4 - -# Convert and Build Llama-3.1-8B-Medusa by ModelOpt -python convert_checkpoint.py --model_dir ./llama3.1-medusa-8b-hf_v0.1 \ - --output_dir ./tllm_checkpoint_1gpu_modelopt_llama_medusa \ - --dtype float16 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_modelopt_llama_medusa \ - --output_dir ./tmp/modelopt/llama-8B-medusa/trt_engines/1-gpu/ \ - --gemm_plugin float16 \ - --speculative_decoding_mode medusa \ - --max_batch_size 4 - - -# Convert and Build Llama-3.1-70B-Medusa by ModelOpt with 2-way tensor parallelism. -python convert_checkpoint.py --model_dir ./llama-3.1-70b-medusa_vfp8-fp8-fp8 \ - --output_dir ./tllm_checkpoint_2gpu_modelopt_llama_medusa_70b \ - --dtype float16 - --tp_size 2 - --workers 2 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_2gpu_modelopt_llama_medusa_70b \ - --output_dir ./tmp/modelopt/llama-70B-medusa/trt_engines/2-gpu/ \ - --gemm_plugin float16 \ - --speculative_decoding_mode medusa \ - --max_batch_size 4 - - -``` - -### FP8 Post-Training Quantization for Base Model -The example below quantizes the base model to FP8, while keeping the weight of the medusa head non-quantize. -```bash -# Quantize base model into FP8 and export trtllm checkpoint -python ../quantization/quantize.py --model_dir /path/to/base-model-hf/ \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir ./tllm_checkpoint_1gpu_base_model_fp8_medusa_fp16 \ - --calib_size 512 \ - --tp_size 1 \ - --medusa_model_dir /path/to/medusa_head/ \ - --num_medusa_heads 4 - -# Build trtllm engines from the trtllm checkpoint -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_base_model_fp8_medusa_fp16 \ - --output_dir ./trt_engine_1gpu_base_model_fp8_medusa_fp16 \ - --gemm_plugin float16 \ - --gpt_attention_plugin float16 \ - --speculative_decoding_mode medusa \ - --max_batch_size 4 -``` - -### Run -To run a TensorRT LLM model with Medusa decoding support, we can use `../run.py` script, with an additional argument `--medusa_choices`. -The `--medusa_choices` is of type `list[list[int]]`. - -Medusa decoding is supported by Python runtime and C++ runtime with inflight-batching. C++ runtime is recommended for performance. -For Python runtime use `--use_py_session` flag to `run.py`. - -Medusa decoding only supporting greedy decoding, indicated by `temperature=1.0` argument. The output is equivalent to the base model inference with `--temperature 0.0` (equivalent to `--temperature 1.0 --top-k 1`). - -```bash -# Medusa decoding using vicuna-7b-v1.3 model with 1 GPU -python ../run.py --engine_dir ./tmp/medusa/7B/trt_engines/fp16/1-gpu/ \ - --tokenizer_dir ./vicuna-7b-v1.3/ \ - --max_output_len=100 \ - --medusa_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --temperature 1.0 \ - --input_text "Once upon" - -# Medusa decoding using vicuna-13b-v1.3 with 4 GPUs -mpirun -np 4 --allow-run-as-root --oversubscribe \ - python ../run.py --engine_dir ./tmp/medusa/13B/trt_engines/fp16/4-gpu/ \ - --tokenizer_dir ./vicuna-13b-v1.3/ \ - --max_output_len=100 \ - --medusa_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --temperature 1.0 \ - --input_text "Once upon" - -# Medusa decoding using Llama-3.1-8B-Medusa by ModelOpt with 1 GPU -python ../run.py --engine_dir ./tmp/modelopt/llama-8B-medusa/trt_engines/1-gpu/ \ - --tokenizer_dir ./llama3.1-medusa-8b-hf_v0.1 \ - --max_output_len=100 \ - --medusa_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [1, 6], [0, 7, 0]]" \ - --temperature 1.0 \ - --input_text "Once upon" - -# Medusa decoding using Llama-3.1-70B-Medusa by ModelOpt with 2 GPUs -mpirun -np 2 --allow-run-as-root --oversubscribe \ - python ../run.py --engine_dir ./tmp/modelopt/llama-70B-medusa/trt_engines/2-gpu/ \ - --tokenizer_dir ./llama-3.1-70b-medusa_vfp8-fp8-fp8 \ - --max_output_len=100 \ - --medusa_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --temperature 1.0 \ - --input_text "Once upon" - -``` - -And you will see output like this if run successfully: -```text -...... -Input [Text 0]: " Once upon" -Output [Text 0 Beam 0]: "a time, there was a young girl who loved to read. She would spend hours in the library, devouring books of all genres. She had a special love for fairy tales, and would often dream of living in a magical world where she could meet princes and princesses, and have adventures with talking animals. -One day, while she was reading a book, she came across a passage that spoke to her heart. It said, "You are the author of" -``` - -### Summarization using Medusa decoding - -```bash -# Medusa decoding using vicuna-7b-v1.3 model with 1 GPU -python ../summarize.py --engine_dir ./tmp/medusa/7B/trt_engines/fp16/1-gpu/ \ - --hf_model_dir ./vicuna-7b-v1.3/ \ - --tokenizer_dir ./vicuna-7b-v1.3/ \ - --test_trt_llm \ - --data_type fp16 \ - --medusa_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --use_py_session \ - --temperature 1.0 \ - --batch_size 1 - -# Medusa decoding using vicuna-13b-v1.3 with 4 GPUs -mpirun -np 4 --allow-run-as-root --oversubscribe \ - python ../summarize.py --engine_dir ./tmp/medusa/13B/trt_engines/fp16/4-gpu/ \ - --hf_model_dir ./vicuna-13b-v1.3/ \ - --tokenizer_dir ./vicuna-13b-v1.3/ \ - --test_trt_llm \ - --data_type fp16 \ - --medusa_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --use_py_session \ - --temperature 1.0 \ - --batch_size 1 - -# Medusa decoding using Llama-3.1-8B-Medusa by ModelOpt with 1 GPU -python ../summarize.py --engine_dir ./tmp/modelopt/llama-8B-medusa/trt_engines/1-gpu/ \ - --hf_model_dir ./llama3.1-medusa-8b-hf_v0.1 \ - --tokenizer_dir ./llama3.1-medusa-8b-hf_v0.1 \ - --test_trt_llm \ - --data_type fp16 \ - --medusa_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [1, 6], [0, 7, 0]]" \ - --use_py_session \ - --temperature 1.0 \ - --batch_size 1 - -# Medusa decoding using Llama-3.1-70B-Medusa by ModelOpt with 2 GPUs -mpirun -np 2 --allow-run-as-root --oversubscribe \ - python ../summarize.py --engine_dir ./tmp/modelopt/llama-70B-medusa/trt_engines/2-gpu/ \ - --hf_model_dir ./llama-3.1-70b-medusa_vfp8-fp8-fp8 \ - --tokenizer_dir ./llama-3.1-70b-medusa_vfp8-fp8-fp8 \ - --test_trt_llm \ - --data_type fp16 \ - --medusa_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --use_py_session \ - --temperature 1.0 \ - --batch_size 1 -``` - -### Medusa with Qwen2 - -To use Medusa with Qwen2 models, specify `--model_type qwen2` to `convert_checkpoint.py`. You have to provide a Qwen2 model checkpoint and the medusa heads. After TRT-LLM checkpoint is generated, trllm-build and `../run.py` use the same arguments as for LLaMA models. diff --git a/examples/medusa/convert_checkpoint.py b/examples/medusa/convert_checkpoint.py deleted file mode 100644 index 09eb55b46105..000000000000 --- a/examples/medusa/convert_checkpoint.py +++ /dev/null @@ -1,487 +0,0 @@ -import argparse -import json -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed -from pathlib import Path - -import safetensors -import torch -from transformers import (LlamaConfig, LlamaForCausalLM, LlamaTokenizer, - Qwen2Config) - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import get_hf_rope_theta, numpy_to_torch -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import (LLaMAForCausalLM, PretrainedConfig, - QWenForCausalLM) -from tensorrt_llm.models.convert_utils import load_calib_dataset -from tensorrt_llm.models.llama.convert import load_weights_from_hf_by_shard -from tensorrt_llm.models.medusa.weight import (capture_activation_range, - convert_hf_llama, load_medusa_hf) -from tensorrt_llm.quantization import QuantAlgo - -try: - from transformers import MixtralForCausalLM -except ImportError: - MixtralForCausalLM = None - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument( - '--model_dir', - type=str, - help="base model or Medusa-enhanced quantized model from ModelOpt") - parser.add_argument('--meta_ckpt_dir', type=str, default=None) - 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('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument('--vocab_size', type=int, default=32000) - parser.add_argument('--n_positions', type=int, default=2048) - parser.add_argument('--n_layer', type=int, default=32) - - 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', 'int4_gptq'], - 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( - '--calib_dataset', - type=str, - default='ccdv/cnn_dailymail', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - 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( - '--per_channel', - action="store_true", - default=False, - 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', - action="store_true", - default=False, - 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( - '--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( - '--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('--load_by_shard', - action='store_true', - help='Load a pretrained model shard-by-shard.') - parser.add_argument('--hidden_act', type=str, default='silu') - - parser.add_argument('--rotary_base', type=float, default=10000.0) - parser.add_argument('--rotary_scaling', nargs=2, type=str, default=None) - - parser.add_argument('--group_size', - type=int, - default=128, - help='Group size used in GPTQ/AWQ quantization.') - - parser.add_argument("--storage-type", - "-t", - 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") - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - - parser.add_argument('--num_medusa_heads', type=int, default=4) - parser.add_argument('--num_medusa_layers', type=int, default=1) - parser.add_argument('--max_medusa_token_len', type=int, default=63) - parser.add_argument('--medusa_hidden_act', type=str, default="silu") - parser.add_argument('--medusa_model_dir', type=str, default=None) - parser.add_argument('--model_type', type=str, default="llama") - args = parser.parse_args() - return args - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - # TODO(qijun): Currently, the convert script depends on a torch op: - # torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix, - # which is included in tensorrt_llm Python package. Otherwise, the convert - # script does not need to import tensorrt_llm. Will remove it after reimplementing - # the op with PyTorch. - print(tensorrt_llm.__version__) - args = parse_arguments() - assert args.model_type in ["llama", "mixtral", - "qwen2"], "Invalid model type" - world_size = args.tp_size * args.pp_size - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - if args.model_dir is not None: - config_cls = Qwen2Config if args.model_type == "qwen2" else LlamaConfig - hf_config = config_cls.from_pretrained(args.model_dir) - - args.model_type = hf_config.model_type - args.n_head = hf_config.num_attention_heads - args.inter_size = hf_config.intermediate_size - args.n_layer = hf_config.num_hidden_layers - args.n_embd = hf_config.hidden_size - args.n_kv_head = hf_config.num_key_value_heads - args.rms_norm_eps = hf_config.rms_norm_eps - args.vocab_size = hf_config.vocab_size - args.n_positions = hf_config.max_position_embeddings - args.rotary_base = get_hf_rope_theta(hf_config, 10000.0) - args.rotary_scaling = hf_config.rope_scaling - - elif args.meta_ckpt_dir is not None: - - with open(Path(args.meta_ckpt_dir, "params.json")) as fp: - meta_config: dict = json.load(fp) - args.n_embd = meta_config["dim"] - args.n_head = meta_config["n_heads"] - args.n_layer = meta_config["n_layers"] - args.n_kv_head = meta_config.get("n_kv_heads", args.n_head) - - if "hidden_dim" in meta_config: - args.inter_size = meta_config["hidden_dim"] - else: - args.multiple_of = meta_config.get("multiple_of", 1) - n_embd = int(4 * args.n_embd * 2 / 3) - args.ffn_dim_multiplier = meta_config.get("ffn_dim_multiplier", 1) - args.inter_size = args.multiple_of * ( - (int(n_embd * args.ffn_dim_multiplier) + args.multiple_of - 1) - // args.multiple_of) - args.rms_norm_eps = meta_config["norm_eps"] - - 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["rope_type"], - } - args.rotary_scaling = rotary_scaling - - # ModelOpt quantized ckpt - quant_config_file = Path(args.model_dir) / "hf_quant_config.json" - - if quant_config_file.exists(): - with open(quant_config_file, 'r') as f: - quant_config = json.load(f) - - is_modelopt_ckpt = quant_config.get("producer", - {}).get("name") == "modelopt" - if is_modelopt_ckpt: # quantized ckpt - args.medusa_model_dir = args.model_dir - args.num_medusa_heads = hf_config.medusa["num_medusa_heads"] - args.num_medusa_layers = hf_config.medusa["num_medusa_layers"] - if args.smoothquant or args.calib_dataset: - logger.warning( - "Checkpoint is already quantized. All quantization-related flags will be ignored." - ) - else: - is_modelopt_ckpt = False - - config = { - 'architecture': 'MedusaForCausalLM', - 'dtype': args.dtype, - 'logits_dtype': 'float32', - 'num_hidden_layers': args.n_layer, - 'num_attention_heads': args.n_head, - 'hidden_size': args.n_embd, - 'intermediate_size': args.inter_size, - 'num_key_value_heads': args.n_kv_head, - 'vocab_size': args.vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'max_position_embeddings': args.n_positions, - 'hidden_act': args.hidden_act, - 'rotary_base': args.rotary_base, - 'rotary_scaling': args.rotary_scaling, - 'norm_epsilon': args.rms_norm_eps, - 'quantization': { - 'quant_algo': None, - 'kv_cache_quant_algo': None, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'max_draft_len': args.max_medusa_token_len, - 'num_medusa_heads': args.num_medusa_heads, - 'num_medusa_layers': args.num_medusa_layers, - 'model_type': args.model_type, - } - if args.model_type == "qwen2": - config['qwen_type'] = args.model_type - - if is_modelopt_ckpt: - with open(quant_config_file, "r") as f: - hf_quant_config = json.load(f) - for key, value in hf_quant_config.get("quantization", {}).items(): - config['quantization'][key] = value - elif args.use_weight_only: - if args.weight_only_precision == 'int8': - config['quantization']['quant_algo'] = QuantAlgo.W8A16 - elif args.weight_only_precision == 'int4': - config['quantization']['quant_algo'] = QuantAlgo.W4A16 - elif args.smoothquant: - if args.per_channel: - if args.per_token: - config['quantization'][ - 'quant_algo'] = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN - else: - config['quantization'][ - 'quant_algo'] = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN - else: - if args.per_token: - config['quantization'][ - 'quant_algo'] = QuantAlgo.W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN - else: - config['quantization'][ - 'quant_algo'] = QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN - - if args.int8_kv_cache: - config['quantization']['kv_cache_quant_algo'] = QuantAlgo.INT8 - - if args.weight_only_precision == 'int4_gptq': - config['quantization'].update({ - "group_size": args.group_size, - "has_zero_point": True, - "pre_quant_scale": False, - 'quant_algo': QuantAlgo.W4A16_GPTQ - }) - - with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: - json.dump(config, f, indent=4) - - if args.weight_only_precision == 'int8': - plugin_weight_only_quant_type = torch.int8 - elif args.weight_only_precision == 'int4': - plugin_weight_only_quant_type = torch.quint4x2 - - act_range = {} - llama_qkv_para = {} - # smoother for inputs of self_attn.o_proj and mlp.down_proj - llama_smoother = {} - model = None - if args.model_dir is not None: - if args.model_type == "qwen2": - model = QWenForCausalLM.from_hugging_face(args.model_dir, - args.dtype) - elif not is_modelopt_ckpt: - hf_model = LlamaForCausalLM if args.model_type != "mixtral" else MixtralForCausalLM - model = hf_model.from_pretrained( - args.model_dir, - dtype='auto', - device_map='auto' if not args.load_model_on_cpu else 'cpu', - trust_remote_code=True) - - if args.smoothquant is not None or args.int8_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 slow." - ) - tokenizer = LlamaTokenizer.from_pretrained(args.model_dir, - padding_side='left') - dataset = load_calib_dataset(args.calib_dataset, - cache_dir=args.dataset_cache_dir) - - act_range = capture_activation_range(model, tokenizer, dataset) - if args.smoothquant is not None: - smooth_llama_model(model, act_range, args.smoothquant, - llama_qkv_para, llama_smoother) - convert_args = { - 'hf_model': model, - 'act_range': act_range, - 'llama_qkv_para': llama_qkv_para, - 'llama_smoother': llama_smoother, - 'config': config, - 'is_modelopt_ckpt': is_modelopt_ckpt - } - - def covert_and_save(rank, convert_args): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - - if convert_args["is_modelopt_ckpt"]: - convert_args["hf_model"] = LLaMAForCausalLM.from_hugging_face( - args.model_dir, - args.dtype, - mapping=mapping, - quant_config=convert_args["config"]["quantization"], - load_by_shard=args.load_by_shard) - - if args.use_weight_only and args.weight_only_precision == 'int4_gptq': - assert False, "Never supported" - else: - if args.load_by_shard: - weights = load_weights_from_hf_by_shard( - args.model_dir, PretrainedConfig.from_dict(config)) - - else: - if args.model_type == "qwen2" or convert_args[ - "is_modelopt_ckpt"]: - weights = { - name: numpy_to_torch(param.raw_value) - for name, param in - convert_args['hf_model'].named_parameters() - } - else: - weights = convert_hf_llama( - convert_args['hf_model'], - mapping, - rank, - dtype=args.dtype, - use_weight_only=args.use_weight_only, - plugin_weight_only_quant_type= - plugin_weight_only_quant_type, - use_parallel_embedding=args.use_parallel_embedding, - sharding_dim=args.embedding_sharding_dim, - use_smooth_quant=args.smoothquant, - per_channel=args.per_channel, - per_token=args.per_token, - int8_kv_cache=args.int8_kv_cache, - act_range=convert_args['act_range'], - qkv_para=convert_args['llama_qkv_para'], - smoother=convert_args['llama_smoother']) - - if args.medusa_model_dir is not None: - config_file = Path(args.medusa_model_dir) / "config.json" - with open(config_file) as fp: - config = json.load(fp) - if not convert_args["is_modelopt_ckpt"]: - num_medusa_heads_from_config = config.get( - 'medusa_num_heads', args.num_medusa_heads) - args.num_medusa_layers = config.get( - 'medusa_num_layers', args.num_medusa_layers) - if args.num_medusa_heads is None: - args.num_medusa_heads = num_medusa_heads_from_config - - assert args.max_medusa_token_len > 0, "should have max_medusa_token_len > 0" - - medusa_weights = load_medusa_hf( - medusa_path=args.medusa_model_dir, - num_medusa_heads=args.num_medusa_heads, - num_medusa_layers=args.num_medusa_layers, - mapping=mapping, - dtype=args.dtype, - use_weight_only=args.use_weight_only, - plugin_weight_only_quant_type= - plugin_weight_only_quant_type, - is_modelopt_ckpt=convert_args["is_modelopt_ckpt"]) - weights.update(medusa_weights) - - safetensors.torch.save_file( - weights, os.path.join(args.output_dir, f'rank{rank}.safetensors')) - - if args.workers == 1: - for rank in range(world_size): - covert_and_save(rank, convert_args) - else: - with ThreadPoolExecutor(max_workers=args.workers) as p: - futures = [ - p.submit(covert_and_save, rank, convert_args) - for rank in range(world_size) - ] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/medusa/requirements.txt b/examples/medusa/requirements.txt deleted file mode 100644 index 423ba2d4b0e0..000000000000 --- a/examples/medusa/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ --c ../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -rouge_score -sentencepiece>=0.1.99 -evaluate diff --git a/examples/models/contrib/baichuan/README.md b/examples/models/contrib/baichuan/README.md deleted file mode 100644 index b1da0a6b0319..000000000000 --- a/examples/models/contrib/baichuan/README.md +++ /dev/null @@ -1,314 +0,0 @@ -# Baichuan - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a Baichuan models (including `v1_7b`/`v1_13b`/`v2_7b`/`v2_13b`) in TensorRT LLM on both single GPU and single node multi-GPU. - -## Table of Contents - -- [Baichuan](#baichuan) - - [Table of Contents](#table-of-contents) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [Build TensorRT engine(s)](#build-tensorrt-engines) - - [SmoothQuant](#smoothquant) - - [FP8 Post-Training Quantization](#fp8-post-training-quantization) - - [Groupwise quantization (AWQ/GPTQ)](#groupwise-quantization-awqgptq) - - [AWQ](#awq) - - [GPTQ](#gptq) - - [INT8 KV cache](#int8-kv-cache) - - [Run](#run) - - [Summarization using the Baichuan model](#summarization-using-the-baichuan-model) - -## 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/models/contrib/baichuan`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert supported checkpoints into TensorRT LLM format. - -The script accepts an argument named model_version, whose value should be `v1_7b`/`v1_13b`/`v2_7b`/`v2_13b` and the default value is `v1_13b`. - -In addition, there are two shared files in the folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 - * FP8 - * BF16 - * INT4 & INT8 Weight-Only - * INT8 SmoothQuant - * Groupwise quantization (AWQ/GPTQ) - * INT8 KV CACHE (+ AWQ/per-channel weight-only/SmoothQuant) - -## Usage - -The TensorRT LLM Baichuan example code locates at [examples/models/contrib/baichuan](./). It takes HF weights as input, and builds the corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Build TensorRT engine(s) - -Please install required packages first: - -```bash -pip install -r requirements.txt -``` - -Need to specify the HF Baichuan checkpoint path. For `v1_13b`, you should use whether [baichuan-inc/Baichuan-13B-Chat](https://huggingface.co/baichuan-inc/Baichuan-13B-Chat) or [baichuan-inc/Baichuan-13B-Base](https://huggingface.co/baichuan-inc/Baichuan-13B-Base). For `v2_13b`, you should use whether [baichuan-inc/Baichuan2-13B-Chat](https://huggingface.co/baichuan-inc/Baichuan2-13B-Chat) or [baichuan-inc/Baichuan2-13B-Base](https://huggingface.co/baichuan-inc/Baichuan2-13B-Base). More Baichuan models could be found on [baichuan-inc](https://huggingface.co/baichuan-inc). - -TensorRT LLM Baichuan builds TensorRT engine(s) from HF checkpoint. If no checkpoint directory is specified, TensorRT LLM will build engine(s) with dummy weights. - -***For all kinds of checkpoints, they share the same trtllm-build command like:*** - -```bash -# Enable several TensorRT LLM plugins to increase runtime performance. It also helps with build time. - -# The TensorRT LLM GPT Attention plugin (--gpt_attention_plugin) is -# enabled by default to increase runtime performance. -# 7B models should always enable `gpt_attention_plugin`` since RoPE is only -# supported with GPTAttention plugin now. -# Try gemm_plugin to prevent accuracy issue. -trtllm-build --checkpoint_dir ./tmp/baichuan_v1_13b/trt_ckpts/fp16/1-gpu/ \ - --output_dir ./tmp/baichuan_v1_13b/trt_engines/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size=32 \ - --max_input_len=1024 \ - --max_seq_len=1536 -``` - - -Here're some examples for checkpoint conversion that take `v1_13b` as example: - -```bash -# Convert the Baichuan V1 13B model using a single GPU and FP16. -python convert_checkpoint.py --model_version v1_13b \ - --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype float16 \ - --output_dir ./tmp/baichuan_v1_13b/trt_ckpts/fp16/1-gpu/ - -# Convert the Baichuan V1 13B model using a single GPU and BF16. -python convert_checkpoint.py --model_version v1_13b \ - --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype bfloat16 \ - --output_dir ./tmp/baichuan_v1_13b/trt_ckpts/bf16/1-gpu/ - -# Convert the Baichuan V1 13B model using a single GPU and apply INT8 weight-only quantization. -python convert_checkpoint.py --model_version v1_13b \ - --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype float16 \ - --use_weight_only \ - --output_dir ./tmp/baichuan_v1_13b/trt_ckpts/int8_weight_only/1-gpu/ - -# Convert the Baichuan V1 13B model using a single GPU and apply INT4 weight-only quantization. -# Note that Baichuan V1 7B performs not well when using INT4 weight-only quantization. -python convert_checkpoint.py --model_version v1_13b \ - --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int4 \ - --output_dir ./tmp/baichuan_v1_13b/trt_ckpts/int4_weight_only/1-gpu/ - -# Convert Baichuan V1 13B using 2-way tensor parallelism. -python convert_checkpoint.py --model_version v1_13b \ - --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype float16 \ - --output_dir ./tmp/baichuan_v1_13b/trt_ckpts/fp16/1-gpu/ \ - --tp_size 2 -``` - -#### SmoothQuant - -The SmoothQuant supports all Baichuan model variants. 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. - -`--smoothquant` is the starting point of INT8 inference. By default, it -will run the model in the _per-tensor_ mode. - -Then, you can add any combination of `--per-token` and `--per-channel` to get the corresponding behaviors. - -Examples of build invocations: -```bash -python convert_checkpoint.py --model_version v1_13b \ - --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype float16 \ - --smoothquant 0.8 \ - --per_channel \ - --per_token \ - --output_dir ./tmp/baichuan_v1_13b/sq0.8/1-gpu/ -``` - -#### FP8 Post-Training Quantization - -The examples below uses the NVIDIA Modelopt (AlgorithMic Model Optimization) toolkit for the model quantization process. - -First make sure Modelopt toolkit is installed (see [examples/quantization/README.md](/examples/quantization/README.md#preparation)) - -```bash -# Quantize HF Baichuan v2 13B into FP8 and export a single-rank checkpoint -python ../../../quantization/quantize.py --model_dir /code/model/Baichuan2-13B-Chat/ \ - --dtype float16 \ - --qformat fp8 \ - --output_dir ./quantized_fp8 \ - --calib_size 256 -``` - -The quantized model checkpoint is saved to `./quantized_fp8/` for future TensorRT LLM engine build directly with the `trtllm-build` command mentioned above. -Note that you can enable fp8 context fmha to get further acceleration by setting `--use_fp8_context_fmha enable` when building the engines. - -#### Groupwise quantization (AWQ/GPTQ) -##### AWQ -NVIDIA Modelopt toolkit is used for AWQ weight quantization. Please see [examples/quantization/README.md](/examples/quantization/README.md#preparation) for Modelopt installation instructions. -```bash -# Quantize HF Baichuan 13B checkpoint into INT4 AWQ format -python ../../../quantization/quantize.py --model_dir /code/model/Baichuan2-13B-Chat/ \ - --dtype float16 \ - --qformat int4_awq \ - --output_dir ./quantized_int4-awq_gs128 \ - --calib_size 32 -``` -The quantized model checkpoint is saved to `./quantized_int4-awq_gs128/` for future TensorRT LLM engine build directly with the `trtllm-build` command mentioned above. - -##### GPTQ -To run the GPTQ Baichuan example, the following steps are required: - -1. Weight quantization: - - Quantized weights for GPTQ can be generated using an open source project such as [GPTQ-for-LLaMa](https://github.com/qwopqwop200/GPTQ-for-LLaMa.git). - - Let us build the TensorRT LLM engine with the saved `./baichuan-2-13b-4bit-gs64.safetensors`. - -2. Checkpoint conversion: - - ```bash - # Build the Baichuan2 13B model using 2-way tensor parallelism and apply INT4 GPTQ quantization. - # Compressed checkpoint safetensors are generated separately from GPTQ. - python convert_checkpoint.py --model_version v2_13b \ - --quant_ckpt_path ./baichuan-2-13b-4bit-gs64.safetensors \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int4_gptq \ - --group_size 64 \ - --tp_size 2 \ - --output_dir ./tmp/baichuan_v2_13b/trt_ckpts/int4_gptq_gs64/2-gpu/ - ``` - The quantized model checkpoint is saved for future TensorRT LLM engine build directly with the `trtllm-build` command mentioned above. - -#### INT8 KV cache -INT8 KV cache could be enabled to reduce memory footprint. It will bring more performance gains when batch size gets larger. - -You can get the INT8 scale of KV cache through NVIDIA Modelopt (AlgorithMic Model Optimization) toolkit, which features a -`--kv_cache_dtype` option. - -Example: - -```bash -python ../../../quantization/quantize.py --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype float16 \ - --kv_cache_dtype int8 \ - --output_dir ./trt_ckpt/baichuan_int8kv_tp1 \ - --calib_size 512 -``` - -**INT8 KV cache + per-channel weight-only quantization** - -INT8 KV cache could be combined with per-channel weight-only quantization, as follows: -```bash -python ../../../quantization/quantize.py --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype float16 \ - --qformat int4_wo \ - --kv_cache_dtype int8 \ - --output_dir ./trt_ckpt/baichuan_int4wo_int8kv_tp1 \ - --calib_size 512 \ -``` - -**INT8 KV cache + AWQ** - -In addition, you can enable INT8 KV cache together with AWQ (per-group INT4 weight-only quantization), as follows: - -```bash -python ../../../quantization/quantize.py --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype float16 \ - --qformat int4_awq \ - --kv_cache_dtype int8 \ - --output_dir ./trt_ckpt/baichuan_int4awq_int8kv_tp1 \ - --calib_size 512 -``` - -**INT8 KV cache + INT8 SmoothQuant** - -In addition, you can enable INT8 KV cache together with INT8 SmoothQuant, as follows: - -```bash -python convert_checkpoint.py --model_version v1_13b \ - --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype float16 \ - --smoothquant 0.8 \ - --per_channel \ - --per_token \ - --int8_kv_cache \ - --output_dir ./tmp/baichuan_v1_13b/sq0.8/1-gpu/ -``` - -### Run - -To run a TensorRT LLM Baichuan model using the engines generated by `trtllm-build` - -```bash -# With fp16 inference -python ../../../run.py --input_text "世界上第二高的山峰是哪座？" \ - --max_output_len=50 \ - --tokenizer_dir baichuan-inc/Baichuan-13B-Chat \ - --engine_dir=./tmp/baichuan_v1_13b/trt_engines/fp16/1-gpu/ - -# With bf16 inference -python ../../../run.py --input_text "世界上第二高的山峰是哪座？" \ - --max_output_len=50 \ - --tokenizer_dir baichuan-inc/Baichuan-13B-Chat \ - --engine_dir=./tmp/baichuan_v1_13b/trt_engines/bf16/1-gpu/ - -# With INT8 weight-only quantization inference -python ../../../run.py --input_text "世界上第二高的山峰是哪座？" \ - --max_output_len=50 \ - --tokenizer_dir=baichuan-inc/Baichuan-13B-Chat \ - --engine_dir=./tmp/baichuan_v1_13b/trt_engines/int8_weight_only/1-gpu/ - -# With INT4 weight-only quantization inference -python ../../../run.py --input_text "世界上第二高的山峰是哪座？" \ - --max_output_len=50 \ - --tokenizer_dir=baichuan-inc/Baichuan-13B-Chat \ - --engine_dir=./tmp/baichuan_v1_13b/trt_engines/int8_weight_only/1-gpu/ - -# With 2-way tensor parallelism inference -mpirun -n 2 --allow-run-as-root \ - python ../run.py --input_text "世界上第二高的山峰是哪座？" \ - --max_output_len=50 \ - --tokenizer_dir=baichuan-inc/Baichuan-13B-Chat \ - --engine_dir=./tmp/baichuan_v1_13b/trt_engines/fp16/2-gpu/ -``` - -### Summarization using the Baichuan model - -```bash -# Run summarization using the Baichuan V1 13B model in FP16. -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 --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 --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/ -``` diff --git a/examples/models/contrib/baichuan/convert_checkpoint.py b/examples/models/contrib/baichuan/convert_checkpoint.py deleted file mode 100644 index f5a296deee87..000000000000 --- a/examples/models/contrib/baichuan/convert_checkpoint.py +++ /dev/null @@ -1,263 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import argparse -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.baichuan.config import BaichuanConfig -from tensorrt_llm.models.baichuan.convert import load_weights_from_gptq -from tensorrt_llm.models.baichuan.model import BaichuanForCausalLM -from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument('--quant_ckpt_path', type=str, default=None) - parser.add_argument('--tp_size', type=int, default=1) - parser.add_argument('--pp_size', type=int, default=1) - parser.add_argument('--model_version', - type=str, - default='v1_13b', - choices=['v1_7b', 'v1_13b', 'v2_7b', 'v2_13b']) - parser.add_argument( - '--dtype', - type=str, - default='auto', - choices=['auto', 'float16', 'bfloat16', 'float32'], - help= - "The data type for the model weights and activations if not quantized. " - "If 'auto', the data type is automatically inferred from the source model; " - "however, if the source dtype is float32, it is converted to float16.") - parser.add_argument('--logits_dtype', - type=str, - default='float32', - choices=['float16', 'float32']) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument( - '--calib_dataset', - type=str, - default='ccdv/cnn_dailymail', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - 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( - "--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( - '--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', 'int4_gptq'], - 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('--group_size', - type=int, - default=128, - help='Group size used in GPTQ/AWQ quantization.') - 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' - ) - args = parser.parse_args() - return args - - -def args_to_quant_config(args: argparse.Namespace) -> QuantConfig: - config = QuantConfig(group_size=args.group_size) - - if args.smoothquant: - config.smoothquant_val = args.smoothquant - if args.per_token and args.per_channel: - config.quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN - elif not args.per_token and not args.per_channel: - config.quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN - elif not args.per_token and args.per_channel: - config.quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN - elif args.per_token and not args.per_channel: - config.quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN - else: - if args.use_weight_only and args.weight_only_precision == 'int8': - config.quant_algo = QuantAlgo.W8A16 - elif args.use_weight_only and args.weight_only_precision == 'int4': - config.quant_algo = QuantAlgo.W4A16 - elif args.use_weight_only and args.weight_only_precision == 'int4_gptq': - config.quant_algo = QuantAlgo.W4A16_GPTQ - - if args.int8_kv_cache: - config.kv_cache_quant_algo = QuantAlgo.INT8 - - if args.use_weight_only and args.weight_only_precision == 'int4_gptq': - config.has_zero_point = True - - return config - - -def convert_and_save_hf(args): - world_size = args.tp_size * args.pp_size - quantization_config = args_to_quant_config(args) - - if args.smoothquant is not None or args.int8_kv_cache: - mapping = Mapping( - world_size=world_size, - tp_size=args.tp_size, - pp_size=args.pp_size, - ) - BaichuanForCausalLM.quantize( - args.model_dir, - args.output_dir, - dtype=args.dtype, - mapping=mapping, - quant_config=quantization_config, - calib_dataset=args.calib_dataset, - model_version=args.model_version, - ) - else: - - def convert_and_save_rank(args, rank): - mapping = Mapping( - world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size, - ) - - model = BaichuanForCausalLM.from_hugging_face( - args.model_dir, - dtype=args.dtype, - mapping=mapping, - quant_config=quantization_config, - model_version=args.model_version, - logits_dtype=args.logits_dtype, - ) - model.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del model - - execute(args.workers, [convert_and_save_rank] * world_size, args) - - -def convert_and_save_gptq(args, rank): - mapping = Mapping(world_size=args.tp_size * args.pp_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - config = BaichuanConfig.from_hugging_face( - args.model_dir, - dtype=args.dtype, - mapping=mapping, - quant_config=args_to_quant_config(args), - model_version=args.model_version, - logits_dtype=args.logits_dtype) - - config.vocab_size = int((config.vocab_size + 63) / 64) * 64 - model = BaichuanForCausalLM(config) - weights = load_weights_from_gptq(config, args.quant_ckpt_path) - model.load(weights) - model.save_checkpoint(args.output_dir, save_config=(rank == 0)) - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - world_size = args.tp_size * args.pp_size - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - if args.use_weight_only and args.weight_only_precision == 'int4_gptq': - assert args.model_dir is not None - assert args.quant_ckpt_path is not None - execute(args.workers, [convert_and_save_gptq] * world_size, args) - else: - assert args.model_dir is not None - assert args.quant_ckpt_path is None, "only gptq weights only needs this option" - convert_and_save_hf(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/baichuan/requirements.txt b/examples/models/contrib/baichuan/requirements.txt deleted file mode 100644 index d234cb9d00fa..000000000000 --- a/examples/models/contrib/baichuan/requirements.txt +++ /dev/null @@ -1,8 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score -sentencepiece>=0.1.99 -cpm-kernels~=1.0.11 -transformers_stream_generator~=0.0.4 diff --git a/examples/models/contrib/bloom/.gitignore b/examples/models/contrib/bloom/.gitignore deleted file mode 100644 index 49fd9fe3a8fa..000000000000 --- a/examples/models/contrib/bloom/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -__pycache__/ -bloom/ diff --git a/examples/models/contrib/bloom/README.md b/examples/models/contrib/bloom/README.md deleted file mode 100644 index a086294d5165..000000000000 --- a/examples/models/contrib/bloom/README.md +++ /dev/null @@ -1,242 +0,0 @@ -# BLOOM - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a BLOOM model in TensorRT LLM on both single GPU, single node multi-GPU and multi-node multi-GPU. - -## Table of Contents - -- [BLOOM](#bloom) - - [Table of Contents](#table-of-contents) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [Build TensorRT engine(s)](#build-tensorrt-engines) - - [INT8 weight only + INT8 KV cache](#int8-weight-only--int8-kv-cache) - - [SmoothQuant](#smoothquant) - - [FP8 Post-Training Quantization](#fp8-post-training-quantization) - - [Run](#run) - -## 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/models/contrib/bloom`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 - * INT8 & INT4 Weight-Only - * INT8 KV CACHE - * Smooth Quant - * Tensor Parallel - * FP8 and FP8 KV cache - -## Usage - -The TensorRT LLM BLOOM example code locates at [examples/models/contrib/bloom](./). It takes HF weights as input, and builds the corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Build TensorRT engine(s) - -Please install required packages first: - -```bash -pip install -r requirements.txt -``` - -Need to prepare the HF BLOOM checkpoint by following the guides here https://huggingface.co/docs/transformers/main/en/model_doc/bloom. - -e.g. To install BLOOM-560M - -```bash -# Setup git-lfs -git lfs install -rm -rf ./bloom/560M -mkdir -p ./bloom/560M && git clone https://huggingface.co/bigscience/bloom-560m ./bloom/560M -``` - -TensorRT LLM BLOOM builds TensorRT engine(s) from HF checkpoint. If no checkpoint directory is specified, TensorRT LLM will build engine(s) with dummy weights. - -Normally `trtllm-build` only requires single GPU, but if you've already got all the GPUs needed for inference, you could enable parallel building to make the engine building process faster by adding `--workers` argument. Please note that currently `workers` feature only supports single node. - -Here're some examples: - -```bash -# Build a single-GPU float16 engine from HF weights. -# Try gemm_plugin to prevent accuracy issue. TODO check this holds for BLOOM - -# Single GPU on BLOOM 560M -python convert_checkpoint.py --model_dir ./bloom/560M/ \ - --dtype float16 \ - --output_dir ./bloom/560M/trt_ckpt/fp16/1-gpu/ -trtllm-build --checkpoint_dir ./bloom/560M/trt_ckpt/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --output_dir ./bloom/560M/trt_engines/fp16/1-gpu/ - -# Build the BLOOM 560M using a single GPU and apply INT8 weight-only quantization. -python convert_checkpoint.py --model_dir ./bloom/560M/ \ - --dtype float16 \ - --use_weight_only \ - --output_dir ./bloom/560M/trt_ckpt/int8_weight_only/1-gpu/ -trtllm-build --checkpoint_dir ./bloom/560M/trt_ckpt/int8_weight_only/1-gpu/ \ - --gemm_plugin float16 \ - --output_dir ./bloom/560M/trt_engines/int8_weight_only/1-gpu/ - -# Use 2-way tensor parallelism on BLOOM 560M -python convert_checkpoint.py --model_dir ./bloom/560M/ \ - --dtype float16 \ - --output_dir ./bloom/560M/trt_ckpt/fp16/2-gpu/ \ - --tp_size 2 -trtllm-build --checkpoint_dir ./bloom/560M/trt_ckpt/fp16/2-gpu/ \ - --gemm_plugin float16 \ - --output_dir ./bloom/560M/trt_engines/fp16/2-gpu/ - -# Use 8-way tensor parallelism on BLOOM 176B -# Currently, TensorRT does not support tensors with more than 2^31-1 elements, -# so we have to shard the embedding table to multi-GPUs. - -# sharding embedding table in the vocab dimension (the lookup plugin is optional) -python convert_checkpoint.py --model_dir ./bloom/176B/ \ - --dtype float16 \ - --output_dir ./bloom/176B/trt_ckpt/fp16/8-gpu/ \ - --tp_size 8 \ - --use_parallel_embedding \ - --embedding_sharding_dim 0 -trtllm-build --checkpoint_dir ./bloom/176B/trt_ckpt/fp16/8-gpu/ \ - --gemm_plugin float16 \ - --output_dir ./bloom/176B/trt_engines/fp16/8-gpu/ \ - --workers 2 - -# sharding embedding table in the hidden dimension -python convert_checkpoint.py --model_dir ./bloom/176B/ \ - --dtype float16 \ - --output_dir ./bloom/176B/trt_ckpt/fp16/8-gpu/ \ - --tp_size 8 \ - --use_parallel_embedding \ - --embedding_sharding_dim 1 -trtllm-build --checkpoint_dir ./bloom/176B/trt_ckpt/fp16/8-gpu/ \ - --gemm_plugin float16 \ - --output_dir ./bloom/176B/trt_engines/fp16/8-gpu/ \ - --workers 2 - -# share embedding table between embedding() and lm_head() layers -# To reduce the generated engine size, we can turn off gemm plugin and shard the embedding table in the vocab dimension. -python convert_checkpoint.py --model_dir ./bloom/176B/ \ - --dtype float16 \ - --output_dir ./bloom/176B/trt_ckpt/fp16/8-gpu/ \ - --tp_size 8 \ - --use_parallel_embedding \ - --embedding_sharding_dim 0 -trtllm-build --checkpoint_dir ./bloom/176B/trt_ckpt/fp16/8-gpu/ \ - --output_dir ./bloom/176B/trt_engines/fp16/8-gpu/ \ - --workers 2 -``` - -#### INT8 weight only + INT8 KV cache - - -For INT8 KV cache, [`convert_checkpoint.py`](./convert_checkpoint.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 convert_checkpoint.py --model_dir ./bloom/560m/ \ - --dtype float16 \ - --int8_kv_cache \ - --use_weight_only --output_dir ./bloom/560m/trt_ckpt/int8/1-gpu/ -trtllm-build --checkpoint_dir ./bloom/560m/trt_ckpt/int8/1-gpu/ \ - --gemm_plugin float16 \ - --output_dir ./bloom/560m/trt_engines/int8/1-gpu/ \ -``` - - -#### SmoothQuant - -Unlike the FP16 build where the HF weights are processed and loaded into the TensorRT LLM directly, the SmoothQuant needs to load INT8 weights which should be pre-processed before building an engine. - -Example: -```bash -python convert_checkpoint.py --model_dir bloom/560M/ --output_dir tllm_checkpoint_1gpu --smoothquant 0.5 --per_token --per_channel -``` - -[`convert_checkpoint.py`](./convert_checkpoint.py) add new options for the support of INT8 inference of SmoothQuant models. - -`--smoothquant` is the starting point of INT8 inference. By default, it -will run the model in the _per-tensor_ mode. - -Then, you can add any combination of `--per-token` and `--per-channel` to get the corresponding behaviors. - - - -```bash -# Build model for SmoothQuant with below command. - -trtllm-build --checkpoint_dir tllm_checkpoint_1gpu --output_dir ./engine_outputs -``` -Note that GPT attention plugin is required to be enabled for SmoothQuant for now. - - -Note we use `--bin_model_dir` instead of `--model_dir` since SmoothQuant model needs INT8 weights and various scales from the binary files. - -#### FP8 Post-Training Quantization - -``` -# Quantize HF Bloom 3B into FP8 and export trtllm checkpoint -python ../../../quantization/quantize.py --model_dir /home/scratch.trt_llm_data_ci/llm-models/bloom-3b \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir /tmp/bloom/3b/trt_ckpts/fp8/1-gpu/ \ - --calib_size 512 \ - --tp_size 1 - -trtllm-build --checkpoint_dir /tmp/bloom/3b/trt_ckpts/fp8/1-gpu/ \ - --output_dir /tmp/bloom/3b/trt_engines/fp8/1-gpu/ \ - --gemm_plugin float16 \ - --use_fp8_context_fmha enable \ - --workers 1 -``` - -### Run - -```bash -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_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_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_dir ./bloom/176B/ \ - --data_type fp16 \ - --engine_dir ./bloom/176B/trt_engines/fp16/8-gpu/ - -python ../../../summarize.py --test_trt_llm \ - --hf_model_dir /home/scratch.trt_llm_data_ci/llm-models/bloom-3b \ - --data_type fp16 \ - --engine_dir /tmp/bloom/3b/trt_engines/fp8/1-gpu/ -``` diff --git a/examples/models/contrib/bloom/convert_checkpoint.py b/examples/models/contrib/bloom/convert_checkpoint.py deleted file mode 100644 index 8ab16f2b23bf..000000000000 --- a/examples/models/contrib/bloom/convert_checkpoint.py +++ /dev/null @@ -1,962 +0,0 @@ -import argparse -import functools -import json -import os -import time -from collections import defaultdict -from pathlib import Path -from typing import Any, Dict, Iterable, Optional, Union - -import numpy as np -import safetensors -import torch -import torch.nn as nn -from tqdm import tqdm -from transformers import BloomConfig, BloomForCausalLM, BloomTokenizerFast -from transformers.models.bloom.modeling_bloom import BloomBlock -from transformers.pytorch_utils import Conv1D - -# isort: off -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm import logger -from tensorrt_llm.quantization import QuantAlgo, QuantMode -from tensorrt_llm.models.convert_utils import iterate_shard_files, load_state_dict, \ - load_calib_dataset, split_matrix_tp, get_weight_and_bias, split, smooth_gemm, \ - generate_int8,get_weight -# isort: on - - -@torch.no_grad() -def capture_activation_range(model, - tokenizer, - dataset, - num_samples=512, - seq_len=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))) - - for i in tqdm(range(num_samples), desc="calibrating model"): - input_ids = tokenizer(dataset[i], - return_tensors="pt", - max_length=seq_len, - truncation=True).input_ids.to(device) - model(input_ids) - - for h in hooks: - h.remove() - - return act_scales - - -def reorder_torch_qkv_weight_or_bias(v, model, is_bias=False): - """ Reorder the qkv weight. - - Note that the shape of the fused QKV weights in HF is different from the - shape that TRT-LLM requires. - HF: (num_heads x 3 x head_dim, hidden_size) - TRT-LLM: (3 x num_heads x head_dim, hidden_size) - This is unlike to the other models in HF e.g. GPT where they have the - same shape with TRT-LLM, i.e., (3 x num_heads x head_dim, hidden_size). We reshape the qkv - weight: (3 x num_heads x head_dim, hidden). - bias : (3 x num_heads x head_dim). - """ - - n_head = model.transformer.num_heads - hidden_size = model.transformer.embed_dim - head_dim = hidden_size // n_head - - # (3 x hidden, ...) view as (num_heads, 3, head_dim, ...) - v = v.reshape(n_head, 3, head_dim, -1) - # permute to (3, num_heads, head_dim, ...) - v = v.permute((1, 0, 2, 3)) - # final shape: weight=(3 x hidden, hidden) or bias=(3 x hidden) - if is_bias: - return v.reshape(3 * hidden_size) - return v.reshape(3 * hidden_size, hidden_size) - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=Path, default=None) - 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('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - 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_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - parser.add_argument('--output_dir', - type=Path, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--calib_dataset', - type=str, - default='lambada', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - 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( - '--per_channel', - action="store_true", - default=False, - 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', - action="store_true", - default=False, - 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( - '--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( - '--workers', - type=int, - default=1, - help='The number of workers to convert checkpoint in parallel') - parser.add_argument('--log_level', type=str, default='info') - args = parser.parse_args() - - return args - - -def reorder_qkv_weight_or_bias(v, n_head, n_hidden, is_bias=False): - """ Reorder the qkv weight. - - Note that the shape of the fused QKV weights in HF is different from the - shape that TRT-LLM requires. - HF: (num_heads x 3 x head_dim, hidden_size) - TRT-LLM: (3 x num_heads x head_dim, hidden_size) - This is unlike to the other models in HF e.g. GPT where they have the - same shape with TRT-LLM, i.e., (3 x num_heads x head_dim, hidden_size). Also, - to split across attention heads in tensor parallel, we reshape the qkv - weight: (3, num_heads x head_dim, hidden). - bias : (3, num_heads x head_dim). - """ - - head_dim = n_hidden // n_head - - # (3 x hidden, ...) view as (num_heads, 3, head_dim, ...) - v = v.reshape(n_head, 3, head_dim, -1) - # permute to (3, num_heads, head_dim, ...) - v = v.transpose(0, 1) - # final shape: weight=(3, hidden, hidden) or bias=(3, hidden) - if is_bias: - return v.reshape(3, n_hidden) - return v.reshape(3, n_hidden, n_hidden) - - -def split_qkv_tp(v, n_head, n_hidden, tensor_parallel, rank): - """ - Splits the QKV matrix according to tensor parallelism - """ - v = reorder_qkv_weight_or_bias(v, n_head, n_hidden, is_bias=False) - split_v = split(v, tensor_parallel, rank, dim=1) - split_v = split_v.reshape(3 * (n_hidden // tensor_parallel), n_hidden) - return split_v.contiguous() - - -def split_qkv_bias_tp(v, n_head, n_hidden, tensor_parallel, rank): - """ - Splits the QKV bias according to tensor parallelism - """ - v = reorder_qkv_weight_or_bias(v, n_head, n_hidden, is_bias=True) - split_v = split(v, tensor_parallel, rank, dim=1) - split_v = split_v.reshape(3 * (n_hidden // tensor_parallel)) - return split_v.contiguous() - - -def get_tllm_linear_weight(weight, - prefix, - bias=None, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8): - results = {} - if use_weight_only: - v = weight.cpu().t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[prefix + 'weight'] = processed_torch_weights - results[prefix + 'per_channel_scale'] = torch_weight_scales - else: - results[prefix + 'weight'] = weight.contiguous() - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def add_tllm_weight( - weights: Dict[str, torch.Tensor], - name: str, - param: torch.Tensor, - quant_mode: QuantMode = QuantMode(0), -): - assert name not in weights, f'{name} is already added.' - - if name.endswith('.weight') and quant_mode.is_weight_only(): - if quant_mode.is_int8_weight_only(): - quant_dtype = torch.int8 - elif quant_mode.is_int4_weight_only(): - quant_dtype = torch.quint4x2 - else: - raise ValueError( - f'Invalid configuration, got quant_mode={quant_mode}') - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - param.t().contiguous(), quant_dtype) - weights[name] = processed_torch_weights - scale_name = name.replace('.weight', '.per_channel_scale') - weights[scale_name] = torch_weight_scales - else: - weights[name] = param.contiguous() - - -@torch.no_grad() -def smooth_bloom_model(model, scales, alpha, bloom_qkv_param, bloom_smoother): - # Smooth the activation and weights with smoother = $\diag{s}$ - for name, module in model.named_modules(): - if not isinstance(module, BloomBlock): - continue - - # reorder qkv weight/bias and scales - param = module.self_attention.query_key_value.weight - param = reorder_torch_qkv_weight_or_bias(param, model, is_bias=False) - - layer_name = name + ".self_attention.query_key_value" - act_range_qkv = scales.get(layer_name) - # (n_head x 3 x head_dim) -> (3 x n_head x head_dim) - act_range_qkv['w'] = reorder_torch_qkv_weight_or_bias( - act_range_qkv['w'], model, is_bias=True) - act_range_qkv['y'] = reorder_torch_qkv_weight_or_bias( - act_range_qkv['y'], model, is_bias=True) - scales[layer_name] = act_range_qkv - - # qkv_proj - smoother = smooth_gemm(param, scales[layer_name]["x"], - module.input_layernorm.weight, - module.input_layernorm.bias, alpha) - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = param.abs().max(dim=1)[0] - bloom_qkv_param[layer_name] = param - - # dense - # enabled for better accuracy with perf overhead of quantization - layer_name = name + ".self_attention.dense" - smoother = smooth_gemm(module.self_attention.dense.weight, - scales[layer_name]["x"], None, None, alpha) - bloom_smoother[layer_name] = smoother - - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.self_attention.dense.weight.abs().max( - dim=1)[0] - - # fc1 - layer_name = name + ".mlp.dense_h_to_4h" - smoother = smooth_gemm(module.mlp.dense_h_to_4h.weight, - scales[layer_name]["x"], - module.post_attention_layernorm.weight, - module.post_attention_layernorm.bias, alpha) - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.mlp.dense_h_to_4h.weight.abs().max( - dim=1)[0] - - # fc2 - # enabled for better accuracy with perf overhead of quantization - layer_name = name + ".mlp.dense_4h_to_h" - smoother = smooth_gemm(module.mlp.dense_4h_to_h.weight, - scales[layer_name]["x"], None, None, alpha) - bloom_smoother[layer_name] = smoother - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.mlp.dense_4h_to_h.weight.abs().max( - dim=1)[0] - - -def get_tllm_linear_sq_weight( - vals, - prefix, - shape, - tensor_parallel, - is_qkv=False, - per_token=False, - per_channel=False, - last_prefix=None, - bias=None, - smoother_value=None, - smoother_shape=None, - rank=0, - cat_dim=0, -): - results = {} - - col_shape = shape if (is_qkv or per_channel) else [1, 1] - - if per_token: - - original_weights = vals["weight.int8.col"] - cur_weights = np.split(original_weights, tensor_parallel, - axis=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + 'weight'] = cur_weights.t().contiguous() - if smoother_value is None: - results[last_prefix] = torch.from_numpy( - np.array([1.0], dtype=np.float32)) - - if smoother_value is None: - cur_per_channel_value = np.split(vals["scale_w_quant_orig.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_w_quant_orig.col"] - results[prefix + 'per_channel_scale'] = cur_per_channel_value.reshape( - col_shape).contiguous() - else: - original_weights = vals["weight.int8"] - cur_weights = np.split(original_weights, tensor_parallel, - axis=cat_dim)[rank] - - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + 'weight'] = cur_weights.t().contiguous() - - cur_per_channel_value = vals["scale_y_accum_quant"] - - results[prefix + 'per_channel_scale'] = cur_per_channel_value.reshape( - col_shape).contiguous() - - results[last_prefix] = vals['scale_x_orig_quant'].contiguous() - - results[prefix + 'act_scale'] = vals["scale_y_quant_orig"].contiguous() - - if smoother_value is not None: - cur_smoother_value = np.split(smoother_value, - tensor_parallel, - axis=cat_dim)[rank] - results[prefix + 'smoother'] = cur_smoother_value.reshape( - smoother_shape).contiguous().to(torch.float32) - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def convert_hf_bloom(hf_bloom, - rank=0, - tensor_parallel=1, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8, - use_smooth_quant=False, - bloom_qkv_param={}, - act_range=None, - smoother=None, - per_channel=False, - per_token=False, - int8_kv_cache=False): - - weights = {} - tik = time.time() - - model_params = dict(hf_bloom.named_parameters()) - dtype = getattr(torch, dtype) - num_attention_heads = hf_bloom.config.n_head - hidden_size = hf_bloom.config.hidden_size - - for l in range(hf_bloom.config.num_hidden_layers): - prefix = f'transformer.h.{l}.' - tllm_prex = f'transformer.layers.{l}.' - qkv_weight, qkv_bias = get_weight_and_bias( - model_params, prefix + 'self_attention.query_key_value', dtype) - split_v = split_qkv_tp(qkv_weight, num_attention_heads, hidden_size, - tensor_parallel, rank) - bias = split_qkv_bias_tp(qkv_bias, num_attention_heads, hidden_size, - tensor_parallel, rank) - - if use_smooth_quant: - qkv_weight = bloom_qkv_param[prefix + - 'self_attention.query_key_value'].t() - - qkv_weight = qkv_weight.reshape(hidden_size, 3, hidden_size) - int8_weights = generate_int8( - qkv_weight, - act_range.get( - (tllm_prex + 'self_attention.query_key_value').replace( - ".layers.", ".h.")), - is_qkv=True) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'attention.qkv.', - [1, 3 * hidden_size // tensor_parallel], - tensor_parallel, - is_qkv=True, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'input_layernorm.scale_to_int', - bias=bias, - smoother_value=None, - smoother_shape=None, - rank=rank, - cat_dim=-1)) - else: - split_v = split_qkv_tp(qkv_weight, num_attention_heads, hidden_size, - tensor_parallel, rank) - bias = split_qkv_bias_tp(qkv_bias, num_attention_heads, hidden_size, - tensor_parallel, rank) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.qkv.', - bias, use_weight_only, - plugin_weight_only_quant_type)) - if int8_kv_cache: - qkv_weight = qkv_weight.reshape(hidden_size, 3, hidden_size) - - int8_weights = generate_int8( - qkv_weight, - act_range.get( - (tllm_prex + 'self_attention.query_key_value').replace( - ".layers.", ".h.")), - is_qkv=True) - - kv_cache_weights = {} - - kv_cache_weights[ - tllm_prex + - 'attention.kv_cache_scaling_factor'] = torch.from_numpy( - np.array([int8_weights['scale_y_quant_orig']], - dtype=np.float32)).contiguous() - weights.update(kv_cache_weights) - attn_dense_weight, attn_dense_bias = get_weight_and_bias( - model_params, prefix + 'self_attention.dense', dtype) - if use_smooth_quant: - attn_dense_weight = attn_dense_weight.t() - int8_weights = generate_int8( - attn_dense_weight, - act_range.get((tllm_prex + 'self_attention.dense').replace( - ".layers.", ".h."))) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'attention.dense.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'attention.quantization_scaling_factor', - bias=attn_dense_bias, - smoother_value=smoother[(tllm_prex + - 'self_attention.dense').replace( - ".layers.", ".h.")], - smoother_shape=[1, hidden_size // tensor_parallel], - rank=rank, - cat_dim=0)) - else: - - split_v = split_matrix_tp(attn_dense_weight, - tensor_parallel, - rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.dense.', - attn_dense_bias, use_weight_only, - plugin_weight_only_quant_type)) - mlp_fc_weight, mlp_fc_bias = get_weight_and_bias( - model_params, prefix + 'mlp.dense_h_to_4h', dtype) - bias = split_matrix_tp(mlp_fc_bias, tensor_parallel, rank, dim=0) - - if use_smooth_quant: - mlp_fc_weight = mlp_fc_weight.t() - int8_weights = generate_int8( - mlp_fc_weight, - act_range.get((tllm_prex + 'mlp.dense_h_to_4h').replace( - ".layers.", ".h."))) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.fc.', - [1, 4 * hidden_size // tensor_parallel], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - bias=bias, - smoother_value=None, - smoother_shape=None, - rank=rank, - cat_dim=-1)) - else: - split_v = split_matrix_tp(mlp_fc_weight, - tensor_parallel, - rank, - dim=0) - bias = split_matrix_tp(mlp_fc_bias, tensor_parallel, rank, dim=0) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.fc.', bias, - use_weight_only, - plugin_weight_only_quant_type)) - mlp_proj_weight, mlp_proj_bias = get_weight_and_bias( - model_params, prefix + 'mlp.dense_4h_to_h', dtype) - if use_smooth_quant: - mlp_proj_weight = mlp_proj_weight.t() - int8_weights = generate_int8( - mlp_proj_weight, - act_range.get((tllm_prex + 'mlp.dense_4h_to_h').replace( - ".layers.", ".h."))) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.proj.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'mlp.quantization_scaling_factor', - bias=mlp_proj_bias, - smoother_value=smoother[(tllm_prex + - 'mlp.dense_4h_to_h').replace( - ".layers.", ".h.")], - smoother_shape=[1, 4 * hidden_size // tensor_parallel], - rank=rank, - cat_dim=0)) - else: - split_v = split_matrix_tp(mlp_proj_weight, - tensor_parallel, - rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.proj.', - mlp_proj_bias, use_weight_only, - plugin_weight_only_quant_type)) - - # Layer norms do not use tensor parallelism - input_ln_weight, input_ln_bias = get_weight_and_bias( - model_params, prefix + 'input_layernorm', dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - weights[tllm_prex + 'input_layernorm.bias'] = input_ln_bias - - post_ln_weight, post_ln_bias = get_weight_and_bias( - model_params, prefix + 'post_attention_layernorm', dtype) - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - weights[tllm_prex + 'post_layernorm.bias'] = post_ln_bias - - embed_w = get_weight(model_params, 'transformer.word_embeddings', dtype) - weights['lm_head.weight'] = split_matrix_tp(embed_w.clone(), - tensor_parallel, - rank, - dim=0) - - if not use_parallel_embedding: - weights['transformer.vocab_embedding.weight'] = embed_w - else: - assert hf_bloom.config.vocab_size % tensor_parallel == 0 - weights['transformer.vocab_embedding.weight'] = split_matrix_tp( - embed_w, tensor_parallel, rank, dim=sharding_dim) - - embed_f_w, embed_f_b = get_weight_and_bias( - model_params, 'transformer.word_embeddings_layernorm', dtype) - weights['transformer.ln_embed.weight'] = embed_f_w - weights['transformer.ln_embed.bias'] = embed_f_b - - ln_f_w, ln_f_b = get_weight_and_bias(model_params, 'transformer.ln_f', - dtype) - weights['transformer.ln_f.weight'] = ln_f_w - weights['transformer.ln_f.bias'] = ln_f_b - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def rename_hf_to_tllm(name: str): - """ Rename a HF parameter name by the corresponding TRT-LLM style name. """ - if 'word_embeddings_layernorm.' in name: - name = name.replace('word_embeddings_layernorm', 'ln_embed') - if not name.startswith('transformer.'): - name = f'transformer.{name}' - elif 'word_embeddings.' in name: - name = name.replace('word_embeddings', 'vocab_embedding') - if name.startswith(('ln_embed.', 'vocab_embedding.', 'ln_f.')): - name = f'transformer.{name}' - - # Parameter names in layers - if name.startswith(('transformer.h.', 'h.')): - import re - name = re.sub(r'^(transformer.h.|h.)', 'transformer.layers.', name, 1) - if 'post_attention_layernorm' in name: - name = name.replace('post_attention_layernorm', 'post_layernorm') - elif 'self_attention.query_key_value' in name: - name = name.replace('self_attention.query_key_value', 'attention.qkv') - elif 'self_attention.dense' in name: - name = name.replace('self_attention.dense', 'attention.dense') - elif 'mlp.dense_h_to_4h' in name: - name = name.replace('mlp.dense_h_to_4h', 'mlp.fc') - elif 'mlp.dense_4h_to_h' in name: - name = name.replace('mlp.dense_4h_to_h', 'mlp.proj') - return name - - -def contain_any(name: str, words: Iterable[str]): - for word in words: - if word in name: - return True - return False - - -def convert_from_hf_checkpoint( - model_dir: Union[str, Path], - rank=0, - tensor_parallel=1, - dtype: Union[str, torch.dtype] = torch.float32, - use_parallel_embedding: bool = False, - sharding_dim: int = 0, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8, - use_smooth_quant: bool = False, - bloom_qkv_param: Optional[Dict] = None, - act_range: Optional[Any] = None, - smoother: Optional[Any] = None, - per_channel: bool = False, - per_token: bool = False, - int8_kv_cache: bool = False, -): - logger.info('Loading weights from HF BLOOM...') - tik = time.time() - - weights = {} - hf_config = BloomConfig.from_pretrained(model_dir) - num_heads = hf_config.n_head - hidden_size = hf_config.hidden_size - if isinstance(dtype, str): - dtype = tensorrt_llm.str_dtype_to_torch(dtype) - tp_rank = rank - tp_size = tensor_parallel - - if use_smooth_quant: - quant_mode = QuantMode.use_smooth_quant(per_token, per_channel) - elif use_weight_only: - quant_mode = QuantMode.from_description( - quantize_weights=True, - quantize_activations=False, - per_token=False, - per_channel=False, - use_int8_kv_cache=int8_kv_cache, - use_int4_weights=plugin_weight_only_quant_type == torch.quint4x2) - else: - quant_mode = QuantMode(0) - - def is_bias(_name): - return 'bias' in _name - - for model_file in iterate_shard_files(model_dir, tp_rank): - logger.debug(f'Loading file {str(model_file)}...') - model_params = load_state_dict(model_file, dtype=dtype) - for name, param in model_params.items(): - logger.debug(f'Converting weight {name}...') - tllm_name = rename_hf_to_tllm(name) - param = param.detach().cpu() - - # TODO: Support SmmothQuant. - - if 'self_attention.query_key_value' in name: - if not is_bias(name): - param = split_qkv_tp(param, num_heads, hidden_size, tp_size, - tp_rank) - # TODO: Add KV scalers when quantizing KV cache. - else: - param = split_qkv_bias_tp(param, num_heads, hidden_size, - tp_size, tp_rank) - add_tllm_weight(weights, tllm_name, param, quant_mode) - elif 'self_attention.dense' in name: - if not is_bias(name): - param = split_matrix_tp(param, tp_size, tp_rank, dim=1) - add_tllm_weight(weights, tllm_name, param, quant_mode) - elif 'mlp.dense_h_to_4h' in name: - if not is_bias(name): - param = split_matrix_tp(param, tp_size, tp_rank, dim=0) - else: - param = split_matrix_tp(param, tp_size, tp_rank, dim=0) - add_tllm_weight(weights, tllm_name, param, quant_mode) - elif 'mlp.dense_4h_to_h' in name: - if not is_bias(name): - param = split_matrix_tp(param, tp_size, tp_rank, dim=1) - add_tllm_weight(weights, tllm_name, param, quant_mode) - elif 'word_embeddings.' in name: - # TODO: safetensor doesn't allow to save a shared tensor. - # Currently, we clone the weight but to save the disk, it - # would be better to skip saving lm_head weights and - # handle it at the loading phase. - lm_head = split_matrix_tp(param, tp_size, tp_rank, dim=0) - weights['lm_head.weight'] = lm_head.clone() - - if not use_parallel_embedding: - weights[tllm_name] = param - else: - assert hf_config.vocab_size % tp_size == 0 - weights[tllm_name] = split_matrix_tp(param, - tp_size, - tp_rank, - dim=sharding_dim) - elif contain_any(name, - ('input_layernorm', 'post_attention_layernorm', - 'word_embeddings_layernorm.', 'ln_f.')): - weights[tllm_name] = 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}') - return weights - - -def do_convert_from_ckpt(args): - return (args.model_dir.exists() and args.smoothquant is None - and not args.use_weight_only and not args.int8_kv_cache) - - -def convert(worker_rank, world_size, args, convert_args): - convert_from_ckpt = do_convert_from_ckpt(args) - for rank in range(worker_rank, world_size, args.workers): - if convert_from_ckpt: - weights = convert_from_hf_checkpoint(rank=rank, **convert_args) - else: - weights = convert_hf_bloom(rank=rank, **convert_args) - safetensors.torch.save_file(weights, - args.output_dir / f'rank{rank}.safetensors') - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - # TODO(qijun): Currently, the convert script depends on a torch op: - # torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix, - # which is included in tensorrt_llm Python package. Otherwise, the convert - # script does not need to import tensorrt_llm. Will remove it after reimplementing - # the op with PyTorch. - print(tensorrt_llm.__version__) - - args = parse_arguments() - world_size = args.tp_size * args.pp_size - assert args.pp_size == 1, "Pipeline parallelism is not supported." - - logger.set_level(args.log_level) - tik = time.time() - - args.output_dir.mkdir(exist_ok=True, parents=True) - - quant_algo = None - plugin_weight_only_quant_type = None - if args.use_weight_only and args.weight_only_precision == 'int8': - plugin_weight_only_quant_type = torch.int8 - quant_algo = QuantAlgo.W8A16 - elif args.use_weight_only and args.weight_only_precision == 'int4': - plugin_weight_only_quant_type = torch.quint4x2 - quant_algo = QuantAlgo.W4A16 - elif args.smoothquant: - if args.per_channel and args.per_token: - quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN - elif args.per_channel and not args.per_token: - quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN - elif not args.per_channel and args.per_token: - quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN - else: - quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN - - kv_cache_quant_algo = None - if args.int8_kv_cache: - kv_cache_quant_algo = QuantAlgo.INT8 - - hf_config = BloomConfig.from_pretrained(args.model_dir) - config = { - 'architecture': hf_config.architectures[0], - 'dtype': args.dtype, - 'logits_dtype': 'float32', - 'num_hidden_layers': hf_config.num_hidden_layers, - 'num_attention_heads': hf_config.num_attention_heads, - 'hidden_size': hf_config.hidden_size, - 'norm_epsilon': hf_config.layer_norm_epsilon, - 'vocab_size': hf_config.vocab_size, - 'position_embedding_type': 'alibi', - 'hidden_act': 'gelu', - 'intermediate_size': hf_config.hidden_size * 4, - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - } - - with (args.output_dir / 'config.json').open('w') as f: - json.dump(config, f, indent=4) - - # TODO: convert_from_hf_checkpoint is memory efficient but has not - # supported quantization yet. Will enable once implemented. - convert_from_ckpt = do_convert_from_ckpt(args) - if not convert_from_ckpt: - logger.info(f'Convert by using model') - hf_bloom = BloomForCausalLM.from_pretrained(args.model_dir, - dtype="auto", - device_map="auto", - trust_remote_code=True) - else: - logger.info(f'Convert by using checkpoint') - hf_bloom = None - - act_range = {} - bloom_qkv_param = {} - bloom_smoother = {} - - if args.smoothquant is not None or args.int8_kv_cache: - os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get( - "TOKENIZERS_PARALLELISM", "false") - tokenizer = BloomTokenizerFast.from_pretrained(args.model_dir) - dataset = load_calib_dataset(args.calib_dataset) - - act_range = capture_activation_range(hf_bloom, tokenizer, dataset) - if args.smoothquant is not None: - smooth_bloom_model(hf_bloom, act_range, args.smoothquant, - bloom_qkv_param, bloom_smoother) - - convert_args = dict( - tensor_parallel=args.tp_size, - dtype=args.dtype, - use_weight_only=args.use_weight_only, - plugin_weight_only_quant_type=plugin_weight_only_quant_type, - use_parallel_embedding=args.use_parallel_embedding, - sharding_dim=args.embedding_sharding_dim, - use_smooth_quant=args.smoothquant, - act_range=act_range, - bloom_qkv_param=bloom_qkv_param, - smoother=bloom_smoother, - per_channel=args.per_channel, - per_token=args.per_token, - int8_kv_cache=args.int8_kv_cache, - ) - if convert_from_ckpt: - convert_args['model_dir'] = args.model_dir - else: - convert_args['hf_bloom'] = hf_bloom - - if args.workers == 1: - convert(0, world_size, args, convert_args) - else: - if args.workers > world_size: - args.workers = world_size - logger.info(f'Convert checkpoint using {args.workers} workers.') - import torch.multiprocessing as mp - mp.spawn(convert, - nprocs=args.workers, - args=(world_size, args, convert_args)) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/bloom/requirements.txt b/examples/models/contrib/bloom/requirements.txt deleted file mode 100644 index 88232baef811..000000000000 --- a/examples/models/contrib/bloom/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score -sentencepiece>=0.1.99 diff --git a/examples/models/contrib/cogvlm/convert_checkpoint.py b/examples/models/contrib/cogvlm/convert_checkpoint.py deleted file mode 100644 index 26a56de58438..000000000000 --- a/examples/models/contrib/cogvlm/convert_checkpoint.py +++ /dev/null @@ -1,501 +0,0 @@ -import argparse -import copy -import json -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed -from pathlib import Path - -import safetensors -import torch -from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import PretrainedConfig -from tensorrt_llm.models.cogvlm.convert import convert_hf_cogvlm -from tensorrt_llm.models.convert_utils import load_calib_dataset -from tensorrt_llm.models.llama.convert import (capture_activation_range, - load_weights_from_gptq, - load_weights_from_hf_by_shard, - load_weights_from_meta_ckpt, - smooth_llama_model) - -try: - from transformers import LlavaConfig, LlavaForConditionalGeneration -except ImportError: - pass - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument('--meta_ckpt_dir', type=str, default=None) - - 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('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument('--vocab_size', type=int, default=32000) - parser.add_argument('--n_positions', type=int, default=2048) - parser.add_argument('--n_layer', type=int, default=32) - parser.add_argument('--n_head', type=int, default=32) - parser.add_argument('--n_kv_head', type=int, default=None) - parser.add_argument('--n_embd', type=int, default=4096) - parser.add_argument('--inter_size', type=int, default=11008) - parser.add_argument('--rms_norm_eps', type=float, default=1e-06) - - 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( - '--disable_weight_only_quant_plugin', - default=False, - action="store_true", - help= - 'By default, using plugin implementation for weight quantization. Enabling disable_weight_only_quant_plugin flag will use ootb implementation instead of plugin.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8', 'int4', 'int4_gptq'], - 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( - '--calib_dataset', - type=str, - default='ccdv/cnn_dailymail', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - 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( - '--per_channel', - action="store_true", - default=False, - 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', - action="store_true", - default=False, - 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( - '--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( - '--quant_ckpt_path', - type=str, - default=None, - help='Path of a quantized model checkpoint in .npz format') - - 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( - '--enable_fp8', - default=False, - action='store_true', - help='Use FP8 Linear layer for Attention QKV/Dense and MLP.') - parser.add_argument( - '--fp8_kv_cache', - default=False, - action="store_true", - help='By default, we use dtype for KV cache. fp8_kv_cache chooses int8 ' - 'quantization for KV') - parser.add_argument('--load_by_shard', - action='store_true', - help='Load a pretrained model shard-by-shard.') - parser.add_argument('--hidden_act', type=str, default='silu') - - parser.add_argument('--rotary_base', type=float, default=10000.0) - parser.add_argument('--rotary_scaling', nargs=2, type=str, default=None) - - parser.add_argument('--group_size', - type=int, - default=128, - help='Group size used in GPTQ/AWQ quantization.') - - parser.add_argument("--storage-type", - "-t", - 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( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - parser.add_argument('--use_prompt_tuning', - action="store_true", - default=False) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument('--enable_pos_shift', - default=False, - action='store_true', - help='Enable position shift for streamingllm method') - parser.add_argument( - '--dense_context_fmha', - default=False, - action='store_true', - help= - 'Enable dense fmha in context phase, otherwise sliding window attention.' - 'If dense_context_fmha=False, the sliding window size is the max attention window size.' - ) - args = parser.parse_args() - return args - - -def update_quantization_from_args(config: dict, args: argparse.Namespace): - '''update the given config dict in-place based on the command line args - ''' - if args.use_weight_only: - if args.weight_only_precision == 'int8': - config['quantization']['quant_algo'] = 'W8A16' - elif args.weight_only_precision == 'int4': - config['quantization']['quant_algo'] = 'W4A16' - elif args.smoothquant: - config['quantization']['sq_use_plugin'] = True - if args.per_channel: - if args.per_token: - config['quantization'][ - 'quant_algo'] = 'W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN' - else: - config['quantization'][ - 'quant_algo'] = 'W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN' - else: - if args.per_token: - config['quantization'][ - 'quant_algo'] = 'W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN' - else: - config['quantization'][ - 'quant_algo'] = 'W8A8_SQ_PER_TENSOR_PLUGIN' - - if args.int8_kv_cache: - config['quantization']['kv_cache_quant_algo'] = 'INT8' - - if args.weight_only_precision == 'int4_gptq': - config['quantization'].update({ - "group_size": args.group_size, - "has_zero_point": True, - "pre_quant_scale": False, - 'quant_algo': 'W4A16_GPTQ' - }) - - -def create_config_from_args(args: argparse.Namespace): - config = { - 'architecture': args.architecture, - 'dtype': args.dtype, - 'logits_dtype': 'float32', - 'num_hidden_layers': args.n_layer, - 'num_attention_heads': args.n_head, - 'hidden_size': args.n_embd, - 'intermediate_size': args.inter_size, - 'num_key_value_heads': args.n_kv_head, - 'vocab_size': args.vocab_size, - 'position_embedding_type': 'learned_absolute', - 'max_position_embeddings': args.n_positions, - 'hidden_act': args.hidden_act, - 'rotary_base': args.rotary_base, - 'rotary_scaling': args.rotary_scaling, - 'norm_epsilon': args.rms_norm_eps, - 'quantization': { - 'quant_algo': None, - 'kv_cache_quant_algo': None, - }, - 'mapping': { - 'world_size': args.tp_size * args.pp_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'use_prompt_tuning': args.use_prompt_tuning, - 'enable_pos_shift': args.enable_pos_shift, - 'dense_context_fmha': args.dense_context_fmha, - } - update_quantization_from_args(config, args) - return config - - -def smooth_quant(model, args): - assert model is not None - act_range = {} - llama_qkv_para = {} - # smoother for inputs of self_attn.o_proj and mlp.down_proj - llama_smoother = {} - - 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 slow." - ) - tokenizer = AutoTokenizer.from_pretrained(args.model_dir, - trust_remote_code=True, - use_fast=False, - padding_side='left') - dataset = load_calib_dataset(args.calib_dataset, - cache_dir=args.dataset_cache_dir) - - act_range = capture_activation_range(model, tokenizer, dataset) - if args.smoothquant is not None: - smooth_llama_model(model, act_range, args.smoothquant, llama_qkv_para, - llama_smoother) - return act_range, llama_qkv_para, llama_smoother - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - # TODO(qijun): Currently, the convert script depends on a torch op: - # torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix, - # which is included in tensorrt_llm Python package. Otherwise, the convert - # script does not need to import tensorrt_llm. Will remove it after reimplementing - # the op with PyTorch. - logger.info(tensorrt_llm.__version__) - args = parse_arguments() - if args.model_dir is None and args.meta_ckpt_dir is None: - raise AssertionError( - "One of the model_dir or meta_ckpt_dir must be specified to generate the checkpoint" - ) - - world_size = args.tp_size * args.pp_size - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - hf_config = None - if args.model_dir is not None: - hf_config = AutoConfig.from_pretrained(args.model_dir, - trust_remote_code=True) - if hf_config.model_type == "llava": - # LLaVA = Vision model + Llama LLM - # We load a llava config and use its' text config as llama config - hf_config = LlavaConfig.from_pretrained(args.model_dir).text_config - hf_config.model_type = "llava" # Replace llama with llava - - if hf_config.architectures[0] == "CogVLMForCausalLM": - hf_config.model_type = 'cogvlm' - args.model_type = hf_config.model_type - args.n_head = hf_config.num_attention_heads - args.inter_size = hf_config.intermediate_size - args.n_layer = hf_config.num_hidden_layers - args.n_embd = hf_config.hidden_size - if hasattr(hf_config, "num_key_value_heads"): - args.n_kv_head = hf_config.num_key_value_heads - if args.n_kv_head is None: - args.n_kv_head = args.n_head - args.rms_norm_eps = hf_config.rms_norm_eps - args.vocab_size = hf_config.vocab_size - args.n_positions = hf_config.max_position_embeddings - - args.architecture = hf_config.architectures[0] - - elif args.meta_ckpt_dir is not None: - with open(Path(args.meta_ckpt_dir, "params.json")) as fp: - meta_config: dict = json.load(fp) - args.n_embd = meta_config["dim"] - args.n_head = meta_config["n_heads"] - args.n_layer = meta_config["n_layers"] - args.n_kv_head = meta_config.get("n_kv_heads", args.n_head) - - if "hidden_dim" in meta_config: - args.inter_size = meta_config["hidden_dim"] - else: - args.multiple_of = meta_config.get("multiple_of", 1) - n_embd = int(4 * args.n_embd * 2 / 3) - args.ffn_dim_multiplier = meta_config.get("ffn_dim_multiplier", 1) - args.inter_size = args.multiple_of * ( - (int(n_embd * args.ffn_dim_multiplier) + args.multiple_of - 1) - // args.multiple_of) - args.rms_norm_eps = meta_config["norm_eps"] - args.architecture = "LlamaForCausalLM" - else: - args.n_kv_head = args.n_kv_head or args.n_head - args.architecture = "LlamaForCausalLM" - - 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 - - config = create_config_from_args(args) - - with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: - json.dump(config, f, indent=4) - - act_range = {} - llama_qkv_para = {} - # smoother for inputs of self_attn.o_proj and mlp.down_proj - llama_smoother = {} - model = None - if args.model_dir is not None: - - if args.model_type == "llava": - hf_llava = LlavaForConditionalGeneration.from_pretrained( - args.model_dir, dtype="auto") - model = hf_llava.language_model - else: - model = AutoModelForCausalLM.from_pretrained( - args.model_dir, - device_map='auto' if not args.load_model_on_cpu else 'cpu', - dtype='auto' if not args.smoothquant else torch.float16, - trust_remote_code=True, - ) - if args.smoothquant is not None or args.int8_kv_cache: - act_range, llama_qkv_para, llama_smoother = smooth_quant( - model, args) - - def covert_and_save(rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - - if args.use_weight_only and args.weight_only_precision == 'int4_gptq': - weights = load_weights_from_gptq( - args.quant_ckpt_path, - PretrainedConfig.from_dict(copy.deepcopy(config)), - ) - - elif args.meta_ckpt_dir is not None: - weights = load_weights_from_meta_ckpt( - args.meta_ckpt_dir, - PretrainedConfig.from_dict(copy.deepcopy(config)), - ) - - else: - if args.load_by_shard: - weights = load_weights_from_hf_by_shard( - args.model_dir, - PretrainedConfig.from_dict(copy.deepcopy(config)), - ) - - else: - if args.weight_only_precision == 'int8': - plugin_weight_only_quant_type = torch.int8 - elif args.weight_only_precision == 'int4': - plugin_weight_only_quant_type = torch.quint4x2 - weights = convert_hf_cogvlm( - model, - mapping, - vocab_size=args.vocab_size, - dtype=args.dtype, - use_weight_only=args.use_weight_only, - use_gemm_woq_plugin=not args. - disable_weight_only_quant_plugin, - plugin_weight_only_quant_type=plugin_weight_only_quant_type, - use_parallel_embedding=args.use_parallel_embedding, - sharding_dim=args.embedding_sharding_dim, - use_smooth_quant=args.smoothquant, - per_channel=args.per_channel, - per_token=args.per_token, - int8_kv_cache=args.int8_kv_cache, - act_range=act_range, - qkv_para=llama_qkv_para, - smoother=llama_smoother) - - safetensors.torch.save_file( - weights, os.path.join(args.output_dir, f'rank{rank}.safetensors')) - - if args.workers == 1: - - for rank in range(world_size): - covert_and_save(rank) - else: - with ThreadPoolExecutor(max_workers=args.workers) as p: - futures = [ - p.submit(covert_and_save, rank) for rank in range(world_size) - ] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/dbrx/README.md b/examples/models/contrib/dbrx/README.md deleted file mode 100644 index 85002397a83c..000000000000 --- a/examples/models/contrib/dbrx/README.md +++ /dev/null @@ -1,238 +0,0 @@ -# DBRX - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a DBRX model in TensorRT-LLM. DBRX is a leading large language model trained by Databricks. Read more details about the model: [DBRX Technical Blog](https://www.databricks.com/blog/introducing-dbrx-new-state-art-open-llm). - -- [DBRX](#dbrx) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [Download weights from HuggingFace Transformers](#download-weights-from-huggingface-transformers) - - [Build TensorRT engine(s)](#build-tensorrt-engines) - - [Weight Only Quantization](#weight-only-quantization) - - [INT8 KV Cache](#int8-kv-cache) - - [Run inference](#run-inference) - -## Overview - -The TensorRT LLM DBRX implementation can be found in [tensorrt_llm/models/dbrx/model.py](../../../../tensorrt_llm/models/dbrx/model.py). - -## Support Matrix - * BF16 - * FP16 - * INT8 Weight-Only - * INT4 Weight-Only - * INT8 KV CACHE - * Tensor Parallel - * Pipeline Parallel - * Expert Parallel - -## Usage - -### Download weights from HuggingFace Transformers - -Install the dependencies and setup `git-lfs`. - -```bash -# Install dependencies -# DBRX uses tiktoken as the tokenizer; make sure it is installed -pip install -r requirements.txt - -# Setup git-lfs -git lfs install -``` - -Download one or more DBRX models that you would like to build to TensorRT LLM engines. You can download from the [HuggingFace](https://huggingface.co) hub: - -```bash -# Download dbrx-base -git clone https://huggingface.co/databricks/dbrx-base - -# Download dbrx-instruct -git clone https://huggingface.co/databricks/dbrx-instruct -``` - -### Build TensorRT engine(s) - -The [`convert_checkpoint.py`](./convert_checkpoint.py) script converts HF weights to TensorRT LLM checkpoints. A DBRX model has 132B parameters, so you need at least 4 x 80GB GPUs to load the model in 16-bit precision for weight conversion. - -The `trtllm-build` command builds TensorRT LLM engines from TensorRT LLM checkpoints. The number of engine files is same to the number of GPUs used to run inference. Normally, `trtllm-build` uses one GPU by default, but if you have already more GPUs available at build time, you may enable parallel builds to make the engine building process faster by adding the `--workers` argument. - -Here are some examples: - -```bash -# 8-way tensor parallelism, dtype bfloat16 -python convert_checkpoint.py --model_dir dbrx-base \ - --dtype bfloat16 \ - --tp_size 8 \ - --workers 8 \ - --output_dir dbrx/trt_ckpt/bf16/tp8 - -trtllm-build --checkpoint_dir dbrx/trt_ckpt/bf16/tp8 \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 \ - --moe_plugin bfloat16 \ - --workers 8 \ - --output_dir dbrx/trt_engines/bf16/tp8 -``` - -```bash -# 8-way tensor parallelism, dtype float16 -python convert_checkpoint.py --model_dir dbrx-base \ - --dtype float16 \ - --tp_size 8 \ - --workers 8 \ - --output_dir dbrx/trt_ckpt/fp16/tp8 - -trtllm-build --checkpoint_dir dbrx/trt_ckpt/fp16/tp8 \ - --gpt_attention_plugin float16 \ - --gemm_plugin float16 \ - --moe_plugin float16 \ - --workers 8 \ - --output_dir dbrx/trt_engines/fp16/tp8 -``` - -```bash -# 4-way tensor parallelism and 2-way pipeline parallelism, dtype bfloat16 -python convert_checkpoint.py --model_dir dbrx-base \ - --dtype bfloat16 \ - --tp_size 4 \ - --pp_size 2 \ - --workers 8 \ - --output_dir dbrx/trt_ckpt/bf16/tp4pp2 - -trtllm-build --checkpoint_dir dbrx/trt_ckpt/bf16/tp4pp2 \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 \ - --moe_plugin bfloat16 \ - --workers 8 \ - --output_dir dbrx/trt_engines/bf16/tp4pp2 -``` - - -```bash -# Build DBRX with expert parallelism for DbrxExperts layer and tensor parallelism for rest -python convert_checkpoint.py --model_dir dbrx-base \ - --dtype bfloat16 \ - --tp_size 8 \ - --moe_tp_size 1 \ - --moe_ep_size 8 \ - --workers 8 \ - --output_dir dbrx/trt_ckpt/bf16/ep8 - -trtllm-build --checkpoint_dir dbrx/trt_ckpt/bf16/ep8 \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 \ - --moe_plugin bfloat16 \ - --workers 8 \ - --output_dir dbrx/trt_engines/bf16/ep8 -``` - -#### Weight Only Quantization - -[`convert_checkpoint.py`](./convert_checkpoint.py) features a `--use_weight_only` option that can enable weight-only quantization. You can further set the weight-only precision by passing `int8` or `int4` to the `--weight_only_precision` flag. - -```bash -# 4-way tensor parallelism, int8 weight-only -python convert_checkpoint.py --model_dir dbrx-base \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int8 \ - --tp_size 4 \ - --workers 4 \ - --output_dir dbrx/trt_ckpt/int8-wo/tp4 - -trtllm-build --checkpoint_dir dbrx/trt_ckpt/int8-wo/tp4 \ - --gpt_attention_plugin float16 \ - --gemm_plugin float16 \ - --moe_plugin float16 \ - --workers 4 \ - --output_dir dbrx/trt_engines/int8-wo/tp4 -``` - -```bash -# 4-way tensor parallelism, int4 weight-only -python convert_checkpoint.py --model_dir dbrx-base \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int4 \ - --tp_size 4 \ - --workers 4 \ - --output_dir dbrx/trt_ckpt/int4-wo/tp4 - -trtllm-build --checkpoint_dir dbrx/trt_ckpt/int4-wo/tp4 \ - --gpt_attention_plugin float16 \ - --gemm_plugin float16 \ - --moe_plugin float16 \ - --workers 4 \ - --output_dir dbrx/trt_engines/int4-wo/tp4 -``` - -#### INT8 KV Cache -INT8 KV cache can be enabled to reduce the memory footprint. It will bring performance gains at large batch sizes and long sequence lengths. - -For INT8 KV cache, [`convert_checkpoint.py`](./convert_checkpoint.py) features a `--int8_kv_cache` option. Setting `--int8_kv_cache` will calibrate the model, and then export the scaling factors needed for INT8 KV cache inference. - -```bash -# 4-way tensor parallelism, int8 kv-cache -python convert_checkpoint.py --model_dir dbrx-base \ - --dtype float16 \ - --int8_kv_cache \ - --tp_size 4 \ - --workers 4 \ - --output_dir dbrx/trt_ckpt/int8kv/tp4 - -trtllm-build --checkpoint_dir dbrx/trt_ckpt/int8kv/tp4 \ - --gpt_attention_plugin float16 \ - --gemm_plugin float16 \ - --moe_plugin float16 \ - --workers 4 \ - --output_dir dbrx/trt_engines/int8kv/tp4 -``` - -### Run inference - -You can test your engines with the [run.py](../../../run.py) script: - -```bash -mpirun -n 8 \ - python3 ../run.py --engine_dir dbrx/trt_engines/bf16/tp8 \ - --tokenizer_dir dbrx-base \ - --max_output_len 10 \ - --input_text "What is AGI?" -``` - -If the engines are run successfully, you will see output like: -``` -...... -Input [Text 0]: "What is AGI?" -Output [Text 0 Beam 0]: " How do I find it? -AGI stands for" -``` - - -You can also evaluate with the [summarize.py](../../../summarize.py) script: -```bash -mpirun -n 8 \ - python ../summarize.py --engine_dir dbrx/trt_engines/bf16/tp8 \ - --hf_model_dir dbrx-base \ - --test_trt_llm -``` - -If the engines are run successfully, you will see output like: -``` -...... -[04/02/2024-11:16:37] [TRT-LLM] [I] TensorRT LLM (total latency: 9.962657451629639 sec) -[04/02/2024-11:16:37] [TRT-LLM] [I] TensorRT LLM (total output tokens: 1189) -[04/02/2024-11:16:37] [TRT-LLM] [I] TensorRT LLM (tokens per second: 119.34566713477734) -[04/02/2024-11:16:37] [TRT-LLM] [I] TensorRT LLM beam 0 result -[04/02/2024-11:16:37] [TRT-LLM] [I] rouge1 : 26.842471264679535 -[04/02/2024-11:16:37] [TRT-LLM] [I] rouge2 : 9.979512100961314 -[04/02/2024-11:16:37] [TRT-LLM] [I] rougeL : 19.50336050538688 -[04/02/2024-11:16:37] [TRT-LLM] [I] rougeLsum : 22.00400189383231 -``` diff --git a/examples/models/contrib/dbrx/convert_checkpoint.py b/examples/models/contrib/dbrx/convert_checkpoint.py deleted file mode 100644 index 1ca287f2588a..000000000000 --- a/examples/models/contrib/dbrx/convert_checkpoint.py +++ /dev/null @@ -1,654 +0,0 @@ -import argparse -import copy -import functools -import json -import os -import time -import traceback -from collections import defaultdict -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Dict, Optional, Tuple - -import safetensors -import torch -import torch.nn as nn -from tqdm import tqdm -from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer -from transformers.pytorch_utils import Conv1D - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import get_hf_rope_theta, release_gc -from tensorrt_llm.layers import MoeConfig -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.convert_utils import (generate_int8, - load_calib_dataset, split) -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - 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('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument('--logits_dtype', - type=str, - default='float32', - choices=['float16', 'float32']) - 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( - '--calib_dataset', - type=str, - default='ccdv/cnn_dailymail', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - parser.add_argument("--dataset_cache_dir", - type=str, - default=None, - help="cache dir to load the hugging face dataset") - 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_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('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument('--rotary_base', type=float, default=10000.0) - parser.add_argument('--rotary_scaling', nargs=2, type=str, default=None) - parser.add_argument('--vocab_size', type=int, default=32000) - parser.add_argument('--n_positions', type=int, default=2048) - parser.add_argument('--n_layer', type=int, default=32) - parser.add_argument('--n_head', type=int, default=32) - parser.add_argument('--n_kv_head', type=int, default=None) - parser.add_argument('--n_embd', type=int, default=4096) - parser.add_argument('--inter_size', type=int, default=11008) - parser.add_argument('--max_seq_len', type=int, default=4096) - parser.add_argument('--clip_qkv', type=int, default=None) - parser.add_argument('--hidden_act', - type=str, - default='gelu', - help='Set to swiglu to use GLU in MoEs') - parser.add_argument( - '--moe_num_experts', - default=0, - type=int, - help='Specify the number of experts to use for MOE layers') - parser.add_argument( - '--moe_top_k', - default=0, - type=int, - help= - 'Specify the top_k value to use for MOE layers. Default to 1 if --moe_num_experts is set' - ) - parser.add_argument( - '--moe_tp_size', - type=int, - default=-1, - help= - 'N-way tensor parallelism size for MOE, default is tp_size, which will do tp-only for MoE' - ) - parser.add_argument( - '--moe_ep_size', - type=int, - default=-1, - help= - 'N-way expert parallelism size for MOE, default is 1, which will do tp-only for MoE' - ) - parser.add_argument( - '--moe_renorm_mode', - default=MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE, - type=int, - help= - 'Controls renormalization after gate logits. Check layers/moe.py for accepted values', - ) - parser.add_argument( - '--disable_weight_only_quant_plugin', - default=False, - action="store_true", - help= - 'By default, using plugin implementation for weight quantization. Enabling disable_weight_only_quant_plugin flag will use ootb implementation instead of plugin.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument( - '--dense_context_fmha', - default=False, - action='store_true', - help= - 'Enable dense fmha in context phase, otherwise sliding window attention.' - 'If dense_context_fmha=False, the sliding window size is the max attention window size.' - ) - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - parser.add_argument('--use_prompt_tuning', - action="store_true", - default=False) - args = parser.parse_args() - - return args - - -def args_to_build_options(args): - return { - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'disable_weight_only_quant_plugin': - args.disable_weight_only_quant_plugin - } - - -def get_weight(params: Dict[str, torch.Tensor], prefix: str, - dtype: torch.dtype) -> torch.Tensor: - if f'{prefix}' in params: - return params[f'{prefix}'].to(dtype).detach().cpu() - elif f'{prefix}.weight' not in params: - return None - return params[f'{prefix}.weight'].to(dtype).detach().cpu() - - -def get_bias(params: Dict[str, torch.Tensor], prefix: str, - dtype: torch.dtype) -> torch.Tensor: - if f'{prefix}.bias' not in params: - return None - return params[f'{prefix}.bias'].to(dtype).detach().cpu() - - -def get_weight_and_bias(params: Dict[str, torch.Tensor], prefix: str, - dtype: torch.dtype) -> Tuple[torch.Tensor]: - return get_weight(params, prefix, dtype), get_bias(params, prefix, dtype) - - -@torch.no_grad() -def capture_activation_range(model, - tokenizer, - dataset, - num_samples=1, - seq_len=512): - model.eval() - device = next(model.parameters()).device - act_scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) - - tokenizer.pad_token = tokenizer.eos_token - - 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].float() - - 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))) - - for i in tqdm(range(num_samples), desc="calibrating model"): - datapoint = dataset[i:i + 1] - line = copy.copy(datapoint) - line[0] = line[0] + ' TL;DR: ' - line[0] = line[0].strip() - line[0] = line[0].replace(" n't", "n't") - 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() - - return act_scales - - -def split_qkv_tp(qkv, n_head, n_kv_heads, n_hidden, tensor_parallel, rank): - """ - Splits the QKV matrix according to tensor parallelism - """ - kv_head_size = n_kv_heads * (n_hidden // n_head) - q, k, v = torch.split(qkv, [n_hidden, kv_head_size, kv_head_size], dim=0) - q = split(q, tensor_parallel, rank, dim=0) - k = split(k, tensor_parallel, rank, dim=0) - v = split(v, tensor_parallel, rank, dim=0) - return torch.concatenate([q, k, v], dim=0).contiguous() - - -def split_matrix(weight: torch.Tensor, tp_size: int, rank: int, - dim: int) -> torch.Tensor: - return split(weight, tp_size, rank, dim=dim) - - -def get_tllm_linear_weight( - weight: torch.Tensor, - prefix: str, - bias: Optional[torch.Tensor] = None, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8, - postfix='weight', - quant_scale_name=None) -> Dict[str, torch.Tensor]: - results = {} - if use_weight_only: - if weight.dim() > 2: - v = weight.transpose(1, 2).contiguous().clone() - else: - v = weight.t().contiguous().clone() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v.cpu(), plugin_weight_only_quant_type) - results[prefix + postfix] = processed_torch_weights - if quant_scale_name is not None: - results[quant_scale_name] = torch_weight_scales - else: - results[prefix + 'per_channel_scale'] = torch_weight_scales - else: - results[prefix + postfix] = weight.contiguous() - - if bias is not None: - results[f'{prefix}bias'] = bias - - return results - - -def convert_hf_dbrx(model_params: dict, - hf_config: AutoConfig, - mapping: Mapping, - dtype: str = 'float32', - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8, - moe_config: MoeConfig = None, - int8_kv_cache=False, - act_range=[]): - - weights = {} - tik = time.time() - - dtype = getattr(torch, dtype) - num_hidden_layers = hf_config.n_layers - num_head = hf_config.n_heads - num_kv_heads = hf_config.attn_config.kv_n_heads - num_hidden = hf_config.d_model - mlp_hidden_size = hf_config.ffn_config.ffn_hidden_size - layers_range = mapping.pp_layers(num_hidden_layers) - multi_query_mode = (num_kv_heads != num_head) - - for l in layers_range: - prefix = f'transformer.blocks.{l}' - tllm_prex = f'transformer.layers.{l-layers_range[0]}' - # Attention QKV (no bias) - qkv_w = get_weight(model_params, f'{prefix}.norm_attn_norm.attn.Wqkv', - dtype) - qkv_w = split_qkv_tp(qkv_w, num_head, num_kv_heads, num_hidden, - mapping.tp_size, mapping.tp_rank) - weights.update( - get_tllm_linear_weight(qkv_w, f'{tllm_prex}.attention.qkv.', None, - use_weight_only, - plugin_weight_only_quant_type)) - # Attention dense (no bias) - attn_dense_weight = get_weight( - model_params, f'{prefix}.norm_attn_norm.attn.out_proj', dtype) - attn_dense_w = split_matrix(attn_dense_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(attn_dense_w, - f'{tllm_prex}.attention.dense.', None, - use_weight_only, - plugin_weight_only_quant_type)) - - if int8_kv_cache: - qkv_weight = get_weight(model_params, - f'{prefix}.norm_attn_norm.attn.Wqkv', dtype) - qkv_weight = qkv_weight.t() - if not multi_query_mode: - qkv_weight = qkv_weight.reshape(num_hidden, 3, num_hidden) - int8_weights = generate_int8( - qkv_weight, - act_range.get(f'{prefix}.norm_attn_norm.attn.Wqkv'), - is_qkv=True, - multi_query_mode=multi_query_mode) - weights[ - f'{tllm_prex}.attention.kv_cache_scaling_factor'] = int8_weights[ - 'scale_y_quant_orig'].contiguous() - - # input layer_norm - input_ln_weight = get_weight(model_params, - f'{prefix}.norm_attn_norm.norm_1', dtype) - weights[f'{tllm_prex}.input_layernorm.weight'] = input_ln_weight - - # post layer_norm - post_ln_weight = get_weight(model_params, - f'{prefix}.norm_attn_norm.norm_2', dtype) - weights[f'{tllm_prex}.post_layernorm.weight'] = post_ln_weight - - if moe_config and moe_config.has_moe(): - # experts mlp w1 -> mlp gate - mlp_gate_weight = get_weight(model_params, - f'{prefix}.ffn.experts.mlp.w1', dtype) - mlp_gate_weight = mlp_gate_weight.reshape(-1, mlp_hidden_size, - num_hidden) - # moe expert parallel - mlp_gate_weight = split_matrix(mlp_gate_weight, - mapping.moe_ep_size, - mapping.moe_ep_rank, - dim=0) - # moe tensor parallel - mlp_gate_w = split_matrix(mlp_gate_weight, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=1) - - # experts mlp v1 -> mlp fc - mlp_fc_weight = get_weight(model_params, - f'{prefix}.ffn.experts.mlp.v1', dtype) - mlp_fc_weight = mlp_fc_weight.reshape(-1, mlp_hidden_size, - num_hidden) - # moe expert parallel - mlp_fc_weight = split_matrix(mlp_fc_weight, - mapping.moe_ep_size, - mapping.moe_ep_rank, - dim=0) - # moe tensor parallel - mlp_fc_w = split_matrix(mlp_fc_weight, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=1) - mlp_fc_w = torch.concat([mlp_fc_w, mlp_gate_w], dim=-2) - weights.update( - get_tllm_linear_weight(mlp_fc_w, f'{tllm_prex}.mlp.fc.', None, - use_weight_only, - plugin_weight_only_quant_type)) - - # experts mlp w2 -> mlp proj - mlp_proj_weight = get_weight(model_params, - f'{prefix}.ffn.experts.mlp.w2', dtype) - mlp_proj_weight = mlp_proj_weight.reshape(-1, mlp_hidden_size, - num_hidden).transpose( - 1, 2) - # moe expert parallel - mlp_proj_weight = split_matrix(mlp_proj_weight, - mapping.moe_ep_size, - mapping.moe_ep_rank, - dim=0) - # moe tensor parallel - mlp_proj_w = split_matrix(mlp_proj_weight, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=2) - weights.update( - get_tllm_linear_weight(mlp_proj_w, f'{tllm_prex}.mlp.proj.', - None, use_weight_only, - plugin_weight_only_quant_type)) - - # router mlp - router_weights = get_weight(model_params, - f'{prefix}.ffn.router.layer', - torch.float32) - weights[f'{tllm_prex}.mlp.router.weight'] = router_weights - - embed_w = get_weight(model_params, 'transformer.wte', dtype) - lm_head = get_weight(model_params, 'lm_head', dtype) - if mapping.is_first_pp_rank(): - # Embedding - weights['transformer.vocab_embedding.weight'] = embed_w - if mapping.is_last_pp_rank(): - if lm_head is None: - lm_head = embed_w.clone() - ln_f_w = get_weight(model_params, 'transformer.norm_f', dtype) - # ln_f weight and bias - weights['transformer.ln_f.weight'] = ln_f_w - weights['lm_head.weight'] = split_matrix(lm_head, - mapping.tp_size, - mapping.tp_rank, - dim=0) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def execute(workers, func, hf_model): - if workers == 1: - for rank, f in enumerate(func): - f(hf_model, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [ - p.submit(f, hf_model, rank) for rank, f in enumerate(func) - ] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -if __name__ == '__main__': - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - world_size = args.tp_size * args.pp_size - if (args.moe_tp_size == -1 and args.moe_ep_size == -1): - # moe default to tp-only - args.moe_tp_size = args.tp_size - args.moe_ep_size = 1 - elif (args.moe_tp_size == -1): - args.moe_tp_size = args.tp_size // args.moe_ep_size - elif (args.moe_ep_size == -1): - args.moe_ep_size = args.tp_size // args.moe_tp_size - assert (args.moe_tp_size * args.moe_ep_size == args.tp_size - ), "moe_tp_size * moe_ep_size must equal to tp_size" - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - quant_algo = None - kv_cache_quant_algo = None - plugin_weight_only_quant_type = None - if args.use_weight_only: - if args.weight_only_precision == 'int8': - plugin_weight_only_quant_type = torch.int8 - quant_algo = QuantAlgo.W8A16 - elif args.weight_only_precision == 'int4': - plugin_weight_only_quant_type = torch.quint4x2 - quant_algo = QuantAlgo.W4A16 - - if args.int8_kv_cache: - kv_cache_quant_algo = QuantAlgo.INT8 - - hf_config = None - - if args.model_dir is not None: - hf_config = AutoConfig.from_pretrained(args.model_dir, - trust_remote_code=True) - args.n_kv_head = hf_config.attn_config.kv_n_heads - args.n_layer = hf_config.n_layers - args.n_head = hf_config.n_heads - args.vocab_size = hf_config.vocab_size - args.n_embd = hf_config.d_model - args.inter_size = hf_config.ffn_config.ffn_hidden_size - args.max_seq_len = hf_config.max_seq_len - args.moe_num_experts = getattr(hf_config.ffn_config, "moe_num_experts", - 0) - args.moe_top_k = getattr(hf_config.ffn_config, "moe_top_k", 0) - if args.moe_num_experts and args.moe_top_k == 0: - args.moe_top_k = 1 - args.clip_qkv = hf_config.attn_config.clip_qkv - args.hidden_act = 'swiglu' - args.rotary_base = get_hf_rope_theta(hf_config.attn_config, 10000.0) - args.moe_config = MoeConfig( - num_experts=args.moe_num_experts, - top_k=args.moe_top_k, - normalization_mode=args.moe_renorm_mode).validate() - config = { - 'architecture': hf_config.architectures[0], - 'dtype': args.dtype, - 'logits_dtype': args.logits_dtype, - 'vocab_size': args.vocab_size, - 'hidden_size': args.n_embd, - 'intermediate_size': args.inter_size, - 'num_hidden_layers': args.n_layer, - 'num_attention_heads': args.n_head, - 'num_key_value_heads': args.n_kv_head, - 'max_position_embeddings': args.max_seq_len, - 'norm_epsilon': 1e-5, - 'position_embedding_type': 'rope_gpt_neox', - 'hidden_act': args.hidden_act, - 'rotary_base': args.rotary_base, - 'rotary_scaling': args.rotary_scaling, - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo, - }, - 'moe': { - "num_experts": args.moe_num_experts, - "top_k": args.moe_top_k, - "normalization_mode": args.moe_renorm_mode - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - 'moe_tp_size': args.moe_tp_size, - 'moe_ep_size': args.moe_ep_size, - }, - 'clip_qkv': args.clip_qkv, - 'dense_context_fmha': args.dense_context_fmha, - } - - config.update(args_to_build_options(args)) - - with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: - json.dump(config, f, indent=4) - - def load_from_hf(model_dir): - hf_model = AutoModelForCausalLM.from_pretrained(model_dir, - trust_remote_code=True, - device_map="auto", - dtype=getattr( - torch, args.dtype), - config=hf_config) - return hf_model - - def convert_and_save(hf_model, rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size, - moe_tp_size=args.moe_tp_size, - moe_ep_size=args.moe_ep_size) - act_range = {} - if args.int8_kv_cache: - tokenizer = AutoTokenizer.from_pretrained(args.model_dir, - padding_side='left', - trust_remote_code=True) - dataset = load_calib_dataset(args.calib_dataset, - cache_dir=args.dataset_cache_dir) - act_range = capture_activation_range(hf_model, tokenizer, dataset) - - hf_model = dict(hf_model.named_parameters()) - weights = convert_hf_dbrx( - hf_model, - hf_config, - mapping, - dtype=args.dtype, - use_weight_only=args.use_weight_only, - plugin_weight_only_quant_type=plugin_weight_only_quant_type, - moe_config=args.moe_config, - int8_kv_cache=args.int8_kv_cache, - act_range=act_range) - - safetensors.torch.save_file( - weights, os.path.join(args.output_dir, f'rank{rank}.safetensors')) - del weights - release_gc() - - if args.model_dir: - hf_model = load_from_hf(args.model_dir) - execute(args.workers, [convert_and_save] * world_size, hf_model) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') diff --git a/examples/models/contrib/dbrx/requirements.txt b/examples/models/contrib/dbrx/requirements.txt deleted file mode 100644 index 0b8f8e5e0f4f..000000000000 --- a/examples/models/contrib/dbrx/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score -tiktoken==0.6.0 diff --git a/examples/models/contrib/deepseek_v1/README.md b/examples/models/contrib/deepseek_v1/README.md deleted file mode 100755 index 4a49c8266705..000000000000 --- a/examples/models/contrib/deepseek_v1/README.md +++ /dev/null @@ -1,103 +0,0 @@ -# Deepseek-v1 - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run [deepseek-v1](https://arxiv.org/pdf/2401.06066) model in TensorRT-LLM. - -- [Deepseek-v1](#deepseek-v1) - - [Prerequisite](#prerequistie) - - [Hardware](#hardware) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [Build TensorRT engine(s)](#build-tensorrt-engines) - -## Prerequisite - -First, please download Deepseek-v1 weights from HF https://huggingface.co/deepseek-ai/deepseek-moe-16b-base. - -```bash -git lfs install -git clone https://huggingface.co/deepseek-ai/deepseek-moe-16b-base -``` - -## Hardware - -The Deepseek-v1 model requires 1x80G GPU memory. - -## Overview - -The TensorRT LLM Deepseek-v1 implementation can be found in [tensorrt_llm/models/deepseek_v1/model.py](../../tensorrt_llm/models/deepseek_v1/model.py). The TensorRT LLM Deepseek-v1 example code is located in [`examples/models/contrib/deepseek_v1`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert the Deepseek-v1 model into TensorRT LLM checkpoint format. - -In addition, there are three shared files in the parent folder [`examples`](../../../) can be used for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the model inference output by given an input text. -* [`../../../summarize.py`](../../../summarize.py) to summarize the article from [cnn_dailmail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset, it can running the summarize from HF model and TensorRT LLM model. -* [`../../../mmlu.py`](../../../mmlu.py) to running score script from https://github.com/declare-lab/instruct-eval to compare HF model and TensorRT LLM model on the MMLU dataset. - -## Support Matrix - -- [x] FP16 -- [x] TENSOR PARALLEL -- [x] FP8 - -## Usage - -The TensorRT LLM Deepseek-v1 example code locates at [examples/models/contrib/deepseek_v1](./). It takes PyTorch weights as input, and builds corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Build TensorRT engine(s) - -Below is the step-by-step to run Deepseek-v1 with TensorRT-LLM. - -First the checkpoint will be converted to the TensorRT LLM checkpoint format by apply [`convert_checkpoint.py`](./convert_checkpoint.py). After that, the TensorRT engine(s) can be build with TensorRT LLM checkpoint. - -```bash -# Build the bfloat16 engine from Deepseek-v1 HF weights. -python convert_checkpoint.py --model_dir ./deepseek_moe_16b/ \ - --output_dir ./trtllm_checkpoint_deepseek_v1_1gpu_bf16 \ - --dtype bfloat16 \ - --tp_size 1 -trtllm-build --checkpoint_dir ./trtllm_checkpoint_deepseek_v1_1gpu_bf16 \ - --output_dir ./trtllm_engines/deepseek_v1/bf16/tp1 \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 \ - --moe_plugin bfloat16 \ -``` - -Then, test the engine with [run.py](../../../run.py) script: - -```bash -python ../../../run.py --engine_dir ./trtllm_engines/deepseek_v1/bf16/tp1 \ - --tokenizer_dir ./deepseek_moe_16b/ \ - --max_output_len 32 \ - --top_p 0 \ - --input_text "The president of the United States is person who" -``` - -### FP8 Quantization - -The [`../../../quantization/quantize.py`](../../../quantization/quantize.py) script can be used to quantize the models and export TensorRT LLM checkpoints. - -```bash -# Deepseek-v1: single gpu, fp8 quantization -python ../../../quantization/quantize.py --model_dir deepseek_moe_16b \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir trt_ckpt/deepseek_moe_16b/fp8/1-gpu \ - --calib_size 512 - -# Deepseek-v1: single-gpu engine with fp8 quantization, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir ./trt_ckpt/deepseek_moe_16b/fp8/1-gpu \ - --gemm_plugin float16 \ - --gpt_attention_plugin bfloat16 \ - --output_dir ./trt_engines/fp8/1-gpu/ -``` -## Credits -This Deepseek-v1 model example exists thanks to @akhoroshev(https://github.com/akhoroshev) community contribution! diff --git a/examples/models/contrib/deepseek_v1/__init__.py b/examples/models/contrib/deepseek_v1/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/examples/models/contrib/deepseek_v1/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. diff --git a/examples/models/contrib/deepseek_v1/convert_checkpoint.py b/examples/models/contrib/deepseek_v1/convert_checkpoint.py deleted file mode 100644 index 3e7b1aaf2825..000000000000 --- a/examples/models/contrib/deepseek_v1/convert_checkpoint.py +++ /dev/null @@ -1,210 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import argparse -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.layers import MoeConfig -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import DeepseekForCausalLM - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None, required=True) - 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( - '--moe_tp_size', - type=int, - default=-1, - help= - 'N-way tensor parallelism size for MoE, default is tp_size, which will do tp-only for MoE' - ) - parser.add_argument( - '--moe_ep_size', - type=int, - default=-1, - help= - 'N-way expert parallelism size for MoE, default is 1, which will do tp-only for MoE' - ) - parser.add_argument( - '--dtype', - type=str, - default='auto', - choices=['auto', 'float16', 'bfloat16', 'float32'], - help= - "The data type for the model weights and activations if not quantized. " - "If 'auto', the data type is automatically inferred from the source model; " - "however, if the source dtype is float32, it is converted to float16.") - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0)' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim=0') - parser.add_argument('--output_dir', - type=str, - default='trtllm_checkpoint', - required=True, - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument( - '--moe_num_experts', - type=int, - default=0, - help='Specify the number of experts to use for MOE layers') - parser.add_argument( - '--moe_top_k', - type=int, - default=0, - help= - 'Specify the top_k value to use for MOE layers. Default to 1 if --moe_num_experts is set' - ) - parser.add_argument( - '--moe_renorm_mode', - type=int, - default=MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE, - help= - 'Controls renormalization after gate logits. Check layers/moe.py for accepted values' - ) - parser.add_argument( - '--save_config_only', - action="store_true", - default=False, - help= - 'Only save the model config w/o read and converting weights, be careful, this is for debug only' - ) - parser.add_argument( - '--disable_weight_only_quant_plugin', - default=False, - action="store_true", - help= - 'By default, using plugin implementation for weight quantization. Enabling disable_weight_only_quant_plugin flag will use ootb implementation instead of plugin.' - 'You must also use --use_weight_only for that argument to have an impact' - ) - # Add quantization related feature later - args = parser.parse_args() - - return args - - -def args_to_build_options(args): - return { - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'disable_weight_only_quant_plugin': - args.disable_weight_only_quant_plugin - } - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def convert_and_save_hf(args): - world_size = args.tp_size * args.pp_size - # Need to convert the cli args to the kay-value pairs and override them in the generate config dict. - # Ideally these fields will be moved out of the config and pass them into build API, keep them here for compatibility purpose for now, - # before the refactor is done. - override_fields = {} - override_fields.update(args_to_build_options(args)) - - def convert_and_save_rank(args, rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size, - moe_tp_size=args.moe_tp_size, - moe_ep_size=args.moe_ep_size) - deepseekv1 = DeepseekForCausalLM.from_hugging_face( - args.model_dir, args.dtype, mapping, **override_fields) - deepseekv1.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del deepseekv1 - - execute(args.workers, [convert_and_save_rank] * world_size, args) - release_gc() - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - - args.tp_size * args.pp_size - if (args.moe_tp_size == -1 and args.moe_ep_size == -1): - # moe default to tp-only - args.moe_tp_size = args.tp_size - args.moe_ep_size = 1 - elif (args.moe_tp_size == -1): - args.moe_tp_size = args.tp_size // args.moe_ep_size - elif (args.moe_ep_size == -1): - args.moe_ep_size = args.tp_size // args.moe_tp_size - assert (args.moe_tp_size * args.moe_ep_size == args.tp_size - ), "moe_tp_size * moe_ep_size must equal to tp_size" - - tik = time.time() - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - assert args.model_dir is not None - convert_and_save_hf(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/deepseek_v1/requirements.txt b/examples/models/contrib/deepseek_v1/requirements.txt deleted file mode 100644 index 2f8713d8658a..000000000000 --- a/examples/models/contrib/deepseek_v1/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score diff --git a/examples/models/contrib/deepseek_v2/README.md b/examples/models/contrib/deepseek_v2/README.md deleted file mode 100644 index aefbeec4120a..000000000000 --- a/examples/models/contrib/deepseek_v2/README.md +++ /dev/null @@ -1,267 +0,0 @@ -# Deepseek-v2 - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run [deepseek-v2](https://arxiv.org/pdf/2405.04434) model in TensorRT-LLM. - -- [Deepseek-v2](#deepseek-v2) - - [Prerequisite](#prerequisite) - - [Hardware](#hardware) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [Build TensorRT engine(s)](#build-tensorrt-engines) - -## Prerequisite - -First, please download Deepseek-v2 weights from HF https://huggingface.co/deepseek-ai/DeepSeek-V2. - -```bash -git lfs install -git clone https://huggingface.co/deepseek-ai/DeepSeek-V2 -``` - -## Hardware - -The Deepseek-v2 model requires least 8x80G GPU memory, model contains 236B parameters roughly 472GB memory (with BF16 precision). - -***Caution: Current TRT-LLM MLA kernel only support Hopper architecture (SM90). Ampere architecture (SM80 & SM86) will be supported in next release.*** - -## Overview - -The TensorRT LLM Deepseek-v2 implementation can be found in [tensorrt_llm/models/deepseek_v2/model.py](../../../../tensorrt_llm/models/deepseek_v2/model.py). The TensorRT LLM Deepseek-v2 example code is located in [`examples/models/contrib/deepseek_v2`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert the Deepseek-v2 model into TensorRT LLM checkpoint format. - -In addition, there are three shared files in the parent folder [`examples`](../../../) can be used for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the model inference output by given an input text. -* [`../../../summarize.py`](../../../summarize.py) to summarize the article from [cnn_dailmail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset, it can running the summarize from HF model and TensorRT LLM model. -* [`../../../mmlu.py`](../../../mmlu.py) to running score script from https://github.com/declare-lab/instruct-eval to compare HF model and TensorRT LLM model on the MMLU dataset. - -## Support Matrix - -- [x] BF16 -- [ ] FP8 - -***Caution: prefer using BF16 over FP16 for Deepseek-v2 since model original training precision is BF16 and we found direct convert BF16 -> FP16 will suffer accuracy drop.*** - -## Usage - -The TensorRT LLM Deepseek-v2 example code locates at [examples/models/contrib/deepseek_v2](./). It takes PyTorch weights as input, and builds corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Build TensorRT engine(s) - -Below is the step-by-step to run Deepseek-v2 with TensorRT-LLM. - -First the checkpoint will be converted to the TensorRT LLM checkpoint format by apply [`convert_checkpoint.py`](./convert_checkpoint.py). After that, the TensorRT engine(s) can be build with TensorRT LLM checkpoint. - -```bash -# Convert Deepseek-v2 HF weights to TensorRT LLM checkpoint format. -python convert_checkpoint.py --model_dir ./DeepSeek-V2 \ - --output_dir ./trtllm_checkpoint_deepseek_v2_8gpu_bf16 \ - --dtype bfloat16 \ - --tp_size 8 - -# With '--load_model_on_cpu' option if total GPU memory is insufficient -python convert_checkpoint.py --model_dir ./DeepSeek-V2 \ - --output_dir ./trtllm_checkpoint_deepseek_v2_cpu_bf16 \ - --dtype bfloat16 \ - --tp_size 8 \ - --load_model_on_cpu -``` - - -We observe use GPUs(8xH200) the checkpoint conversion time took ~ 34 mints, while use CPUs took ~ 21 mints and CPU memory required >= 770GB. - -After the checkpoint conversion, the TensorRT engine(s) can be built with the TensorRT LLM checkpoint. - -```bash -# Build engine -trtllm-build --checkpoint_dir ./trtllm_checkpoint_deepseek_v2_8gpu_bf16 \ - --output_dir ./trtllm_engines/deepseek_v2/bf16/tp8-sel4096-isl2048-bs4 \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 \ - --max_batch_size 4 \ - --max_seq_len 4096 \ - --max_input_len 2048 \ - --use_paged_context_fmha enable -``` - -***Caution: `--max_batch_size` and `--max_seq_len` are the main factors to determine how many GPU memory will be used during runtime, so later when try to run e.g., `summarize.py` or `mmlu.py` or `gptManagerBenchmark.cpp`may need adjust `--max_batch_size` and `--max_seq_len` accordingly to avoid OOM.(meaning rebuild TensorRT engine with smaller `--max_batch_size` and `--max_seq_len` if needed based on GPU memory size), there is beautiful technical log perf-best-practices.md (https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/performance/perf-best-practices.md) explained the mechanism.*** - -Test the engine with [run.py](../../../run.py) script: - -```bash -mpirun --allow-run-as-root -n 8 python ../../../run.py --engine_dir ./trtllm_engines/deepseek_v2/bf16/tp8-sel4096-isl2048-bs4 \ - --tokenizer_dir ./DeepSeek-V2 \ - --max_output_len 40 \ - --input_text "The president of the United States is person who" -``` - -and the output will be like: - -``` -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31490468978882 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31163835525513 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31164216995239 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31491041183472 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.3116364479065 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.3118085861206 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.3118691444397 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31516337394714 sec -Input [Text 0]: "<｜begin▁of▁sentence｜>The president of the United States is person who" -Output [Text 0 Beam 0]: " is elected by the people of the United States to lead the country. The president is the head of the executive branch of the government. The president is also the commander in chief of the armed forces." -[TensorRT-LLM][INFO] Refreshed the MPI local session -[TensorRT-LLM][INFO] Refreshed the MPI local session -[TensorRT-LLM][INFO] Refreshed the MPI local session -[TensorRT-LLM][INFO] Refreshed the MPI local session -[TensorRT-LLM][INFO] Refreshed the MPI local session -[TensorRT-LLM][INFO] Refreshed the MPI local session -[TensorRT-LLM][INFO] Refreshed the MPI local session -[TensorRT-LLM][INFO] Refreshed the MPI local session -``` - -If we want to evaluate the model summarization ability, we can use [summarize.py](../../../summarize.py) script: - -```bash -mpirun --allow-run-as-root -n 8 python ../../../summarize.py --engine_dir ./trtllm_engines/deepseek_v2/bf16/tp8-sel4096-isl2048-bs4 \ - --hf_model_dir ./DeepSeek-V2 \ - --data_type bfloat16 \ - --batch_size 1 \ - --test_trt_llm \ - --test_hf -``` - -and the output will be like: - - -``` -[10/28/2024-16:46:22] [TRT-LLM] [I] HF Generated : -[10/28/2024-16:46:22] [TRT-LLM] [I] Input : ['(CNN)James Best, best known for his portrayal of bumbling sheriff Rosco P. Coltrane on TV\'s "The Dukes of Hazzard," died Monday after a brief illness. He was 88. Best died in hospice in Hickory, North Carolina, of complications from pneumonia, said Steve Latshaw, a longtime friend and Hollywood colleague. Although he\'d been a busy actor for decades in theater and in Hollywood, Best didn\'t become famous until 1979, when "The Dukes of Hazzard\'s" cornpone charms began beaming into millions of American homes almost every Friday night. For seven seasons, Best\'s Rosco P. Coltrane chased the moonshine-running Duke boys back and forth across the back roads of fictitious Hazzard County, Georgia, although his "hot pursuit" usually ended with him crashing his patrol car. Although Rosco was slow-witted and corrupt, Best gave him a childlike enthusiasm that got laughs and made him endearing. His character became known for his distinctive "kew-kew-kew" chuckle and for goofy catchphrases such as "cuff \'em and stuff \'em!" upon making an arrest. Among the most popular shows on TV in the early \'80s, "The Dukes of Hazzard" ran until 1985 and spawned TV movies, an animated series and video games. Several of Best\'s "Hazzard" co-stars paid tribute to the late actor on social media. "I laughed and learned more from Jimmie in one hour than from anyone else in a whole year," co-star John Schneider, who played Bo Duke, said on Twitter. "Give Uncle Jesse my love when you see him dear friend." "Jimmy Best was the most constantly creative person I have ever known," said Ben Jones, who played mechanic Cooter on the show, in a Facebook post. "Every minute of his long life was spent acting, writing, producing, painting, teaching, fishing, or involved in another of his life\'s many passions." Born Jewel Guy on July 26, 1926, in Powderly, Kentucky, Best was orphaned at 3 and adopted by Armen and Essa Best, who renamed him James and raised him in rural Indiana. Best served in the Army during World War II before launching his acting career. In the 1950s and 1960s, he accumulated scores of credits, playing a range of colorful supporting characters in such TV shows as "The Twilight Zone," "Bonanza," "The Andy Griffith Show" and "Gunsmoke." He later appeared in a handful of Burt Reynolds\' movies, including "Hooper" and "The End." But Best will always be best known for his "Hazzard" role, which lives on in reruns. "Jimmie was my teacher, mentor, close friend and collaborator for 26 years," Latshaw said. "I directed two of his feature films, including the recent \'Return of the Killer Shrews,\' a sequel he co-wrote and was quite proud of as he had made the first one more than 50 years earlier." People we\'ve lost in 2015 . CNN\'s Stella Chan contributed to this story.'] -[10/28/2024-16:46:22] [TRT-LLM] [I] - Reference : ['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 .'] -[10/28/2024-16:46:22] [TRT-LLM] [I] - Output : [[' James Best, best known for his portrayal of bumbling sheriff Rosco P. Coltrane on TV\'s "The Dukes of Hazzard," died Monday after a brief illness. He was 88.']] -[10/28/2024-16:46:22] [TRT-LLM] [I] --------------------------------------------------------- -[10/28/2024-16:49:33] [TRT-LLM] [I] TensorRT LLM (total latency: 32.02327513694763 sec) -[10/28/2024-16:49:33] [TRT-LLM] [I] TensorRT LLM (total output tokens: 1394) -[10/28/2024-16:49:33] [TRT-LLM] [I] TensorRT LLM (tokens per second: 43.53083793080361) -[10/28/2024-16:49:33] [TRT-LLM] [I] TensorRT LLM beam 0 result -[10/28/2024-16:49:33] [TRT-LLM] [I] rouge1 : 17.85755990133811 -[10/28/2024-16:49:33] [TRT-LLM] [I] rouge2 : 6.273032755727469 -[10/28/2024-16:49:33] [TRT-LLM] [I] rougeL : 14.768323033457317 -[10/28/2024-16:49:33] [TRT-LLM] [I] rougeLsum : 15.700915348496391 -[10/28/2024-16:49:33] [TRT-LLM] [I] Hugging Face (total latency: 189.76398921012878 sec) -[10/28/2024-16:49:33] [TRT-LLM] [I] Hugging Face (total output tokens: 1376) -[10/28/2024-16:49:33] [TRT-LLM] [I] Hugging Face (tokens per second: 7.2511123197159) -[10/28/2024-16:49:33] [TRT-LLM] [I] HF beam 0 result -[10/28/2024-16:49:33] [TRT-LLM] [I] rouge1 : 18.542590123197257 -[10/28/2024-16:49:33] [TRT-LLM] [I] rouge2 : 6.345777100488389 -[10/28/2024-16:49:33] [TRT-LLM] [I] rougeL : 15.235695878419156 -[10/28/2024-16:49:33] [TRT-LLM] [I] rougeLsum : 16.64935135356226 -``` - -At last, we can evaluate the model with [mmlu.py](../../../mmlu.py) script: - -```bash -# Download MMLU dataset -mkdir mmlu_data && cd mmlu_data -wget https://people.eecs.berkeley.edu/~hendrycks/data.tar && tar -xf data.tar -# Run MMLU evaluation -mpirun --allow-run-as-root -n 8 python ../../../mmlu.py --engine_dir ./trtllm_engines/deepseek_v2/bf16/tp8-sel4096-isl2048-bs4 \ - --hf_model_dir ./DeepSeek-V2 \ - --data_type bfloat16 \ - --batch_size 1 \ - --test_trt_llm \ - --test_hf \ - --data_dir ./mmlu_data/data/ -``` - -and the output will be like: - -``` -Average accuracy 0.480 - abstract_algebra -Average accuracy 0.741 - anatomy -Average accuracy 0.888 - astronomy -Average accuracy 0.790 - business_ethics -Average accuracy 0.845 - clinical_knowledge -Average accuracy 0.924 - college_biology -Average accuracy 0.600 - college_chemistry -Average accuracy 0.720 - college_computer_science -Average accuracy 0.510 - college_mathematics -Average accuracy 0.751 - college_medicine -Average accuracy 0.618 - college_physics -Average accuracy 0.860 - computer_security -Average accuracy 0.796 - conceptual_physics -Average accuracy 0.675 - econometrics -Average accuracy 0.800 - electrical_engineering -Average accuracy 0.741 - elementary_mathematics -Average accuracy 0.643 - formal_logic -Average accuracy 0.570 - global_facts -Average accuracy 0.910 - high_school_biology -Average accuracy 0.714 - high_school_chemistry -Average accuracy 0.890 - high_school_computer_science -Average accuracy 0.891 - high_school_european_history -Average accuracy 0.909 - high_school_geography -Average accuracy 0.953 - high_school_government_and_politics -Average accuracy 0.826 - high_school_macroeconomics -Average accuracy 0.522 - high_school_mathematics -Average accuracy 0.916 - high_school_microeconomics -Average accuracy 0.576 - high_school_physics -Average accuracy 0.923 - high_school_psychology -Average accuracy 0.704 - high_school_statistics -Average accuracy 0.907 - high_school_us_history -Average accuracy 0.932 - high_school_world_history -Average accuracy 0.834 - human_aging -Average accuracy 0.893 - human_sexuality -Average accuracy 0.909 - international_law -Average accuracy 0.880 - jurisprudence -Average accuracy 0.853 - logical_fallacies -Average accuracy 0.598 - machine_learning -Average accuracy 0.874 - management -Average accuracy 0.953 - marketing -Average accuracy 0.880 - medical_genetics -Average accuracy 0.920 - miscellaneous -Average accuracy 0.850 - moral_disputes -Average accuracy 0.613 - moral_scenarios -Average accuracy 0.830 - nutrition -Average accuracy 0.859 - philosophy -Average accuracy 0.883 - prehistory -Average accuracy 0.635 - professional_accounting -Average accuracy 0.625 - professional_law -Average accuracy 0.857 - professional_medicine -Average accuracy 0.833 - professional_psychology -Average accuracy 0.745 - public_relations -Average accuracy 0.869 - security_studies -Average accuracy 0.935 - sociology -Average accuracy 0.930 - us_foreign_policy -Average accuracy 0.596 - virology -Average accuracy 0.877 - world_religions -Average accuracy 0.632 - math -Average accuracy 0.801 - health -Average accuracy 0.738 - physics -Average accuracy 0.897 - business -Average accuracy 0.914 - biology -Average accuracy 0.677 - chemistry -Average accuracy 0.762 - computer science -Average accuracy 0.832 - economics -Average accuracy 0.800 - engineering -Average accuracy 0.736 - philosophy -Average accuracy 0.821 - other -Average accuracy 0.902 - history -Average accuracy 0.909 - geography -Average accuracy 0.883 - politics -Average accuracy 0.876 - psychology -Average accuracy 0.919 - culture -Average accuracy 0.660 - law -Average accuracy 0.727 - STEM -Average accuracy 0.740 - humanities -Average accuracy 0.873 - social sciences -Average accuracy 0.821 - other (business, health, misc.) -Average accuracy: 0.785 -``` diff --git a/examples/models/contrib/deepseek_v2/convert_checkpoint.py b/examples/models/contrib/deepseek_v2/convert_checkpoint.py deleted file mode 100755 index 5a0c6a7df330..000000000000 --- a/examples/models/contrib/deepseek_v2/convert_checkpoint.py +++ /dev/null @@ -1,231 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import argparse -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.layers import MoeConfig -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import DeepseekV2ForCausalLM -from tensorrt_llm.models.deepseek_v2.convert import load_hf_deepseek - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None, required=True) - 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( - '--moe_tp_size', - type=int, - default=-1, - help= - 'N-way tensor parallelism size for MoE, default is tp_size, which will do tp-only for MoE' - ) - parser.add_argument( - '--moe_ep_size', - type=int, - default=-1, - help= - 'N-way expert parallelism size for MoE, default is 1, which will do tp-only for MoE' - ) - parser.add_argument('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument('--load_model_on_cpu', - default=False, - action="store_true", - help='Choose to load HF cpkt into CPU') - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0)' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim=0') - parser.add_argument('--output_dir', - type=str, - default='trtllm_checkpoint', - required=True, - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument( - '--moe_num_experts', - type=int, - default=0, - help='Specify the number of experts to use for MOE layers') - parser.add_argument( - '--moe_top_k', - type=int, - default=0, - help= - 'Specify the top_k value to use for MOE layers. Default to 1 if --moe_num_experts is set' - ) - parser.add_argument( - '--moe_renorm_mode', - type=int, - default=MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE, - help= - 'Controls renormalization after gate logits. Check layers/moe.py for accepted values' - ) - parser.add_argument('--use_preloading', - action="store_true", - default=False, - help='Loading weights from HF model') - parser.add_argument('--use_safetensors_loading', - action="store_true", - default=False, - help='Loading weights from HF safetensors') - parser.add_argument( - '--save_config_only', - action="store_true", - default=False, - help= - 'Only save the model config w/o read and converting weights, be careful, this is for debug only' - ) - parser.add_argument( - '--disable_weight_only_quant_plugin', - default=False, - action="store_true", - help= - 'By default, using plugin implementation for weight quantization. Enabling disable_weight_only_quant_plugin flag will use ootb implementation instead of plugin.' - 'You must also use --use_weight_only for that argument to have an impact' - ) - # Add quantization related feature later - args = parser.parse_args() - - return args - - -def args_to_build_options(args): - return { - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'disable_weight_only_quant_plugin': - args.disable_weight_only_quant_plugin, - 'load_model_on_cpu': args.load_model_on_cpu - } - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def convert_and_save_hf(args): - world_size = args.tp_size * args.pp_size - # Need to convert the cli args to the kay-value pairs and override them in the generate config dict. - # Ideally these fields will be moved out of the config and pass them into build API, keep them here for compatibility purpose for now, - # before the refactor is done. - override_fields = {} - override_fields.update(args_to_build_options(args)) - use_preloading = args.use_preloading - use_safetensors_loading = args.use_safetensors_loading - - args.load_model_on_cpu - - if os.environ.get("TRTLLM_DISABLE_UNIFIED_CONVERTER" - ) is not None and not use_safetensors_loading: - hf_model = load_hf_deepseek(args.model_dir, args.load_model_on_cpu) - else: - hf_model = None - - def convert_and_save_rank(args, rank): - - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size, - moe_tp_size=args.moe_tp_size, - moe_ep_size=args.moe_ep_size) - deepseekv2 = DeepseekV2ForCausalLM.from_hugging_face( - args.model_dir, args.dtype, hf_model, use_preloading, - use_safetensors_loading, mapping, **override_fields) - deepseekv2.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del deepseekv2 - - execute(args.workers, [convert_and_save_rank] * world_size, args) - release_gc() - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - - args.tp_size * args.pp_size - if (args.moe_tp_size == -1 and args.moe_ep_size == -1): - # moe default to tp-only - args.moe_tp_size = args.tp_size - args.moe_ep_size = 1 - elif (args.moe_tp_size == -1): - args.moe_tp_size = args.tp_size // args.moe_ep_size - elif (args.moe_ep_size == -1): - args.moe_ep_size = args.tp_size // args.moe_tp_size - assert (args.moe_tp_size * args.moe_ep_size == args.tp_size - ), "moe_tp_size * moe_ep_size must equal to tp_size" - - tik = time.time() - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - assert args.model_dir is not None - convert_and_save_hf(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/deepseek_v2/requirements.txt b/examples/models/contrib/deepseek_v2/requirements.txt deleted file mode 100644 index 6be396146581..000000000000 --- a/examples/models/contrib/deepseek_v2/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -datasets==3.1.0 -evaluate~=0.4.1 -rouge_score~=0.1.2 diff --git a/examples/models/contrib/dit/README.md b/examples/models/contrib/dit/README.md deleted file mode 100644 index 868ecbf8adef..000000000000 --- a/examples/models/contrib/dit/README.md +++ /dev/null @@ -1,133 +0,0 @@ -# DiT - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a [DiT](https://arxiv.org/abs/2212.09748) model with TensorRT-LLM. - -## Overview - -The TensorRT LLM DiT implementation can be found in [tensorrt_llm/models/dit/model.py](../../../../tensorrt_llm/models/dit/model.py). The TensorRT LLM DiT example code is located in [`examples/dit`](./). There are main files to build and run DiT with TensorRT-LLM: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert the DiT model into TensorRT LLM checkpoint format. -* [`sample.py`](./sample.py) to generate images with TensorRT engine(s). - -## Support Matrix - -- [x] FP16 -- [x] TP -- [x] FP8 -- [x] CP - -## Usage - -The TensorRT LLM DiT example code locates at [examples/dit](./). It takes PyTorch 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 DiT TensorRT engine(s) - -First, download the pretrained DiT-XL/2 PyTorch checkpoint from the official pytorch implementation repo [here](https://github.com/facebookresearch/DiT/tree/main?tab=readme-ov-file#sampling--), please review its license items before use. - -This checkpoint will be converted to the TensorRT LLM checkpoint format by [`convert_checkpoint.py`](./convert_checkpoint.py). After that, we can build TensorRT engine(s) with the TensorRT LLM checkpoint. - -As for run inference with FP8 quantization, currently only linear layers are supported to be quantized. Make sure that scaling factors for weights are also stored in the quantized checkpoint. - -``` -# Convert to TRT-LLM with float16(by default) -python convert_checkpoint.py -trtllm-build --checkpoint_dir ./tllm_checkpoint/ \ - --max_batch_size 8 \ - --remove_input_padding disable \ - --bert_attention_plugin disable - -# Convert to TRT-LLM with float8 -python convert_checkpoint.py --fp8_linear --timm_ckpt= --output_dir=tllm_checkpoint_fp8 -trtllm-build --checkpoint_dir ./tllm_checkpoint_fp8/ \ - --output_dir ./engine_outputs_fp8/ \ - --max_batch_size 8 \ - --remove_input_padding disable \ - --bert_attention_plugin disable -``` - -The default quantized checkpoint is based on official DiT by Facebook. If the checkpoint is obtained with APIs of `diffusers` (provided by HuggingFace), please enable the option `--diffusers_dit`: - -``` -python convert_checkpoint.py --fp8_linear --diffusers_dit --timm_ckpt= --output_dir=tllm_checkpoint_fp8 -trtllm-build --checkpoint_dir ./tllm_checkpoint_fp8/ \ - --output_dir ./engine_outputs_fp8/ \ - --max_batch_size 8 \ - --remove_input_padding disable \ - --bert_attention_plugin disable -``` - -Set `--max_batch_size` to tell how many images at most you would like to generate. We disable `--remove_input_padding` since we don't need to padding DiT's patches. Besides, we disable `--bert_attention_plugin` for better performance, since the plugin's fmha is not supported for DiT's hidden size (72 for DiT-XL). - -After build, we can find a `./engine_output` directory, it is ready for running DiT model with TensorRT LLM now. - -### Build VAE TensorRT engine -We can further accelerate VAE decoder by TensorRT. -``` -python vae_decoder_trt.py --max_batch_size 8 -``` - -### Generate images - -A [`sample.py`](./sample.py) is provided to generated images with the optimized TensorRT engines. - -We can simply run `python sample.py`. After that, an image named `sample.png` will be generated: -![sample.png](./figs/sample.png). - -Image generated with quantized FP8 DiT: - -![sample.fp8.png](./figs/sample.fp8.png). - -By default, we set batch size to 2, you can set batch size up to `--max-batch-size`, e.g., `python sample.py --batch-size 8` - -We can also run `python sample.py --tllm_model_dir=` to apply specific engines. - -### Tensor Parallel - -We can levaerage tensor parallel to further reduce latency and memory consumption on each GPU. We can also run TP for FP8 DiT. - -``` -# build dit engine -python convert_checkpoint.py --tp_size 4 -trtllm-build --checkpoint_dir ./tllm_checkpoint/ \ - --max_batch_size 8 \ - --remove_input_padding disable \ - --bert_attention_plugin disable -# build vae engine -python vae_decoder_trt.py --max_batch_size 8 -# run -mpirun -n 4 --allow-run-as-root python sample.py -``` - -### Context Parallel - -Context parallel can also be used to reduce latency and memory consumption on each GPU. We can also run CP for FP8 DiT. - -``` -# build dit engine -python convert_checkpoint.py --cp_size 4 -trtllm-build --checkpoint_dir ./tllm_checkpoint/ --max_batch_size 8 --remove_input_padding disable --bert_attention_plugin disable -# build vae engine -python vae_decoder_trt.py --max_batch_size 8 -# run -mpirun -n 4 --allow-run-as-root python sample.py -``` - -### Combine Tensor Parallel and Context Parallel - -Tensor Parallel and Context Parallel can be used together to better balance latency and memory consumption. - -``` -# build dit engine -python convert_checkpoint.py --cp_size 2 --tp_size 2 -trtllm-build --checkpoint_dir ./tllm_checkpoint/ --max_batch_size 8 --remove_input_padding disable --bert_attention_plugin disable -# build vae engine -python vae_decoder_trt.py --max_batch_size 8 -# run -mpirun -n 4 --allow-run-as-root python sample.py -``` diff --git a/examples/models/contrib/dit/convert_checkpoint.py b/examples/models/contrib/dit/convert_checkpoint.py deleted file mode 100644 index e36c9f2b9c2c..000000000000 --- a/examples/models/contrib/dit/convert_checkpoint.py +++ /dev/null @@ -1,334 +0,0 @@ -import argparse -import json -import os -import re -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import safetensors.torch -import torch - -import tensorrt_llm -from tensorrt_llm import str_dtype_to_torch -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.convert_utils import (split, split_matrix_tp, - split_qkv_bias_tp, split_qkv_tp) - -FACEBOOK_DIT_NAME_MAPPING = { - '^t_embedder.mlp.0.weight$': - 't_embedder.mlp1.weight', - '^t_embedder.mlp.0.bias$': - 't_embedder.mlp1.bias', - '^t_embedder.mlp.2.weight$': - 't_embedder.mlp2.weight', - '^t_embedder.mlp.2.bias$': - 't_embedder.mlp2.bias', - '^t_embedder.mlp.0.weights_scaling_factor$': - 't_embedder.mlp1.weights_scaling_factor', - '^t_embedder.mlp.0.activation_scaling_factor$': - 't_embedder.mlp1.activation_scaling_factor', - '^t_embedder.mlp.2.weights_scaling_factor$': - 't_embedder.mlp2.weights_scaling_factor', - '^t_embedder.mlp.2.activation_scaling_factor$': - 't_embedder.mlp2.activation_scaling_factor', - - # Add negative lookhead for scaling matching. - r'^blocks.(\d+).mlp.fc1.weight$': - 'blocks.*.mlp.fc.weight', - r'^blocks.(\d+).mlp.fc1.bias$': - 'blocks.*.mlp.fc.bias', - r'^blocks.(\d+).mlp.fc2.weight$': - 'blocks.*.mlp.proj.weight', - r'^blocks.(\d+).mlp.fc2.bias$': - 'blocks.*.mlp.proj.bias', - r'^blocks.(\d+).mlp.fc1.weights_scaling_factor$': - 'blocks.*.mlp.fc.weights_scaling_factor', - r'^blocks.(\d+).mlp.fc1.activation_scaling_factor$': - 'blocks.*.mlp.fc.activation_scaling_factor', - r'^blocks.(\d+).mlp.fc2.weights_scaling_factor$': - 'blocks.*.mlp.proj.weights_scaling_factor', - r'^blocks.(\d+).mlp.fc2.activation_scaling_factor$': - 'blocks.*.mlp.proj.activation_scaling_factor', - - # Add negative lookhead for scaling matching. - r'^blocks.(\d+).attn.proj.weight$': - 'blocks.*.attn.dense.weight', - r'^blocks.(\d+).attn.proj.bias$': - 'blocks.*.attn.dense.bias', - r'^blocks.(\d+).attn.proj.weights_scaling_factor$': - 'blocks.*.attn.dense.weights_scaling_factor', - r'^blocks.(\d+).attn.proj.activation_scaling_factor$': - 'blocks.*.attn.dense.activation_scaling_factor', - r'^blocks.(\d+).adaLN_modulation.1.weight$': - 'blocks.*.adaLN_modulation.weight', - r'^blocks.(\d+).adaLN_modulation.1.bias$': - 'blocks.*.adaLN_modulation.bias', - r'^blocks.(\d+).adaLN_modulation.1.weights_scaling_factor$': - 'blocks.*.adaLN_modulation.weights_scaling_factor', - r'^blocks.(\d+).adaLN_modulation.1.activation_scaling_factor$': - 'blocks.*.adaLN_modulation.activation_scaling_factor', - '^final_layer.adaLN_modulation.1.weight$': - 'final_layer.adaLN_modulation.weight', - '^final_layer.adaLN_modulation.1.bias$': - 'final_layer.adaLN_modulation.bias', - '^final_layer.adaLN_modulation.1.weights_scaling_factor$': - 'final_layer.adaLN_modulation.weights_scaling_factor', - '^final_layer.adaLN_modulation.1.activation_scaling_factor$': - 'final_layer.adaLN_modulation.activation_scaling_factor', -} - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--timm_ckpt', - type=str, - default="./DiT-XL-2-512x512.pt") - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument('--input_size', - type=int, - default=64, - help='The input latent size') - parser.add_argument('--patch_size', - type=int, - default=2, - help='The patch size for patchify') - parser.add_argument('--in_channels', - type=int, - default=4, - help='The channels of input latent') - parser.add_argument('--hidden_size', - type=int, - default=1152, - help='The hidden size of DiT') - parser.add_argument('--depth', - type=int, - default=28, - help='The number of DiTBlock layers') - parser.add_argument('--num_heads', - type=int, - default=16, - help='The number of heads of attention module') - parser.add_argument( - '--mlp_ratio', - type=float, - default=4.0, - help= - 'The ratio of hidden size compared to input hidden size in MLP layer') - parser.add_argument( - '--class_dropout_prob', - type=float, - default=0.1, - help='The probability to drop class token when training') - parser.add_argument('--num_classes', - type=int, - default=1000, - help='The number of classes for conditional control') - parser.add_argument('--learn_sigma', - type=bool, - default=True, - help='Whether the model learn sigma') - parser.add_argument('--cfg_scale', type=float, default=4.0) - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--cp_size', - type=int, - default=1, - help='Context parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument('--fp8_linear', - action='store_true', - help='Whether use FP8 for linear layers') - parser.add_argument( - '--diffusers_dit', - action='store_true', - help='Convert checkpoint provided by `HuggingFace/diffusers`') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - args = parser.parse_args() - return args - - -def convert_timm_dit(args, mapping, dtype='float32'): - - weights = {} - tik = time.time() - torch_dtype = str_dtype_to_torch(dtype) - tensor_parallel = mapping.tp_size - if args.diffusers_dit and args.fp8_linear: - from utils_modelopt import remap_model - converted_ckpt = "transformer.fp8.converted.state_dict.pt" - remap_model(quantized_ckpt=args.timm_ckpt, output_ckpt=converted_ckpt) - model_params = dict(torch.load(converted_ckpt)) - else: - model_params = dict(torch.load(args.timm_ckpt)) - timm_to_trtllm_name = FACEBOOK_DIT_NAME_MAPPING - - def get_trtllm_name(timm_name): - for k, v in timm_to_trtllm_name.items(): - m = re.match(k, timm_name) - if m is not None: - if "*" in v: - v = v.replace("*", m.groups()[0]) - return v - return timm_name - - weights = dict() - for name, param in model_params.items(): - if param.dtype in [torch.int8, torch.float8_e4m3fn - ] or 'scaling_factor' in name: - if 'scaling_factor' in name: - assert param.dtype == torch.float32 - weights[get_trtllm_name(name)] = param.contiguous() - else: - weights[get_trtllm_name(name)] = param.contiguous().to(torch_dtype) - - assert len(weights) == len(model_params) - - for k, v in weights.items(): - if re.match('^blocks.*.attn.qkv.weight$', k): - weights[k] = split_qkv_tp(v, args.num_heads, args.hidden_size, - tensor_parallel, mapping.tp_rank) - elif re.match('^blocks.*.attn.qkv.bias$', k): - weights[k] = split_qkv_bias_tp(v, args.num_heads, args.hidden_size, - tensor_parallel, mapping.tp_rank) - elif re.match('^blocks.*.attn.dense.weight$', k): - weights[k] = split_matrix_tp(v, - tensor_parallel, - mapping.tp_rank, - dim=1) - elif re.match('^blocks.*.mlp.fc.weight$', k): - weights[k] = split_matrix_tp(v, - tensor_parallel, - mapping.tp_rank, - dim=0) - elif re.match('^blocks.*.mlp.fc.bias$', k): - weights[k] = split(v, tensor_parallel, mapping.tp_rank) - elif re.match('^blocks.*.mlp.proj.weight$', k): - weights[k] = split_matrix_tp(v, - tensor_parallel, - mapping.tp_rank, - dim=1) - elif re.match(r'.*adaLN_modulation.weight$', k): - weights[k] = split_matrix_tp(v, - tensor_parallel, - mapping.tp_rank, - dim=0) - elif re.match(r'.*adaLN_modulation.bias$', k): - weights[k] = split(v, tensor_parallel, mapping.tp_rank) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def save_config(args): - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - config = { - 'architecture': "DiT", - 'dtype': args.dtype, - 'input_size': args.input_size, - 'patch_size': args.patch_size, - 'in_channels': args.in_channels, - 'hidden_size': args.hidden_size, - 'num_hidden_layers': args.depth, - 'num_attention_heads': args.num_heads, - 'mlp_ratio': args.mlp_ratio, - 'class_dropout_prob': args.class_dropout_prob, - 'num_classes': args.num_classes, - 'learn_sigma': args.learn_sigma, - 'cfg_scale': args.cfg_scale, - 'mapping': { - 'world_size': args.cp_size * args.tp_size * args.pp_size, - 'cp_size': args.cp_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - } - } - if args.fp8_linear: - config['quantization'] = { - 'quant_algo': "FP8", - # TODO: add support for exclude modules. - # 'exclude_modules': ["*final_layer*"], - } - - with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: - json.dump(config, f, indent=4) - - -def covert_and_save(args, rank): - if rank == 0: - save_config(args) - - mapping = Mapping(world_size=args.cp_size * args.tp_size * args.pp_size, - rank=rank, - cp_size=args.cp_size, - tp_size=args.tp_size, - pp_size=args.pp_size) - - weights = convert_timm_dit(args, mapping, dtype=args.dtype) - - safetensors.torch.save_file( - weights, os.path.join(args.output_dir, f'rank{rank}.safetensors')) - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - world_size = args.cp_size * args.tp_size * args.pp_size - - assert args.pp_size == 1, "PP is not supported yet." - - tik = time.time() - - if args.timm_ckpt is None: - return - - execute(args.workers, [covert_and_save] * world_size, args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/dit/diffusion.py b/examples/models/contrib/dit/diffusion.py deleted file mode 100644 index 510001b2ede2..000000000000 --- a/examples/models/contrib/dit/diffusion.py +++ /dev/null @@ -1,116 +0,0 @@ -import numpy as np -import torch -from tqdm.auto import tqdm - - -def space_timesteps(num_timesteps, timestep_respacing): - if num_timesteps < timestep_respacing: - raise ValueError( - f"cannot divide section of {num_timesteps} steps into {timestep_respacing}" - ) - frac_stride = (num_timesteps - 1) / (timestep_respacing - 1) - cur_idx = 0.0 - taken_steps = [] - for _ in range(timestep_respacing): - taken_steps.append(round(cur_idx)) - cur_idx += frac_stride - return set(taken_steps) - - -def _extract_into_tensor(arr, timesteps, broadcast_shape): - res = torch.from_numpy(arr).to(device=timesteps.device)[timesteps].float() - while len(res.shape) < len(broadcast_shape): - res = res[..., None] - return res + torch.zeros(broadcast_shape, device=timesteps.device) - - -class DiTDiffusionPipeline: - - def __init__(self, dit_model, timestep_respacing, diffusion_steps=1000): - self.dit_model = dit_model - self.timesteps = space_timesteps(diffusion_steps, timestep_respacing) - self.timestep_map = [] - self.original_num_steps = diffusion_steps - betas = self._setup_betas() - self.betas = betas - self.num_timesteps = int(betas.shape[0]) - - alphas = 1.0 - betas - self.alphas_cumprod = np.cumprod(alphas, axis=0) - self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1]) - self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0) - assert self.alphas_cumprod_prev.shape == (self.num_timesteps, ) - - self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod) - self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - - 1) - - self.posterior_variance = (betas * (1.0 - self.alphas_cumprod_prev) / - (1.0 - self.alphas_cumprod)) - - self.posterior_log_variance_clipped = np.log( - np.append(self.posterior_variance[1], - self.posterior_variance[1:])) if len( - self.posterior_variance) > 1 else np.array([]) - - self.posterior_mean_coef1 = (betas * np.sqrt(self.alphas_cumprod_prev) / - (1.0 - self.alphas_cumprod)) - self.posterior_mean_coef2 = ((1.0 - self.alphas_cumprod_prev) * - np.sqrt(alphas) / - (1.0 - self.alphas_cumprod)) - - def _setup_betas(self): - scale = 1000 / self.original_num_steps - betas = np.linspace(scale * 0.0001, - scale * 0.02, - self.original_num_steps, - dtype=np.float64) - last_alpha_cumprod = 1.0 - new_betas = [] - alphas = 1.0 - betas - alphas_cumprod = np.cumprod(alphas, axis=0) - for i, alpha_cumprod in enumerate(alphas_cumprod): - if i in self.timesteps: - new_betas.append(1 - alpha_cumprod / last_alpha_cumprod) - last_alpha_cumprod = alpha_cumprod - self.timestep_map.append(i) - betas = np.array(new_betas) - return betas - - def run(self, sample, labels): - indices = list(range(self.num_timesteps))[::-1] - - for i in tqdm(indices): - t = torch.tensor([i] * sample.shape[0], device=sample.device) - B, C = sample.shape[:2] - assert t.shape == (B, ) - map_tensor = torch.tensor(self.timestep_map, - device=t.device, - dtype=t.dtype) - model_output = self.dit_model(sample, map_tensor[t], labels) - assert model_output.shape == (B, C * 2, *sample.shape[2:]) - model_output, model_var_values = torch.split(model_output, C, dim=1) - min_log = _extract_into_tensor(self.posterior_log_variance_clipped, - t, sample.shape) - max_log = _extract_into_tensor(np.log(self.betas), t, sample.shape) - - frac = (model_var_values + 1) / 2 - model_log_variance = frac * max_log + (1 - frac) * min_log - pred_xstart = ( - _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, - sample.shape) * sample - - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, - sample.shape) * model_output) - - model_mean = (_extract_into_tensor(self.posterior_mean_coef1, t, - sample.shape) * pred_xstart + - _extract_into_tensor(self.posterior_mean_coef2, t, - sample.shape) * sample) - - assert model_mean.shape == model_log_variance.shape == pred_xstart.shape == sample.shape - noise = torch.randn_like(sample) - nonzero_mask = ((t != 0).float().view( - -1, *([1] * (len(sample.shape) - 1)))) - sample = model_mean + nonzero_mask * torch.exp( - 0.5 * model_log_variance) * noise - return sample diff --git a/examples/models/contrib/dit/figs/sample.fp8.png b/examples/models/contrib/dit/figs/sample.fp8.png deleted file mode 100644 index 333318aaa4de..000000000000 Binary files a/examples/models/contrib/dit/figs/sample.fp8.png and /dev/null differ diff --git a/examples/models/contrib/dit/figs/sample.png b/examples/models/contrib/dit/figs/sample.png deleted file mode 100644 index ba3293298600..000000000000 Binary files a/examples/models/contrib/dit/figs/sample.png and /dev/null differ diff --git a/examples/models/contrib/dit/sample.py b/examples/models/contrib/dit/sample.py deleted file mode 100644 index 2b4b7aebb2bc..000000000000 --- a/examples/models/contrib/dit/sample.py +++ /dev/null @@ -1,268 +0,0 @@ -import argparse -import json -import os -from functools import wraps - -import tensorrt as trt -import torch -from cuda import cudart -from diffusion import DiTDiffusionPipeline -from torchvision.utils import save_image - -import tensorrt_llm -from tensorrt_llm._utils import str_dtype_to_torch, trt_dtype_to_torch -from tensorrt_llm.logger import logger -from tensorrt_llm.plugin.plugin import CustomAllReduceHelper -from tensorrt_llm.runtime.session import Session, TensorInfo - - -def CUASSERT(cuda_ret): - err = cuda_ret[0] - if err != cudart.cudaError_t.cudaSuccess: - raise RuntimeError( - f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t" - ) - if len(cuda_ret) > 1: - return cuda_ret[1:] - return None - - -class TllmDiT(object): - - def __init__(self, - config, - debug_mode=True, - stream: torch.cuda.Stream = None): - self.dtype = config['pretrained_config']['dtype'] - - rank = tensorrt_llm.mpi_rank() - world_size = config['pretrained_config']['mapping']['world_size'] - cp_size = config['pretrained_config']['mapping']['cp_size'] - tp_size = config['pretrained_config']['mapping']['tp_size'] - pp_size = config['pretrained_config']['mapping']['pp_size'] - assert pp_size == 1 - self.mapping = tensorrt_llm.Mapping(world_size=world_size, - rank=rank, - cp_size=cp_size, - tp_size=tp_size, - pp_size=1, - gpus_per_node=args.gpus_per_node) - - local_rank = rank % self.mapping.gpus_per_node - self.device = torch.device(f'cuda:{local_rank}') - torch.cuda.set_device(self.device) - CUASSERT(cudart.cudaSetDevice(local_rank)) - - self.stream = stream - if self.stream is None: - self.stream = torch.cuda.Stream(self.device) - torch.cuda.set_stream(self.stream) - - engine_file = os.path.join(args.tllm_model_dir, f"rank{rank}.engine") - logger.info(f'Loading engine from {engine_file}') - with open(engine_file, "rb") as f: - engine_buffer = f.read() - - assert engine_buffer is not None - - self.session = Session.from_serialized_engine(engine_buffer) - - self.debug_mode = debug_mode - - self.inputs = {} - self.outputs = {} - self.buffer_allocated = False - - expected_tensor_names = ['latent', 'timestep', 'label', 'output'] - - if self.mapping.tp_size > 1: - self.buffer, self.all_reduce_workspace = CustomAllReduceHelper.allocate_workspace( - self.mapping, - CustomAllReduceHelper.max_workspace_size_auto( - self.mapping.tp_size)) - self.inputs['all_reduce_workspace'] = self.all_reduce_workspace - expected_tensor_names += ['all_reduce_workspace'] - - found_tensor_names = [ - self.session.engine.get_tensor_name(i) - for i in range(self.session.engine.num_io_tensors) - ] - if not self.debug_mode and set(expected_tensor_names) != set( - found_tensor_names): - logger.error( - f"The following expected tensors are not found: {set(expected_tensor_names).difference(set(found_tensor_names))}" - ) - logger.error( - f"Those tensors in engine are not expected: {set(found_tensor_names).difference(set(expected_tensor_names))}" - ) - logger.error(f"Expected tensor names: {expected_tensor_names}") - logger.error(f"Found tensor names: {found_tensor_names}") - raise RuntimeError( - "Tensor names in engine are not the same as expected.") - if self.debug_mode: - self.debug_tensors = list( - set(found_tensor_names) - set(expected_tensor_names)) - - def _tensor_dtype(self, name): - # return torch dtype given tensor name for convenience - dtype = trt_dtype_to_torch(self.session.engine.get_tensor_dtype(name)) - return dtype - - def _setup(self, batch_size): - for i in range(self.session.engine.num_io_tensors): - name = self.session.engine.get_tensor_name(i) - if self.session.engine.get_tensor_mode( - name) == trt.TensorIOMode.OUTPUT: - shape = list(self.session.engine.get_tensor_shape(name)) - shape[0] = batch_size // 2 if name in [ - 'cond_eps', 'uncond_eps' - ] else batch_size - self.outputs[name] = torch.empty(shape, - dtype=self._tensor_dtype(name), - device=self.device) - - self.buffer_allocated = True - - def cuda_stream_guard(func): - """Sync external stream and set current stream to the one bound to the session. Reset on exit. - """ - - @wraps(func) - def wrapper(self, *args, **kwargs): - external_stream = torch.cuda.current_stream() - if external_stream != self.stream: - external_stream.synchronize() - torch.cuda.set_stream(self.stream) - ret = func(self, *args, **kwargs) - if external_stream != self.stream: - self.stream.synchronize() - torch.cuda.set_stream(external_stream) - return ret - - return wrapper - - @cuda_stream_guard - def forward(self, latent: torch.Tensor, timestep: torch.Tensor, - label: torch.Tensor): - """ - Forward pass of DiT. - latent: (N, C, H, W) - timestep: (N,) - label: (N,) - """ - self._setup(latent.shape[0]) - if not self.buffer_allocated: - raise RuntimeError('Buffer not allocated, please call setup first!') - - inputs = { - 'latent': latent.to(str_dtype_to_torch(self.dtype)), - "timestep": timestep.int(), - "label": label.int() - } - self.inputs.update(**inputs) - self.session.set_shapes(self.inputs) - ok = self.session.run(self.inputs, self.outputs, - self.stream.cuda_stream) - - if not ok: - raise RuntimeError('Executing TRT engine failed!') - if self.debug_mode: - torch.cuda.synchronize() - for k, v in self.inputs.items(): - print(k, v.sum()) - for k, v in self.outputs.items(): - print(k, v.sum()) - return self.outputs['output'] - - -def vae_decode(samples, engine_path): - # Load standard plugins - TRT_LOGGER = trt.Logger(trt.Logger.WARNING) - trt.init_libnvinfer_plugins(TRT_LOGGER, namespace="") - - logger.info(f'Loading vae engine from {engine_path}') - with open(engine_path, 'rb') as f: - engine_buffer = f.read() - logger.info(f'Creating session from engine {engine_path}') - session_vae = Session.from_serialized_engine(engine_buffer) - inputs = {'input': samples} - output_info = session_vae.infer_shapes( - [TensorInfo('input', trt.DataType.FLOAT, samples.shape)]) - outputs = { - t.name: - torch.empty(tuple(t.shape), - dtype=trt_dtype_to_torch(t.dtype), - device='cuda') - for t in output_info - } - stream = torch.cuda.current_stream().cuda_stream - ok = session_vae.run(inputs, outputs, stream) - - assert ok, "Runtime execution failed for vae session" - - samples = outputs['output'] - return samples - - -def main(args): - tensorrt_llm.logger.set_level(args.log_level) - - torch.manual_seed(args.seed) - assert torch.cuda.is_available() - device = "cuda" - - # Load model: - config_file = os.path.join(args.tllm_model_dir, 'config.json') - with open(config_file) as f: - config = json.load(f) - model = TllmDiT(config, debug_mode=args.debug_mode) - - diffusion = DiTDiffusionPipeline(model.forward, - timestep_respacing=args.num_sampling_steps) - - latent_size = args.image_size // 8 - latent = torch.randn(args.batch_size, - 4, - latent_size, - latent_size, - device=device) - labels = torch.randint(args.num_classes, [args.batch_size], device=device) - - latent = torch.cat([latent, latent], 0) - labels_null = torch.tensor([1000] * args.batch_size, device=device) - labels = torch.cat([labels, labels_null], 0) - - samples = diffusion.run(latent, labels) - samples, _ = samples.chunk(2, dim=0) - - samples = vae_decode(samples / 0.18215, args.vae_decoder_engine) - - save_image(samples, - "sample.png", - nrow=4, - normalize=True, - value_range=(-1, 1)) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument('--vae_decoder_engine', - type=str, - default='vae_decoder/plan/visual_encoder_fp16.plan', - help='') - parser.add_argument("--batch-size", type=int, default=2) - parser.add_argument("--image-size", - type=int, - choices=[256, 512], - default=512) - parser.add_argument("--num-classes", type=int, default=1000) - parser.add_argument("--num-sampling-steps", type=int, default=250) - parser.add_argument("--seed", type=int, default=0) - parser.add_argument("--tllm_model_dir", - type=str, - default='./engine_outputs/') - parser.add_argument("--gpus_per_node", type=int, default=8) - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument("--debug_mode", type=bool, default=False) - args = parser.parse_args() - main(args) diff --git a/examples/models/contrib/dit/utils_modelopt.py b/examples/models/contrib/dit/utils_modelopt.py deleted file mode 100644 index e5f620e05325..000000000000 --- a/examples/models/contrib/dit/utils_modelopt.py +++ /dev/null @@ -1,186 +0,0 @@ -import os -import re -from collections import OrderedDict - -import modelopt.torch.opt as mto -import torch -from diffusers import DiTPipeline -from modelopt.torch.export.layer_utils import (get_activation_scaling_factor, - get_weight_scaling_factor) -from modelopt.torch.export.model_config_utils import to_quantized_weight -from torchvision.datasets.utils import download_url - -HUGGINGFACE_TO_FACEBOOK_DIT_NAME_MAPPING = { - r"^transformer_blocks.(\d+).norm1.emb.class_embedder.embedding_table.weight$": - "y_embedder.embedding_table.weight", - r"^transformer_blocks.(\d+).norm1.emb.timestep_embedder.linear_1.weight$": - "t_embedder.mlp.0.weight", - r"^transformer_blocks.(\d+).norm1.emb.timestep_embedder.linear_1.bias$": - "t_embedder.mlp.0.bias", - r"^transformer_blocks.(\d+).norm1.emb.timestep_embedder.linear_2.weight$": - "t_embedder.mlp.2.weight", - r"^transformer_blocks.(\d+).norm1.emb.timestep_embedder.linear_2.bias$": - "t_embedder.mlp.2.bias", - "^pos_embed.proj.weight$": - "x_embedder.proj.weight", - "^pos_embed.proj.bias$": - "x_embedder.proj.bias", - r"^transformer_blocks.(\d+).attn1.to_qkv.weight$": - "blocks.*.attn.qkv.weight", - r"^transformer_blocks.(\d+).attn1.to_qkv.bias$": - "blocks.*.attn.qkv.bias", - r"^transformer_blocks.(\d+).attn1.to_out.0.weight$": - "blocks.*.attn.proj.weight", - r"^transformer_blocks.(\d+).attn1.to_out.0.bias$": - "blocks.*.attn.proj.bias", - r"^transformer_blocks.(\d+).ff.net.0.proj.weight$": - "blocks.*.mlp.fc1.weight", - r"^transformer_blocks.(\d+).ff.net.0.proj.bias$": - "blocks.*.mlp.fc1.bias", - r"^transformer_blocks.(\d+).ff.net.2.weight$": - "blocks.*.mlp.fc2.weight", - r"^transformer_blocks.(\d+).ff.net.2.bias$": - "blocks.*.mlp.fc2.bias", - r"^transformer_blocks.(\d+).norm1.linear.weight$": - "blocks.*.adaLN_modulation.1.weight", - r"^transformer_blocks.(\d+).norm1.linear.bias$": - "blocks.*.adaLN_modulation.1.bias", - "^proj_out_2.weight$": - "final_layer.linear.weight", - "^proj_out_2.bias$": - "final_layer.linear.bias", - "^proj_out_1.weight$": - "final_layer.adaLN_modulation.1.weight", - "^proj_out_1.bias$": - "final_layer.adaLN_modulation.1.bias", - r"^transformer_blocks.(\d+).norm1.emb.timestep_embedder.linear_1.weights_scaling_factor$": - "t_embedder.mlp.0.weights_scaling_factor", - r"^transformer_blocks.(\d+).norm1.emb.timestep_embedder.linear_1.activation_scaling_factor": - "t_embedder.mlp.0.activation_scaling_factor", - r"^transformer_blocks.(\d+).norm1.emb.timestep_embedder.linear_2.weights_scaling_factor": - "t_embedder.mlp.2.weights_scaling_factor", - r"^transformer_blocks.(\d+).norm1.emb.timestep_embedder.linear_2.activation_scaling_factor": - "t_embedder.mlp.2.activation_scaling_factor", - r"^transformer_blocks.(\d+).attn1.to_qkv.weights_scaling_factor$": - "blocks.*.attn.qkv.weights_scaling_factor", - r"^transformer_blocks.(\d+).attn1.to_qkv.activation_scaling_factor$": - "blocks.*.attn.qkv.activation_scaling_factor", - r"^transformer_blocks.(\d+).attn1.to_out.0.weights_scaling_factor$": - "blocks.*.attn.proj.weights_scaling_factor", - r"^transformer_blocks.(\d+).attn1.to_out.0.activation_scaling_factor$": - "blocks.*.attn.proj.activation_scaling_factor", - r"^transformer_blocks.(\d+).ff.net.0.proj.weights_scaling_factor$": - "blocks.*.mlp.fc1.weights_scaling_factor", - r"^transformer_blocks.(\d+).ff.net.0.proj.activation_scaling_factor$": - "blocks.*.mlp.fc1.activation_scaling_factor", - r"^transformer_blocks.(\d+).ff.net.2.weights_scaling_factor$": - "blocks.*.mlp.fc2.weights_scaling_factor", - r"^transformer_blocks.(\d+).ff.net.2.activation_scaling_factor$": - "blocks.*.mlp.fc2.activation_scaling_factor", - r"^transformer_blocks.(\d+).norm1.linear.weights_scaling_factor$": - "blocks.*.adaLN_modulation.1.weights_scaling_factor", - r"^transformer_blocks.(\d+).norm1.linear.activation_scaling_factor$": - "blocks.*.adaLN_modulation.1.activation_scaling_factor", - "^proj_out_2.weights_scaling_factor$": - "final_layer.linear.weights_scaling_factor", - "^proj_out_2.activation_scaling_factor$": - "final_layer.linear.activation_scaling_factor", - "^proj_out_1.weights_scaling_factor$": - "final_layer.adaLN_modulation.1.weights_scaling_factor", - "^proj_out_1.activation_scaling_factor$": - "final_layer.adaLN_modulation.1.activation_scaling_factor", -} - - -def convert_amax_to_scaling_factor(model, state_dict): - ret_dict = state_dict.copy() - for name, module in model.named_modules(): - if isinstance(module, torch.nn.Linear): - activation_scaling_factor = get_activation_scaling_factor(module) - weight_scaling_factor = get_weight_scaling_factor(module) - if activation_scaling_factor: - ret_dict[ - f'{name}.activation_scaling_factor'] = activation_scaling_factor - if weight_scaling_factor: - ret_dict[ - f'{name}.weights_scaling_factor'] = weight_scaling_factor - - weight = module.weight.detach().cpu() - if weight_scaling_factor: - # only module with valid weight scaling factor - weight = to_quantized_weight( - weight=weight, - weights_scaling_factor=weight_scaling_factor, - quantization="fp8", - ) - # replace the quantized weight - ret_dict[f'{name}.weight'] = weight - return ret_dict - - -def get_weights_map(state_dict): - - def _get_fb_dit_name(dit_name): - for k, v in HUGGINGFACE_TO_FACEBOOK_DIT_NAME_MAPPING.items(): - m = re.match(k, dit_name) - if m is not None: - if "*" in v: - v = v.replace("*", m.groups()[0]) - return v - return dit_name - - weights_map = OrderedDict() - for key, value in state_dict.items(): - if ("to_q." in key) or ("to_k." in key) or ("to_v." in key): - continue - if _get_fb_dit_name(key) in weights_map: - continue - else: - weights_map[_get_fb_dit_name(key)] = value - return weights_map - - -def download_model(ckpt="DiT-XL-2-512x512.pt"): - """ - Downloads a pre-trained DiT model from the web. - """ - if not os.path.isfile(ckpt): - os.makedirs('pretrained_models', exist_ok=True) - web_path = f"https://dl.fbaipublicfiles.com/DiT/models/{ckpt}" - download_url(web_path, '.') - model = torch.load(ckpt, map_location=lambda storage, loc: storage) - return model - - -def remap_model(model_name="facebook/DiT-XL-2-512", - quantized_ckpt="dit.quantized.pt", - output_ckpt="dit.converted.pt", - fuse_qkv=True, - dtype=torch.float32): - pipe = DiTPipeline.from_pretrained(model_name, torch_dtype=dtype) - transformer = pipe.transformer - - # fuse qkv gemm - if fuse_qkv: - from diffusers.models.attention_processor import (Attention, - AttnProcessor2_0) - for module in transformer.modules(): - if isinstance(module, Attention): - assert (isinstance(module.processor, AttnProcessor2_0)) - module.fuse_projections(fuse=True) - pretrained_state_dict = transformer.state_dict() - - mto.restore(transformer, quantized_ckpt) - quantized_dict = convert_amax_to_scaling_factor(transformer, - pretrained_state_dict) - assert set(pretrained_state_dict.keys()).issubset(set( - quantized_dict.keys())) - - remapped_dict = get_weights_map(quantized_dict) - # TODO: currently we use weights of pos_embed and x_embedder from official DiT because - # in the implementation by HuggingFace, there are some modifications for these modules. - stored_params = download_model(ckpt="DiT-XL-2-512x512.pt") - for name in ["pos_embed", "x_embedder.proj.weight", "x_embedder.proj.bias"]: - remapped_dict[name] = stored_params[name] - - torch.save(remapped_dict, output_ckpt) diff --git a/examples/models/contrib/dit/vae_decoder_trt.py b/examples/models/contrib/dit/vae_decoder_trt.py deleted file mode 100644 index 31803a6690d7..000000000000 --- a/examples/models/contrib/dit/vae_decoder_trt.py +++ /dev/null @@ -1,139 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import argparse -import os - -import torch -from diffusers.models import AutoencoderKL - - -class TRT_Exporter(object): - - def __init__(self, pytorch_model, max_batch_size, latent_channel, - latent_shape): - self.pytorch_model = pytorch_model - self.max_batch_size = max_batch_size - self.latent_channel = latent_channel - self.latent_shape = latent_shape - - def export_onnx(self, onnxFile): - print(f"Start exporting ONNX model to {onnxFile}!") - latent = torch.randn(self.max_batch_size, self.latent_channel, - *self.latent_shape).cuda() - self.pytorch_model.cuda().eval() - with torch.inference_mode(): - torch.onnx.export( - self.pytorch_model, - latent, - onnxFile, - opset_version=17, - input_names=['input'], - output_names=['output'], - dynamic_axes={'input': { - 0: 'batch' - }}, - # Required for pytorch>=2.9.0 as dynamo becomes the default and introduces bugs as it does not support opset_version=17 natively - dynamo=False) - - def generate_trt_engine(self, onnxFile, planFile): - print(f"Start exporting TRT model to {planFile}!") - from time import time - - import tensorrt as trt - logger = trt.Logger(trt.Logger.VERBOSE) - builder = trt.Builder(logger) - network = builder.create_network( - 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) - profile = builder.create_optimization_profile() - config = builder.create_builder_config() - config.set_flag(trt.BuilderFlag.FP16) - parser = trt.OnnxParser(network, logger) - - with open(onnxFile, 'rb') as model: - if not parser.parse(model.read(), "/".join(onnxFile.split("/"))): - print("Failed parsing %s" % onnxFile) - for error in range(parser.num_errors): - print(parser.get_error(error)) - print("Succeeded parsing %s" % onnxFile) - - nBS = -1 - nMinBS = 1 - nMaxBS = self.max_batch_size - nOptBS = (nMaxBS + nMinBS) // 2 - inputT = network.get_input(0) - inputT.shape = [nBS, self.latent_channel, *self.latent_shape] - profile.set_shape(inputT.name, - [nMinBS, self.latent_channel, *self.latent_shape], - [nOptBS, self.latent_channel, *self.latent_shape], - [nMaxBS, self.latent_channel, *self.latent_shape]) - - config.add_optimization_profile(profile) - - t0 = time() - engineString = builder.build_serialized_network(network, config) - t1 = time() - if engineString is None: - print("Failed building %s" % planFile) - else: - print("Succeeded building %s in %d s" % (planFile, t1 - t0)) - print("plan file is", planFile) - with open(planFile, 'wb') as f: - f.write(engineString) - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument("--vae", - type=str, - choices=["ema", "mse"], - default="mse") - parser.add_argument('--max_batch_size', type=int, default=8) - parser.add_argument("--image-size", - type=int, - choices=[256, 512], - default=512) - parser.add_argument('--onnxFile', - type=str, - default='vae_decoder/onnx/visual_encoder.onnx', - help='') - parser.add_argument('--planFile', - type=str, - default='vae_decoder/plan/visual_encoder_fp16.plan', - help='') - parser.add_argument('--only_trt', - action='store_true', - help='Run only convert the onnx to TRT engine.') - args = parser.parse_args() - return args - - -if __name__ == '__main__': - args = parse_arguments() - onnx_file_dir = os.path.dirname(args.onnxFile) - if not onnx_file_dir == '' and not os.path.exists(onnx_file_dir): - os.makedirs(onnx_file_dir) - plan_file_dir = os.path.dirname(args.planFile) - if not os.path.exists(plan_file_dir): - os.makedirs(plan_file_dir) - - vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-{args.vae}") - vae.forward = vae.decode - latant_shape = [args.image_size // 8] * 2 - onnx_trt_obj = TRT_Exporter(vae, args.max_batch_size, 4, latant_shape) - if args.only_trt: - onnx_trt_obj.generate_trt_engine(args.onnxFile, args.planFile) - else: - onnx_trt_obj.export_onnx(args.onnxFile) - onnx_trt_obj.generate_trt_engine(args.onnxFile, args.planFile) diff --git a/examples/models/contrib/falcon/README.md b/examples/models/contrib/falcon/README.md deleted file mode 100644 index 3d77aa079cff..000000000000 --- a/examples/models/contrib/falcon/README.md +++ /dev/null @@ -1,364 +0,0 @@ -# Falcon - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a Falcon model in TensorRT LLM on single GPU, single node multi-GPU, and multi-node multi-GPU. - -- [Falcon](#falcon) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [1. Download weights from HuggingFace Transformers](#1-download-weights-from-huggingface-transformers) - - [2. Convert weights from HF Transformers to TensorRT LLM format](#2-convert-weights-from-hf-transformers-to-tensorrt-llm-format) - - [3. Build TensorRT engine(s)](#3-build-tensorrt-engines) - - [4. Run summarization task with the TensorRT engine(s)](#4-run-summarization-task-with-the-tensorrt-engines) - - [FP8 Post-Training Quantization](#fp8-post-training-quantization) - - [Groupwise quantization (AWQ)](#groupwise-quantization-awq) - - [W4A16 AWQ with FP8 GEMM (W4A8 AWQ)](#w4a16-awq-with-fp8-gemm-w4a8-awq) - - [Troubleshooting](#troubleshooting) - - [1. The HuggingFace Falcon may raise an error when using the `accelerate` package.](#1-the-huggingface-falcon-may-raise-an-error-when-using--the-accelerate-package) - -## 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/models/contrib/falcon`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 - * BF16 - * FP8 - * FP8 KV CACHE - * Groupwise quantization (AWQ) - * Tensor Parallel - * STRONGLY TYPED - -## Usage - -The next two sections describe how to convert the weights from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) -format to the TensorRT LLM format. - -### 1. Download weights from HuggingFace Transformers - -Install the dependency packages and setup `git-lfs`. - -```bash -# Install dependencies -pip install -r requirements.txt - -# Setup git-lfs -git lfs install -``` - -There are four HF checkpoints available. Use one of the following commands to fetch the checkpoint you are interested in. Follow the guides here https://huggingface.co/docs/transformers/main/en/model_doc/falcon. - -```bash -# falcon-rw-1b -git clone https://huggingface.co/tiiuae/falcon-rw-1b falcon/rw-1b - -# falcon-7b-instruct -git clone https://huggingface.co/tiiuae/falcon-7b-instruct falcon/7b-instruct - -# falcon-40b-instruct -git clone https://huggingface.co/tiiuae/falcon-40b-instruct falcon/40b-instruct - -# falcon-180b -git clone https://huggingface.co/tiiuae/falcon-180B falcon/180b - -# falcon-11b (Falcon 2) -git clone https://huggingface.co/tiiuae/falcon-11B falcon/11b -``` - -### 2. Convert weights from HF Transformers to TensorRT LLM format -The [`convert_checkpoint.py`](./convert_checkpoint.py) script converts HF weights to TensorRT LLM checkpoints. The number of checkpoint files (in .safetensors format) is same to the number of GPUs used to run inference. - -```bash -# falcon-rw-1b: single gpu, dtype float16 -python3 convert_checkpoint.py --model_dir ./falcon/rw-1b \ - --dtype float16 \ - --output_dir ./falcon/rw-1b/trt_ckpt/fp16/1-gpu/ - -# falcon-7b-instruct: single gpu, dtype bfloat16 -python3 convert_checkpoint.py --model_dir ./falcon/7b-instruct \ - --dtype bfloat16 \ - --output_dir ./falcon/7b-instruct/trt_ckpt/bf16/1-gpu/ - -# falcon-40b-instruct: 2-way tensor parallelism -python3 convert_checkpoint.py --model_dir ./falcon/40b-instruct \ - --dtype bfloat16 \ - --output_dir ./falcon/40b-instruct/trt_ckpt/bf16/tp2-pp1/ \ - --tp_size 2 - -# falcon-40b-instruct: 2-way tensor parallelism and 2-way pipeline parallelism -python3 convert_checkpoint.py --model_dir ./falcon/40b-instruct \ - --dtype bfloat16 \ - --output_dir ./falcon/40b-instruct/trt_ckpt/bf16/tp2-pp2/ \ - --tp_size 2 \ - --pp_size 2 - -# falcon-180b: 8-way tensor parallelism, loading weights shard-by-shard -python3 convert_checkpoint.py --model_dir ./falcon/180b \ - --dtype bfloat16 \ - --output_dir ./falcon/180b/trt_ckpt/bf16/tp8-pp1/ \ - --tp_size 8 \ - --load_by_shard \ - --workers 8 - -# falcon-180b: 4-way tensor parallelism and 2-way pipeline parallelism, loading weights shard-by-shard -python3 convert_checkpoint.py --model_dir ./falcon/180b \ - --dtype bfloat16 \ - --output_dir ./falcon/180b/trt_ckpt/bf16/tp4-pp2/ \ - --tp_size 4 \ - --pp_size 2 \ - --load_by_shard \ - --workers 8 - -# falcon-11b (Falcon 2): single gpu, dtype bfloat16 -python3 convert_checkpoint.py --model_dir ./falcon/11b \ - --dtype bfloat16 \ - --output_dir ./falcon/11b/trt_ckpt/bf16/1-gpu/ -``` - -Note that in order to use N-way tensor parallelism, the number of attention heads must be a multiple of N. -For example, you can't configure 2-way tensor parallelism for [falcon-7b](https://huggingface.co/tiiuae/falcon-7b) or [falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b-instruct), because the number of attention heads is 71 (not divisible by 2). - - -### 3. Build TensorRT engine(s) -The `trtllm-build` command builds TensorRT LLM engines from TensorRT LLM checkpoints. The number of engine files is also same to the number of GPUs used to run inference. - -Normally, the `trtllm-build` command only requires a single GPU, but you can enable parallel building by passing the number of GPUs to the `--workers` argument. - -```bash -# falcon-rw-1b -trtllm-build --checkpoint_dir ./falcon/rw-1b/trt_ckpt/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --output_dir ./falcon/rw-1b/trt_engines/fp16/1-gpu/ - -# falcon-7b-instruct -# Enabling --gpt_attention_plugin is necessary for rotary positional embedding (RoPE) -trtllm-build --checkpoint_dir ./falcon/7b-instruct/trt_ckpt/bf16/1-gpu/ \ - --gemm_plugin bfloat16 \ - --remove_input_padding enable \ - --gpt_attention_plugin bfloat16 \ - --output_dir ./falcon/7b-instruct/trt_engines/bf16/1-gpu/ - -# falcon-40b-instruct: 2-way tensor parallelism -trtllm-build --checkpoint_dir ./falcon/40b-instruct/trt_ckpt/bf16/tp2-pp1/ \ - --gemm_plugin bfloat16 \ - --gpt_attention_plugin bfloat16 \ - --output_dir ./falcon/40b-instruct/trt_engines/bf16/tp2-pp1/ - -# falcon-40b-instruct: 2-way tensor parallelism and 2-way pipeline parallelism -trtllm-build --checkpoint_dir ./falcon/40b-instruct/trt_ckpt/bf16/tp2-pp2/ \ - --gemm_plugin bfloat16 \ - --gpt_attention_plugin bfloat16 \ - --output_dir ./falcon/40b-instruct/trt_engines/bf16/tp2-pp2/ - -# falcon-180b: 8-way tensor parallelism -trtllm-build --checkpoint_dir ./falcon/180b/trt_ckpt/bf16/tp8-pp1/ \ - --gemm_plugin bfloat16 \ - --gpt_attention_plugin bfloat16 \ - --output_dir ./falcon/180b/trt_engines/bf16/tp8-pp1/ \ - --workers 8 - -# falcon-180b: 4-way tensor parallelism and 2-way pipeline parallelism -trtllm-build --checkpoint_dir ./falcon/180b/trt_ckpt/bf16/tp4-pp2/ \ - --gemm_plugin bfloat16 \ - --gpt_attention_plugin bfloat16 \ - --output_dir ./falcon/180b/trt_engines/bf16/tp4-pp2/ \ - --workers 8 - -# falcon-11b (Falcon 2) -trtllm-build --checkpoint_dir ./falcon/11b/trt_ckpt/bf16/1-gpu/ \ - --gemm_plugin bfloat16 \ - --gpt_attention_plugin bfloat16 \ - --output_dir ./falcon/11b/trt_engines/bf16/1-gpu/ -``` - -If the engines are built successfully, you will see output like (falcon-rw-1b as the example): -``` -...... -[12/27/2023-03:46:29] [TRT] [I] Engine generation completed in 35.0677 seconds. -[12/27/2023-03:46:29] [TRT] [I] [MemUsageStats] Peak memory usage of TRT CPU/GPU memory allocators: CPU 393 MiB, GPU 2699 MiB -[12/27/2023-03:46:29] [TRT] [I] [MemUsageChange] TensorRT-managed allocation in building engine: CPU +0, GPU +2699, now: CPU 0, GPU 2699 (MiB) -[12/27/2023-03:46:29] [TRT] [I] [MemUsageStats] Peak memory usage during Engine building and serialization: CPU: 10624 MiB -[12/27/2023-03:46:29] [TRT-LLM] [I] Total time of building Unnamed Network 0: 00:00:36 -[12/27/2023-03:46:31] [TRT-LLM] [I] Serializing engine to ./falcon/rw-1b/trt_engines/fp16/1-gpu/rank0.engine... -[12/27/2023-03:46:59] [TRT-LLM] [I] Engine serialized. Total time: 00:00:28 -[12/27/2023-03:46:59] [TRT-LLM] [I] Total time of building all engines: 00:01:59 -``` - -### 4. Run summarization task with the TensorRT engine(s) -The `../../../summarize.py` script can run the built engines to summarize the articles from the -[cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -```bash -# falcon-rw-1b -python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/rw-1b \ - --engine_dir ./falcon/rw-1b/trt_engines/fp16/1-gpu/ - -# falcon-7b-instruct -python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/7b-instruct \ - --engine_dir ./falcon/7b-instruct/trt_engines/bf16/1-gpu/ - -# falcon-40b-instruct: 2-way tensor parallelism -mpirun -n 2 --allow-run-as-root --oversubscribe \ - python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/40b-instruct \ - --engine_dir ./falcon/40b-instruct/trt_engines/bf16/tp2-pp1/ - -# falcon-40b-instruct: 2-way tensor parallelism and 2-way pipeline parallelism -mpirun -n 4 --allow-run-as-root --oversubscribe \ - python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/40b-instruct \ - --engine_dir ./falcon/40b-instruct/trt_engines/bf16/tp2-pp2/ - -# falcon-180b: 8-way tensor parallelism -mpirun -n 8 --allow-run-as-root --oversubscribe \ - python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/180b \ - --engine_dir ./falcon/180b/trt_engines/bf16/tp8-pp1/ - -# falcon-180b: 4-way tensor parallelism and 2-way pipeline parallelism -mpirun -n 8 --allow-run-as-root --oversubscribe \ - python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/180b \ - --engine_dir ./falcon/180b/trt_engines/bf16/tp4-pp2/ - -# falcon-11b (Falcon 2) -python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/11b \ - --engine_dir ./falcon/11b/trt_engines/bf16/1-gpu/ -``` - -If the engines are run successfully, you will see output like (falcon-rw-1b as the example): -``` -...... -[12/27/2023-03:57:02] [TRT-LLM] [I] TensorRT LLM (total latency: 5.816917419433594 sec) -[12/27/2023-03:57:02] [TRT-LLM] [I] TensorRT LLM beam 0 result -[12/27/2023-03:57:02] [TRT-LLM] [I] rouge1 : 15.061493342516243 -[12/27/2023-03:57:02] [TRT-LLM] [I] rouge2 : 4.495335888974063 -[12/27/2023-03:57:02] [TRT-LLM] [I] rougeL : 11.800002670828547 -[12/27/2023-03:57:02] [TRT-LLM] [I] rougeLsum : 13.458777656925877 -``` - -### FP8 Post-Training Quantization - -The examples below use the NVIDIA Modelopt (AlgorithMic Model Optimization) toolkit for the model quantization process. - -First make sure Modelopt toolkit is installed (see [examples/quantization/README.md](/examples/quantization/README.md#preparation)) - -Now quantize HF Falcon weights and export trtllm checkpoint. - -```bash -# Quantize HF Falcon 180B checkpoint into FP8 and export trtllm checkpoint -python ../../../quantization/quantize.py --model_dir ./falcon/180b \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir ./falcon/180b/trt_ckpt/fp8/tp8-pp1 \ - --tp_size 8 - -# Build trtllm engines from the trtllm checkpoint -trtllm-build --checkpoint_dir ./falcon/180b/trt_ckpt/fp8/tp8-pp1 \ - --gemm_plugin float16 \ - --output_dir ./falcon/180b/trt_engines/fp8/tp8-pp1 \ - --workers 8 - -# Run the summarization task -mpirun -n 8 --allow-run-as-root --oversubscribe \ - python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/180b \ - --engine_dir ./falcon/180b/trt_engines/fp8/tp8-pp1 -``` - -Note that you can enable fp8 context fmha to get further acceleration by setting `--use_fp8_context_fmha enable` when building the engines. - -### Groupwise quantization (AWQ) - -The examples below use the NVIDIA Modelopt (AlgorithMic Model Optimization) toolkit for the model quantization process. - -First make sure Modelopt toolkit is installed (see [examples/quantization/README.md](/examples/quantization/README.md#preparation)) - -Now quantize HF Falcon weights and export trtllm checkpoint. - -```bash -# Quantize HF Falcon 180B checkpoint into INT4-AWQ and export trtllm checkpoint -python ../../../quantization/quantize.py --model_dir ./falcon/180b \ - --dtype float16 \ - --qformat int4_awq \ - --output_dir ./falcon/180b/trt_ckpt/int4_awq/tp2 \ - --tp_size 2 - -# Build trtllm engines from the trtllm checkpoint -trtllm-build --checkpoint_dir ./falcon/180b/trt_ckpt/int4_awq/tp2 \ - --gemm_plugin float16 \ - --output_dir ./falcon/180b/trt_engines/int4_awq/tp2 \ - --workers 2 - -# Run the summarization task -mpirun -n 2 --allow-run-as-root --oversubscribe \ - python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/180b \ - --engine_dir ./falcon/180b/trt_engines/int4_awq/tp2 -``` - -#### W4A16 AWQ with FP8 GEMM (W4A8 AWQ) -For Hopper GPUs, TRT-LLM also supports employing FP8 GEMM for accelerating linear layers. This mode is noted with `w4a8_awq` for Modelopt and TRT-LLM, in which both weights and activations are converted from W4A16 to FP8 for GEMM calculation. - -Please make sure your system contains a Hopper GPU before trying the commands below. - -```bash -# Quantize HF Falcon 180B checkpoint into W4A8-AWQ and export trtllm checkpoint -python ../../../quantization/quantize.py --model_dir ./falcon/180b \ - --dtype float16 \ - --qformat w4a8_awq \ - --output_dir ./falcon/180b/trt_ckpt/w4a8_awq/tp2 \ - --tp_size 2 - -# Build trtllm engines from the trtllm checkpoint -trtllm-build --checkpoint_dir ./falcon/180b/trt_ckpt/w4a8_awq/tp2 \ - --gemm_plugin float16 \ - --output_dir ./falcon/180b/trt_engines/w4a8_awq/tp2 \ - --workers 2 - -# Run the summarization task -mpirun -n 2 --allow-run-as-root --oversubscribe \ - python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/180b \ - --engine_dir ./falcon/180b/trt_engines/w4a8_awq/tp2 -``` - -## Troubleshooting - -### 1. The HuggingFace Falcon may raise an error when using the `accelerate` package. - -One may find the following message. -``` -Traceback (most recent call last): - File "build.py", line 10, in - from transformers import FalconConfig, FalconForCausalLM - File "", line 1039, in _handle_fromlist - File "/usr/local/lib/python3.8/dist-packages/transformers/utils/import_utils.py", line 1090, in __getattr__ - value = getattr(module, name) - File "/usr/local/lib/python3.8/dist-packages/transformers/utils/import_utils.py", line 1089, in __getattr__ - module = self._get_module(self._class_to_module[name]) - File "/usr/local/lib/python3.8/dist-packages/transformers/utils/import_utils.py", line 1101, in _get_module - raise RuntimeError( -RuntimeError: Failed to import transformers.models.falcon.modeling_falcon because of the following error (look up to see its traceback): -``` -It may be resolved by pinning the version of `typing-extensions` package by `4.5.0`. -```bash -pip install typing-extensions==4.5.0 -``` diff --git a/examples/models/contrib/falcon/convert_checkpoint.py b/examples/models/contrib/falcon/convert_checkpoint.py deleted file mode 100644 index ac59316519eb..000000000000 --- a/examples/models/contrib/falcon/convert_checkpoint.py +++ /dev/null @@ -1,202 +0,0 @@ -import argparse -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.falcon.model import FalconForCausalLM -from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str) - 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( - '--dtype', - type=str, - default='auto', - choices=['auto', 'float16', 'bfloat16', 'float32'], - help= - "The data type for the model weights and activations if not quantized. " - "If 'auto', the data type is automatically inferred from the source model; " - "however, if the source dtype is float32, it is converted to float16.") - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - - parser.add_argument( - '--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('--load_by_shard', - action='store_true', - help='Load a pretrained model shard-by-shard.') - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument('--log_level', type=str, default='info') - args = parser.parse_args() - - tensorrt_llm.logger.set_level(args.log_level) - - # WAR for modelopt multithreading issue. - import importlib.metadata - if (importlib.metadata.version('nvidia-modelopt') < '0.27' - and args.workers > 1): - args.workers = 1 - tensorrt_llm.logger.warning( - 'Reducing workers=1 when converting checkpoint because ' - 'modelopt has an issue in multi-threading, which will be ' - 'fixed in 0.27.') - - return args - - -def args_to_quant_config(args: argparse.Namespace) -> QuantConfig: - config = QuantConfig() - if args.use_weight_only and args.weight_only_precision == 'int8': - config.quant_algo = QuantAlgo.W8A16 - elif args.use_weight_only and args.weight_only_precision == 'int4': - config.quant_algo = QuantAlgo.W4A16 - return config - - -def args_to_build_options(args): - return { - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - } - - -def convert_and_save_hf(args: argparse.Namespace): - model_dir = args.model_dir - load_by_shard = args.load_by_shard - world_size = args.tp_size * args.pp_size - # Need to convert the cli args to the kay-value pairs and override them in the generate config dict. - # Ideally these fields will be moved out of the config and pass them into build API, keep them here for compatibility purpose for now, - # before the refactor is done. - override_fields = {} - override_fields.update(args_to_build_options(args)) - - quant_config = args_to_quant_config(args) - hf_model = None - import transformers - if not args.load_by_shard and quant_config.quant_mode.has_any_quant(): - hf_model = transformers.FalconForCausalLM.from_pretrained( - model_dir, trust_remote_code=True, dtype='auto', device_map='auto') - else: - # Initialize huggingface local cache. - # Huggingface copies the external configuration source (`configuration_falcon.py` here) into its local cache at - # `~/.cache/huggingface/modules/transformers_modules/`, - # and if multiple threads attempt to do this concurrently, weird issue can happen: - # Some threads may see an empty configuration_falcon.py file and fail. - # Preload the config once to initialize local cache, so subsequent multithread loading won't fail. - _ = transformers.FalconConfig.from_pretrained(model_dir, - trust_remote_code=True, - dtype='auto', - device_map='auto') - - def convert_and_save_rank(args, rank: int): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - falcon = FalconForCausalLM.from_hugging_face( - model_dir if hf_model is None else hf_model, - dtype=args.dtype, - mapping=mapping, - quant_config=quant_config, - load_by_shard=load_by_shard, - **override_fields, - ) - falcon.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del falcon - - execute(args.workers, [convert_and_save_rank] * world_size, args) - release_gc() - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - convert_and_save_hf(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/falcon/requirements.txt b/examples/models/contrib/falcon/requirements.txt deleted file mode 100644 index e69190e50443..000000000000 --- a/examples/models/contrib/falcon/requirements.txt +++ /dev/null @@ -1,8 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -transformers>=4.56.0 -datasets==3.1.0 -evaluate -rouge_score -sentencepiece>=0.1.99 -tqdm diff --git a/examples/models/contrib/gptj/.gitignore b/examples/models/contrib/gptj/.gitignore deleted file mode 100644 index 18404db569c0..000000000000 --- a/examples/models/contrib/gptj/.gitignore +++ /dev/null @@ -1,4 +0,0 @@ -__pycache__/ -gptj_model/ -*.log -*.json diff --git a/examples/models/contrib/gptj/README.md b/examples/models/contrib/gptj/README.md deleted file mode 100644 index c470109a95f7..000000000000 --- a/examples/models/contrib/gptj/README.md +++ /dev/null @@ -1,269 +0,0 @@ -# GPT-J - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document explains how to build the [GPT-J](https://huggingface.co/EleutherAI/gpt-j-6b) model using TensorRT LLM and run on a single GPU. - -- [GPT-J](#gpt-j) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [1. Download weights from HuggingFace (HF) Transformers](#1-download-weights-from-huggingface-hf-transformers) - - [2. Build TensorRT engine(s)](#2-build-tensorrt-engines) - - [FP8 Post-Training Quantization](#fp8-post-training-quantization) - - [AWQ INT4 weight only quantization](#awq-int4-weight-only-quantization) - - [SmoothQuant (W8A8) quantization](#smoothquant-w8a8-quantization) - - [Fused MultiHead Attention (FMHA)](#fused-multihead-attention-fmha) - - [INT8 KV cache](#int8-kv-cache) - - [3. Run](#3-run) - - [Summarization using the GPT-J model](#summarization-using-the-gpt-j-model) - -## 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/models/contrib/gptj`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 - * INT8 & INT4 per-channel weight-only - * FP8 (with FP8 kv cache) - * Groupwise quantization (AWQ) - * INT8 KV CACHE (+ AWQ/per-channel weight-only) - * INT8 SmoothQuant - -## Usage - -### 1. Download weights from HuggingFace (HF) Transformers - -Please install required packages first: - -```bash -pip install -r requirements.txt -``` - -```bash -# 1. Weights & config -git clone https://huggingface.co/EleutherAI/gpt-j-6b gptj_model -pushd gptj_model && \ - rm -f pytorch_model.bin && \ - wget https://huggingface.co/EleutherAI/gpt-j-6b/resolve/main/pytorch_model.bin && \ -popd - -# 2. Vocab and merge table -wget https://huggingface.co/EleutherAI/gpt-j-6b/resolve/main/vocab.json -wget https://huggingface.co/EleutherAI/gpt-j-6b/resolve/main/merges.txt -``` - -### 2. Build TensorRT engine(s) - -TensorRT LLM builds TensorRT engine(s) using a HF checkpoint. If no checkpoint directory is specified, TensorRT LLM will build engine(s) using -dummy weights. - -Examples of build invocations: - -```bash -# Build a float16 engine using HF weights. -python convert_checkpoint.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --output_dir ./trt_ckpt/gptj_fp16_tp1/ -``` - -***For all kinds of checkpoints, they share the same trtllm-build command like:*** - -```bash -# Enable several TensorRT LLM plugins to increase runtime performance. It also helps with build time. -trtllm-build --checkpoint_dir ./trt_ckpt/gptj_fp16_tp1/ \ - --output_dir ./trt_engines/gptj_fp16_tp1/ \ - --gemm_plugin float16 \ - --max_batch_size=32 \ - --max_input_len=1919 \ - --max_seq_len=2047 -``` - -INT8 weight-only - -```bash -# Build an int8 weight-only engine using HF weights with TP=2 -python convert_checkpoint.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int8 \ - --output_dir ./trt_ckpt/gptj_int8_tp2/ \ - --tp_size 2 -``` -Building command is identical to the common one above. - -INT4 weight-only - -```bash -# Build an int4 weight only quantization engine using int4 weight only quantized weights. -python convert_checkpoint.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int4 \ - --output_dir ./trt_ckpt/gptj_int4wo_tp1/ -``` -Building command is identical to the common one above. - -#### FP8 Post-Training Quantization - -The examples below uses the NVIDIA Modelopt (AlgorithMic Model Optimization) toolkit for the model quantization process. - -First make sure Modelopt toolkit is installed (see [examples/quantization/README.md](/examples/quantization/README.md#preparation)) - -One can quantize HF GPT-J weights in FP8 as follows. - -```bash -# Quantize HF GPT-J 6B checkpoint into FP8 format -python ../../../quantization/quantize.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir ./trt_ckpt/gptj_fp8_tp1 \ - --calib_size 512 -``` -Building command is identical to the common one above. -Note that you can enable fp8 context fmha to get further acceleration by setting `--use_fp8_context_fmha enable` when building the engines. - -#### AWQ INT4 weight only quantization - -One can enable AWQ INT4 weight only quantization like the following command. - -```bash -# Enable AWQ int4 group-wise weight-only quantization. -python ../../../quantization/quantize.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --qformat int4_awq \ - --output_dir ./trt_ckpt/gptj_int4_awq_tp1 \ - --calib_size 512 -``` - -```bash -# Enable AWQ int4 group-wise weight-only quantization with tp2. -python ../../../quantization/quantize.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --qformat int4_awq \ - --tp_size 2 \ - --output_dir ./trt_ckpt/gptj_int4_awq_tp2 \ - --calib_size 512 -``` - -And build command is identical to the common one above. - -#### SmoothQuant (W8A8) quantization - -One can enable smoothquant W8A8 (weight per-channel, activation per-tensor) quantization like the following command. - -```bash -# Enable smoothquant (W8A8 kernel). -python ../../../quantization/quantize.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --qformat int8_sq \ - --output_dir ./trt_ckpt/gptj_sq_tp1 \ - --calib_size 512 -``` -Building command is identical to the common one above. - -#### Fused MultiHead Attention (FMHA) - -You can enable the FMHA kernels for GPT by adding `--context_fmha` to the invocation of `trtllm-build`. Note that it is enabled by default. - -If you find that the default fp16 accumulation (`--context_fmha`) cannot meet the requirement, you can try to enable fp32 accumulation by adding `--enable_context_fmha_fp32_acc` to the inference command (`run.py` or `summarize.py`). However, it is expected to see performance drop. - -Note `--context_fmha` has to be used together with `--gpt_attention_plugin float16`. - -#### INT8 KV cache -INT8 KV cache could be enabled to reduce memory footprint. It will bring more performance gains when batch size gets larger. - -You can get the INT8 scale of KV cache through Modelopt: - -```bash -# INT8 calibration -python ../../../quantization/quantize.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --kv_cache_dtype int8 \ - --output_dir ./trt_ckpt/gptj_fp16_int8kv_tp1 \ - --calib_size 512 -``` - -And build command is identical to the common one above. - -**INT8 KV cache + per-channel weight-only quantization** - -For example, you can enable INT8 KV cache together with per-channel INT8/INT4 weight-only quantization like the following command. - -```bash -# Enable INT8 KV cache together with per-channel INT8 weight-only quantization -python ../../../quantization/quantize.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --qformat int4_wo \ - --kv_cache_dtype int8 \ - --output_dir ./trt_ckpt/gptj_int4wo_int8kv_tp1 \ - --calib_size 512 -``` - -Building command is identical to the common one above. - -**INT8 KV cache + AWQ** - -In addition, you can enable INT8 KV cache together with AWQ (per-group INT4 weight-only quantization)like the following command. - -```bash -# Enable INT8 KV cache together with group-wise 4bit AWQ quantization -python ../../../quantization/quantize.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --qformat int4_awq \ - --kv_cache_dtype int8 \ - --output_dir ./trt_ckpt/gptj_int4awq_int8kv_tp1 \ - --calib_size 512 -``` - -Building command is identical to the common one above. - - -### 3. Run - - -To run a TensorRT LLM GPT-J model: - -```bash -python3 ../../../run.py --max_output_len=50 --engine_dir=gptj_engine --tokenizer_dir=gptj_model -``` - -## Summarization using the GPT-J model - -The following section describes how to run a TensorRT LLM GPT-J model to summarize the articles from the -[cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. For each summary, the script can compute the -[ROUGE](https://en.wikipedia.org/wiki/ROUGE_(metric)) scores and use the `ROUGE-1` score to validate the implementation. -The script can also perform the same summarization using the HF GPT-J model. - -As previously explained, the first step is to build the TensorRT engine as described above using HF weights. You also have to install the requirements: - -```bash -pip install -r requirements.txt -``` - -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_engines/gptj_fp16_tp1 \ - --hf_model_dir ./gpt-j-6b \ - --test_hf \ - --batch_size 1 \ - --test_trt_llm \ - --tensorrt_llm_rouge1_threshold 14 \ - --data_type fp16 \ - --check_accuracy -``` diff --git a/examples/models/contrib/gptj/convert_checkpoint.py b/examples/models/contrib/gptj/convert_checkpoint.py deleted file mode 100644 index 885b74c6a926..000000000000 --- a/examples/models/contrib/gptj/convert_checkpoint.py +++ /dev/null @@ -1,166 +0,0 @@ -import argparse -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -from transformers import AutoModelForCausalLM - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm.llmapi import QuantConfig -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import GPTJConfig, GPTJForCausalLM -from tensorrt_llm.models.convert_utils import infer_dtype -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - 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( - '--dtype', - type=str, - default='auto', - choices=['auto', 'float16', 'bfloat16', 'float32'], - help= - "The data type for the model weights and activations if not quantized. " - "If 'auto', the data type is automatically inferred from the source model; " - "however, if the source dtype is float32, it is converted to float16.") - parser.add_argument('--vocab_size', type=int, default=50400) - parser.add_argument('--n_positions', type=int, default=2048) - parser.add_argument('--n_layer', type=int, default=28) - parser.add_argument('--n_head', type=int, default=16) - parser.add_argument('--n_embd', type=int, default=4096) - parser.add_argument('--norm_eps', type=float, default=1e-05) - parser.add_argument('--rotary_dim', type=int, default=64) - 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('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - args = parser.parse_args() - - return args - - -def args_to_quant_config(args): - quant_algo = None - if args.use_weight_only and args.weight_only_precision == 'int8': - quant_algo = QuantAlgo.W8A16 - elif args.use_weight_only and args.weight_only_precision == 'int4': - quant_algo = QuantAlgo.W4A16 - return QuantConfig(quant_algo=quant_algo) - - -def convert_and_save_hf(args): - model_dir = args.model_dir - world_size = args.tp_size * args.pp_size - quant_config = args_to_quant_config(args) - - hf_model = AutoModelForCausalLM.from_pretrained(model_dir, - dtype='auto', - trust_remote_code=True) - - def convert_and_save_rank(args, rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - model = GPTJForCausalLM.from_hugging_face(hf_model, - args.dtype, - mapping=mapping, - quant_config=quant_config) - model.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del model - - if args.workers == 1: - for rank in range(world_size): - convert_and_save_rank(args, rank) - else: - with ThreadPoolExecutor(max_workers=args.workers) as p: - futures = [ - p.submit(convert_and_save_rank, args, rank) - for rank in range(world_size) - ] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - del hf_model - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - if args.model_dir is None: - config = GPTJConfig(architecture='GPTJForCausalLM', - dtype=infer_dtype(args.dtype), - num_hidden_layers=args.n_layer, - num_attention_heads=args.n_head, - hidden_size=args.n_embd, - norm_epsilon=args.norm_eps, - vocab_size=args.vocab_size, - position_embedding_type='rope_gptj', - max_position_embeddings=args.n_positions, - hidden_act='gelu', - rotary_dim=args.rotary_dim, - mapping=Mapping(world_size=args.tp_size * - args.pp_size, - tp_size=args.tp_size, - pp_size=args.pp_size), - quantization=args_to_quant_config(args)) - config.to_json_file(os.path.join(args.output_dir, 'config.json')) - else: - convert_and_save_hf(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/gptj/requirements.txt b/examples/models/contrib/gptj/requirements.txt deleted file mode 100644 index 2f8713d8658a..000000000000 --- a/examples/models/contrib/gptj/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score diff --git a/examples/models/contrib/gptneox/.gitignore b/examples/models/contrib/gptneox/.gitignore deleted file mode 100644 index d51d01117f3a..000000000000 --- a/examples/models/contrib/gptneox/.gitignore +++ /dev/null @@ -1,3 +0,0 @@ -__pycache__/ -gptneox_model/ -*.log diff --git a/examples/models/contrib/gptneox/README.md b/examples/models/contrib/gptneox/README.md deleted file mode 100644 index d1928bfdd8d9..000000000000 --- a/examples/models/contrib/gptneox/README.md +++ /dev/null @@ -1,235 +0,0 @@ -# GPT-NeoX - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document explains how to build the [GPT-NeoX](https://huggingface.co/EleutherAI/gpt-neox-20b) model using TensorRT LLM and run on a single GPU and a single node with -multiple GPUs. - -- [GPT-NeoX](#gpt-neox) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [1. Download weights from HuggingFace (HF) Transformers](#1-download-weights-from-huggingface-hf-transformers) - - [2. Convert weights from HF Transformers to TensorRT LLM format](#2-convert-weights-from-hf-transformers-to-tensorrt-llm-format) - - [3. Build TensorRT engine(s)](#3-build-tensorrt-engines) - - [4. Summarization using the GPT-NeoX model](#4-summarization-using-the-gpt-neox-model) - - [Apply groupwise quantization GPTQ](#apply-groupwise-quantization-gptq) - - [1. Download weights from HuggingFace (HF)](#1-download-weights-from-huggingface-hf) - - [2. Generating quantized weights](#2-generating-quantized-weights) - - [3. Convert weights from HF Transformers to TensorRT LLM format](#3-convert-weights-from-hf-transformers-to-tensorrt-llm-format) - - [4. Build TensorRT engine(s)](#4-build-tensorrt-engines) - - [5. Summarization using the GPT-NeoX model](#5-summarization-using-the-gpt-neox-model) - -## 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/models/contrib/gptneox`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 - * INT8 Weight-Only - * INT4 GPTQ - * Tensor Parallel - -## Usage - -The TensorRT LLM GPT-NeoX example code locates at [examples/models/contrib/gptneox](./). 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. - -### 1. Download weights from HuggingFace (HF) Transformers - -Please install required packages first: - -```bash -pip install -r requirements.txt -``` - -```bash -# Weights & config -git clone https://huggingface.co/EleutherAI/gpt-neox-20b gptneox_model -``` - -### 2. Convert weights from HF Transformers to TensorRT LLM format - -If you want to use Int8 weight only quantization, just need to add `--use_weight_only` flag. - -```bash -# Single GPU -python3 convert_checkpoint.py --model_dir ./gptneox_model \ - --dtype float16 \ - --output_dir ./gptneox/20B/trt_ckpt/fp16/1-gpu/ -# With 2-way Tensor Parallel -python3 convert_checkpoint.py --model_dir ./gptneox_model \ - --dtype float16 \ - --tp_size 2 \ - --workers 2 \ - --output_dir ./gptneox/20B/trt_ckpt/fp16/2-gpu/ -# Single GPU with int8 weight only -python3 convert_checkpoint.py --model_dir ./gptneox_model \ - --dtype float16 \ - --use_weight_only \ - --output_dir ./gptneox/20B/trt_ckpt/int8_wo/1-gpu/ -# With 2-way Tensor Parallel with int8 weight only -python3 convert_checkpoint.py --model_dir ./gptneox_model \ - --dtype float16 \ - --use_weight_only \ - --tp_size 2 \ - --workers 2 \ - --output_dir ./gptneox/20B/trt_ckpt/int8_wo/2-gpu/ -``` - -### 3. Build TensorRT engine(s) -```bash -# Single GPU -trtllm-build --checkpoint_dir ./gptneox/20B/trt_ckpt/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./gptneox/20B/trt_engines/fp16/1-gpu/ -# With 2-way Tensor Parallel -trtllm-build --checkpoint_dir ./gptneox/20B/trt_ckpt/fp16/2-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --workers 2 \ - --output_dir ./gptneox/20B/trt_engines/fp16/2-gpu/ -# Single GPU with int8 weight only -trtllm-build --checkpoint_dir ./gptneox/20B/trt_ckpt/int8_wo/1-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./gptneox/20B/trt_engines/int8_wo/1-gpu/ -# With 2-way Tensor Parallel with int8 weight only -trtllm-build --checkpoint_dir ./gptneox/20B/trt_ckpt/int8_wo/2-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --workers 2 \ - --output_dir ./gptneox/20B/trt_engines/int8_wo/2-gpu/ -``` - -### 4. Summarization using the GPT-NeoX model - -The following section describes how to run a TensorRT LLM GPT-NeoX model to summarize the articles from the -[cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. For each summary, the script can compute the -[ROUGE](https://en.wikipedia.org/wiki/ROUGE_(metric)) scores and use the `ROUGE-1` score to validate the implementation. -The script can also perform the same summarization using the HF GPT-NeoX model. - -```bash -# Single GPU -python3 ../../../summarize.py --engine_dir ./gptneox/20B/trt_engines/fp16/1-gpu/ \ - --test_trt_llm \ - --hf_model_dir gptneox_model \ - --data_type fp16 -# With 2-way Tensor Parallel -mpirun -np 2 --oversubscribe --allow-run-as-root \ - python3 ../../../summarize.py --engine_dir ./gptneox/20B/trt_engines/fp16/2-gpu/ \ - --test_trt_llm \ - --hf_model_dir gptneox_model \ - --data_type fp16 -# Single GPU with int8 weight only -python3 ../../../summarize.py --engine_dir ./gptneox/20B/trt_engines/int8_wo/1-gpu/ \ - --test_trt_llm \ - --hf_model_dir gptneox_model \ - --data_type fp16 -# With 2-way Tensor Parallel with int8 weight only -mpirun -np 2 --oversubscribe --allow-run-as-root \ - python3 ../../../summarize.py --engine_dir ./gptneox/20B/trt_engines/int8_wo/2-gpu/ \ - --test_trt_llm \ - --hf_model_dir gptneox_model \ - --data_type fp16 -``` - -## Apply groupwise quantization GPTQ - -### 1. Download weights from HuggingFace (HF) - -```bash -# Weights & config -sh get_weights.sh -``` - -### 2. Generating quantized weights - -In this example, the weights are quantized using [GPTQ-for-LLaMa](https://github.com/qwopqwop200/GPTQ-for-LLaMa). Note that the parameter `--act-order` referring to whether to apply the activation order GPTQ heuristic is **not supported** by TRT-LLM. - -```bash -sh gptq_convert.sh -``` - -### 3. Convert weights from HF Transformers to TensorRT LLM format - -To apply groupwise quantization GPTQ, addition command-line flags need to be passed to `convert_checkpoint.py`: -Here `--quant_ckpt_path` flag specifies the output safetensors of `gptq_convert.sh` script. - -```bash -# Single GPU -python3 convert_checkpoint.py --model_dir ./gptneox_model \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int4_gptq \ - --quant_ckpt_path ./gptneox_model/gptneox-20b-4bit-gs128.safetensors \ - --output_dir ./gptneox/20B/trt_ckpt/int4_gptq/1-gpu/ -# With 2-way Tensor Parallel -python3 convert_checkpoint.py --model_dir ./gptneox_model \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int4_gptq \ - --tp_size 2 \ - --workers 2 \ - --quant_ckpt_path ./gptneox_model/gptneox-20b-4bit-gs128.safetensors \ - --output_dir ./gptneox/20B/trt_ckpt/int4_gptq/2-gpu/ -``` - -### 4. Build TensorRT engine(s) - -The command to build TensorRT engines to apply GPTQ does not change: - -```bash -# Single GPU -trtllm-build --checkpoint_dir ./gptneox/20B/trt_ckpt/int4_gptq/1-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./gptneox/20B/trt_engines/int4_gptq/1-gpu/ -# With 2-way Tensor Parallel -trtllm-build --checkpoint_dir ./gptneox/20B/trt_ckpt/int4_gptq/2-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --workers 2 \ - --output_dir ./gptneox/20B/trt_engines/int4_gptq/2-gpu/ -``` - -### 5. Summarization using the GPT-NeoX model - -The command to run summarization with GPTQ quantized model also does not change: - -```bash -# Single GPU -python3 ../../../summarize.py --engine_dir ./gptneox/20B/trt_engines/int4_gptq/1-gpu/ \ - --test_trt_llm \ - --hf_model_dir gptneox_model \ - --data_type fp16 -# With 2-way Tensor Parallel -mpirun -np 2 --oversubscribe --allow-run-as-root \ - python3 ../../../summarize.py --engine_dir ./gptneox/20B/trt_engines/int4_gptq/2-gpu/ \ - --test_trt_llm \ - --hf_model_dir gptneox_model \ - --data_type fp16 -``` diff --git a/examples/models/contrib/gptneox/convert_checkpoint.py b/examples/models/contrib/gptneox/convert_checkpoint.py deleted file mode 100644 index 7c4c76413c42..000000000000 --- a/examples/models/contrib/gptneox/convert_checkpoint.py +++ /dev/null @@ -1,732 +0,0 @@ -import argparse -import json -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import List, Optional - -import safetensors -import safetensors.torch -import torch -from safetensors import safe_open -from transformers import AutoConfig, AutoModelForCausalLM - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import str_dtype_to_torch -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.convert_utils import (get_weight, get_weight_and_bias, - split_matrix_tp) -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - 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('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - 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', 'int4_gptq'], - 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('--quant_ckpt_path', - type=str, - default=None, - help='Path of a quantized model checkpoint') - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - args = parser.parse_args() - - return args - - -# TODO: Seems all convert checkpoints may use following utility functions. -# Maybe in one common version. -def reorder_qkv_weight_or_bias(weight: torch.Tensor, - head_dim: int, - num_heads: int, - num_kv_heads: Optional[int] = None, - tp_size: int = 1, - is_bias: bool = False) -> torch.Tensor: - """ Reorder the qkv weight for TRT-LLM use. - - The shape of the fused QKV weights in HF is different from the shape that - TRT-LLM requires. In particular, the weight of HF consists of interleaved - q, k, v head weights, while that of TRT-LLM is contiguous. - HF : [q1, k1, v1, ..., qh, kh, vh] - TRT-LLM: [q1, ..., qh, k1, ..., kh, v1, vh] - where qi, vi, ki are weight vectors corresponding to attention head i. - It's similar to multi/grouped query attention cases. - - We reorder and split the weight of an attention layer to fit into TRT-LLM. - The reordered weight and bias will be - weight: (T, Qh * D + 2 * KVh * D, H) - bias : (T, Qh * D + 2 * KVh * D) - where T=tp_size, Qh=local_num_q_heads, KVh=local_num_kv_heads, D=head_dim, - H=hidden_dim. In the multi/grouped query attention, the number of K/V - attention heads are less than that of Q attention, so that K/V attention - heads may be shared across different ranks if necessary. - - For tensor parallelism, we use the first dimension to select the - corresponding weights. - """ - - # Query types and expected kv heads. - # - Conventional MHA: num_heads = num_kv_heads - # - Multi-Query Attention: num_kv_heads = 1 - # - Grouped-Query Attention: num_heads % num_kv_heads = 0 - num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads - assert num_heads % num_kv_heads == 0, \ - f'num_heads({num_heads}) must be divisible by ' \ - f'num_kv_heads({num_kv_heads})).' - - # The number of attention heads per group: N q head + 1 k head + 1 v head. - num_group_heads = num_heads // num_kv_heads + 2 - assert weight.shape[0] == num_kv_heads * num_group_heads * head_dim, \ - f'{weight.shape[0]} != {num_kv_heads} * {num_group_heads} * {head_dim}' - - qkv_in = num_heads * head_dim if not is_bias else 1 - - # Split Q/K/V weights - weight = weight.reshape(num_kv_heads, num_heads // num_kv_heads + 2, - head_dim, qkv_in) - q_w = weight[:, :-2, ...] # (nKV, num_heads // nKV, head_dim, qkv_in) - k_w = weight[:, -2:-1, ...] # (nKV, 1, head_dim, qkv_in) - v_w = weight[:, -1:, ...] # (nKV, 1, head_dim, qkv_in) - - if num_kv_heads < num_heads and num_kv_heads < tp_size: - # Duplicate K/V heads to make sure that each rank has at least one - # K/V heads. For instance, num_heads=8, num_kv_heads=2, tp_size=4, - # we will make the qkv weight as below. - # Orig: [q0 q1 q2 q3 k0 v0 q4 q5 q6 q7 k1 v0 v1] - # >>>> [[q0 q1 k0 v0], [q2 q3 k0 v0], [q4 q5 k1 v1], [q6 q7 k1 v1]] - assert tp_size % num_kv_heads == 0 - num_dups = tp_size // num_kv_heads - - # k_w and v_w have the same shape. - new_shape = (num_kv_heads, num_dups) + k_w.shape[2:] - k_w = torch.broadcast_to(k_w, size=new_shape) - v_w = torch.broadcast_to(v_w, size=new_shape) - - # Update the number of kv heads. - num_kv_heads = tp_size - - reordered = torch.concat( - [ - q_w.reshape(tp_size, num_heads // tp_size, head_dim, qkv_in), - k_w.reshape(tp_size, num_kv_heads // tp_size, head_dim, qkv_in), - v_w.reshape(tp_size, num_kv_heads // tp_size, head_dim, qkv_in), - ], - dim=1, - ) - - qkv_out = (num_heads + 2 * num_kv_heads) // tp_size * head_dim - return reordered.reshape((tp_size, qkv_out, -1)) - - -def get_gptq_gptneox_group_size(quant_ckpt_path, hf_config): - gptq_model = safe_open(quant_ckpt_path, framework="pt", device=0) - gptq_prefix = "gpt_neox." - split_sym = "." - - def load(key, no_prefix=0): - if no_prefix: - return gptq_model.get_tensor(key).cpu() - else: - return gptq_model.get_tensor(gptq_prefix + key).cpu() - - hidden_size = hf_config.hidden_size - prefix = "layers" + split_sym + "0" + split_sym - - scales_fp16 = load(prefix + 'attention.query_key_value.scales') - return hidden_size // scales_fp16.shape[0] - - -def load_from_gptq_gptneox(quant_ckpt_path, - hf_config=None, - use_parallel_embedding=False, - sharding_dim=0, - mapping=Mapping(), - dtype='float16'): - tensorrt_llm.logger.info( - 'Loading weights from groupwise GPTQ LLaMA safetensors...') - weights = {} - tik = time.time() - - gptq_model = safe_open(quant_ckpt_path, framework="pt", device=0) - gptq_prefix = "gpt_neox." - gptq_suffix_list = [".qweight", ".qzeros", ".scales"] - split_sym = "." - - packer = torch.ops.trtllm.pack_int8_tensor_to_packed_int4 - preprocessor = torch.ops.trtllm.preprocess_weights_for_mixed_gemm - torch_dtype = str_dtype_to_torch(dtype) - - def load(key, no_prefix=0): - if no_prefix: - return gptq_model.get_tensor(key).cpu() - else: - return gptq_model.get_tensor(gptq_prefix + key).cpu() - - 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)) - assert False, "Invalid TP size" - return v.split(v.shape[dim] // mapping.tp_size, - dim=dim)[mapping.tp_rank].contiguous() - - def unpack_int32_into_int8(w_packed): - # Unpack inputs packed in int32/float32 into uint4 and store them in int8 format - w_packed_int4x2 = w_packed.contiguous().view(torch.uint8) - w_unpacked = torch.zeros(w_packed_int4x2.shape[0], - w_packed_int4x2.shape[1] * 2, - dtype=torch.int8, - device=w_packed.device) - w_unpacked[:, ::2] = w_packed_int4x2 % 16 - w_unpacked[:, 1::2] = w_packed_int4x2 // 16 - return w_unpacked.contiguous() - - def process_and_assign_weight(v: List[torch.Tensor], - tllm_prex: str, - tp_dim: int = -1): - if tp_dim == -1: - qweight_int32, qzeros_int32, scales_fp16 = [ - item.cpu() for item in v - ] - else: - qweight_int32, qzeros_int32, scales_fp16 = [ - torch_split(item, tp_dim).cpu() for item in v - ] - - USE_UINT4_INPUT = 1 # Set to true if checkpoint store UINT4 weights - USE_GPTQ_FOR_LLAMA = 1 # GPTQ-for-LLaMA added 1 to zeros - - qweight_unpacked_int8 = unpack_int32_into_int8( - qweight_int32.T).T.contiguous() - 8 - qweight_interleaved = preprocessor(packer(qweight_unpacked_int8), - torch.quint4x2, - torch.float16).view(torch.float16) - # zeros = zeros * scales - qzeros_unpacked_int32 = unpack_int32_into_int8(qzeros_int32) - if not USE_UINT4_INPUT: - # Correcting UINT4 values back to INT4 order - mask_negative = qzeros_unpacked_int32[qzeros_unpacked_int32 < 0] - mask_positive = qzeros_unpacked_int32[qzeros_unpacked_int32 >= 0] - qzeros_unpacked_int32 = qzeros_unpacked_int32 + 16 * mask_negative - 16 * mask_positive - zeros_x_scales_fp16 = (-qzeros_unpacked_int32 + 8 * USE_UINT4_INPUT - - USE_GPTQ_FOR_LLAMA) * scales_fp16 - zeros_x_scales_fp16 = zeros_x_scales_fp16.half() - - results = { - f'{tllm_prex}.weight': qweight_interleaved, - f'{tllm_prex}.weights_scaling_factor': scales_fp16, - f'{tllm_prex}.zero': zeros_x_scales_fp16, - } - return results - - def preprocess_groupwise_weight_params(qweight_unpacked_int8, scales_fp16, - qzeros_unpacked_int8): - UINT4_TO_INT4_FLAG = 1 - GPTQ_FLAG = 1 - - qweight_interleaved = preprocessor(packer(qweight_unpacked_int8), - torch.quint4x2, - torch.float16).view(torch.float16) - - # zeros = zeros * scales - zeros_x_scales_fp16 = (-qzeros_unpacked_int8 + 8 * UINT4_TO_INT4_FLAG - - GPTQ_FLAG) * scales_fp16 - zeros_x_scales_fp16 = zeros_x_scales_fp16.half() - - # return processed interleaved weight, original scales and zeros * scales - return qweight_interleaved.contiguous(), scales_fp16.contiguous( - ), zeros_x_scales_fp16.contiguous() - - # Load weights from GPTQ checkpoint into TRT-LLM module - # 1. vocab_embedding - v = load('embed_in.weight') - if mapping.is_first_pp_rank(): - if not use_parallel_embedding: - weights['transformer.vocab_embedding.weight'] = v.to(torch_dtype) - else: - assert hf_config.vocab_size % mapping.tp_size == 0 - weights['transformer.vocab_embedding.weight'] = torch_split( - v, sharding_dim).to(torch_dtype) - # 2. lm_head - if mapping.is_last_pp_rank(): - v = load('embed_out.weight', no_prefix=1) - weights['lm_head.weight'] = torch_split(v, 0).to(torch_dtype) - - # 3. ln_f - v = load('final_layer_norm.weight') - b = load('final_layer_norm.bias') - if mapping.is_last_pp_rank(): - weights['transformer.ln_f.weight'] = v.to(torch_dtype) - weights['transformer.ln_f.bias'] = b.to(torch_dtype) - # 4. Weights inside each layer - num_hidden_layers = hf_config.num_hidden_layers - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - layer_idx = l - layers_range[0] - prefix = "layers" + split_sym + str(l) + split_sym - tensorrt_llm.logger.info(f'Process weights in layer: {layer_idx}') - # layer = tensorrt_llm_llama.layers[layer_idx] - tllm_prex = f'transformer.layers.{l - layers_range[0]}' - # 4.1 attention.qkv - num_heads = hf_config.num_attention_heads - hidden_size = hf_config.hidden_size - head_size = hidden_size // num_heads - - qweight_int32 = load(prefix + 'attention.query_key_value.qweight') - scales_fp16 = load(prefix + 'attention.query_key_value.scales') - qzeros_int32 = load(prefix + 'attention.query_key_value.qzeros') - biases_fp16 = load(prefix + 'attention.query_key_value.bias') - GROUP_SIZE = hidden_size // scales_fp16.shape[0] - - # [hidden_size // 8, hidden_size * 3] -> [hidden_size * 3, hidden_size] - qweight_unpacked_int8 = unpack_int32_into_int8( - qweight_int32.T).contiguous() - 8 - # [hidden_size // GROUP_SIZE, hidden_size * 3 // 8] -> - # [hidden_size // GROUP_SIZE, hidden_size * 3] - qzeros_unpacked_int8 = unpack_int32_into_int8(qzeros_int32) - - # qkv_weights [num_heads x (q|k|v), hidden_size] -> - # [(num_heads x q)|(num_heads x k)|(num_heads x v), hidden_size] - new_qkv_weight_shape = torch.Size( - [num_heads, 3, head_size * qweight_unpacked_int8.size()[-1]]) - # [hidden_size * 3, hidden_size] - qweight_unpacked_int8 = qweight_unpacked_int8.view( - new_qkv_weight_shape).permute(1, 0, 2).reshape( - [hidden_size * 3, hidden_size]).contiguous() - - new_qkv_scale_shape = torch.Size( - [num_heads, 3, head_size * (hidden_size // GROUP_SIZE)]) - # [hidden_size * 3, hidden_size // GROUP_SIZE] - scales_fp16 = scales_fp16.T.contiguous().view( - new_qkv_scale_shape).permute(1, 0, 2).reshape( - [hidden_size * 3, hidden_size // GROUP_SIZE]).contiguous() - - new_qkv_zero_shape = torch.Size( - [num_heads, 3, head_size * (hidden_size // GROUP_SIZE)]) - # [hidden_size * 3, hidden_size // GROUP_SIZE] - qzeros_unpacked_int8 = qzeros_unpacked_int8.T.contiguous().view( - new_qkv_zero_shape).permute(1, 0, 2).reshape( - [hidden_size * 3, hidden_size // GROUP_SIZE]).contiguous() - - new_qkv_bias_shape = torch.Size([num_heads, 3, head_size]) - biases_fp16 = biases_fp16.view(new_qkv_bias_shape).permute( - 1, 0, 2).reshape([hidden_size * 3]) - - tp_size = mapping.tp_size - - if tp_size > 1: - qweight_unpacked_int8 = qweight_unpacked_int8.reshape( - [3, hidden_size, hidden_size]) - qweight_unpacked_int8 = torch_split(qweight_unpacked_int8, dim=1) - qweight_unpacked_int8 = qweight_unpacked_int8.reshape( - [3 * hidden_size // tp_size, hidden_size]) - - scales_fp16 = scales_fp16.reshape( - [3, hidden_size, hidden_size // GROUP_SIZE]) - scales_fp16 = torch_split(scales_fp16, dim=1) - scales_fp16 = scales_fp16.reshape( - [3 * hidden_size // tp_size, hidden_size // GROUP_SIZE]) - - qzeros_unpacked_int8 = qzeros_unpacked_int8.reshape( - [3, hidden_size, hidden_size // GROUP_SIZE]) - qzeros_unpacked_int8 = torch_split(qzeros_unpacked_int8, dim=1) - qzeros_unpacked_int8 = qzeros_unpacked_int8.reshape( - [3 * hidden_size // tp_size, hidden_size // GROUP_SIZE]) - - biases_fp16 = biases_fp16.reshape([3, hidden_size]) - biases_fp16 = torch_split(biases_fp16, dim=1) - biases_fp16 = biases_fp16.reshape([3 * hidden_size // tp_size]) - - qweight_fp32, scales_fp16, zeros_fp16 = preprocess_groupwise_weight_params( - qweight_unpacked_int8.T.contiguous(), scales_fp16.T.contiguous(), - qzeros_unpacked_int8.T.contiguous()) - weights.update({ - f'{tllm_prex}.attention.qkv.weight': qweight_fp32, - f'{tllm_prex}.attention.qkv.weights_scaling_factor': scales_fp16, - f'{tllm_prex}.attention.qkv.zero': zeros_fp16, - f'{tllm_prex}.attention.qkv.bias': biases_fp16, - }) - - # 4.2 attention.dense - v = [load(prefix + 'attention.dense' + suf) for suf in gptq_suffix_list] - # pre scaling down for duplicated bias add between different tp ranks - b = load(prefix + 'attention.dense.bias') / mapping.tp_size - - weights.update( - process_and_assign_weight(v, - f'{tllm_prex}.attention.dense', - tp_dim=0)) - weights.update({f'{tllm_prex}.attention.dense.bias': b.to(torch_dtype)}) - # 4.3 mlp.fc - v = [ - load(prefix + 'mlp.dense_h_to_4h' + suf) for suf in gptq_suffix_list - ] - b = load(prefix + 'mlp.dense_h_to_4h.bias') - weights.update( - process_and_assign_weight(v, f'{tllm_prex}.mlp.fc', tp_dim=1)) - weights.update( - {f'{tllm_prex}.mlp.fc.bias': torch_split(b, dim=0).to(torch_dtype)}) - # 4.4 mlp.proj - v = [ - load(prefix + 'mlp.dense_4h_to_h' + suf) for suf in gptq_suffix_list - ] - # pre scaling down for duplicated bias add between different tp ranks - b = load(prefix + 'mlp.dense_4h_to_h.bias') / mapping.tp_size - - weights.update( - process_and_assign_weight(v, f'{tllm_prex}.mlp.proj', tp_dim=0)) - weights.update({f'{tllm_prex}.mlp.proj.bias': b.to(torch_dtype)}) - # 4.5 input_layernorm - v = load(prefix + 'input_layernorm.weight') - b = load(prefix + 'input_layernorm.bias') - weights[f'{tllm_prex}.input_layernorm.weight'] = v.to(torch_dtype) - weights[f'{tllm_prex}.input_layernorm.bias'] = b.to(torch_dtype) - - # 4.6 post_layernorm - v = load(prefix + 'post_attention_layernorm.weight') - b = load(prefix + 'post_attention_layernorm.bias') - weights[f'{tllm_prex}.post_attention_layernorm.weight'] = v.to( - torch_dtype) - weights[f'{tllm_prex}.post_attention_layernorm.bias'] = b.to( - torch_dtype) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - tensorrt_llm.logger.info(f'Weights loaded. Total time: {t}') - - return weights - - -def split_qkv_weight(weight: torch.Tensor, - hidden_size: int, - num_heads: int, - tp_size: int, - rank: int, - is_bias: bool, - num_kv_heads: Optional[int] = None) -> torch.Tensor: - """ Splits the QKV matrix according to tensor parallelism """ - head_dim = hidden_size // num_heads - weight = reorder_qkv_weight_or_bias(weight, - head_dim=head_dim, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - tp_size=tp_size, - is_bias=is_bias) - - # Copy a sliced tensor to prevent memory leak. A sliced tensor shares the - # memory buffer of the original tensor. So, returning without copying makes - # the buffer of a loaded "qkv" be referenced, resulting GC can't release - # those weights until the whole process ends. - if not is_bias: - return weight[rank, ...].clone().contiguous() - else: - return weight[rank, ...].ravel().clone().contiguous() - - -def get_tllm_linear_weight(weight, - prefix, - bias=None, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8): - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[prefix + 'weight'] = processed_torch_weights - results[prefix + 'per_channel_scale'] = torch_weight_scales - else: - results[prefix + 'weight'] = weight.contiguous() - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def convert_hf_gptneox(hf_model, - mapping: Mapping, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8): - weights = {} - tik = time.time() - - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - num_attention_heads = hf_model.config.num_attention_heads - hidden_size = hf_model.config.hidden_size - tensor_parallel = mapping.tp_size - rank = mapping.rank - - for l in range(hf_model.config.num_hidden_layers): - prefix = f'gpt_neox.layers.{l}.' - tllm_prex = f'transformer.layers.{l}.' - - qkv_weight, qkv_bias = get_weight_and_bias( - model_params, prefix + 'attention.query_key_value', dtype) - qkv_w = split_qkv_weight(qkv_weight, - hidden_size, - num_attention_heads, - mapping.tp_size, - mapping.tp_rank, - is_bias=False, - num_kv_heads=num_attention_heads) - if qkv_bias is None: - qkv_b = None - else: - qkv_b = split_qkv_weight(qkv_bias, - hidden_size, - num_attention_heads, - mapping.tp_size, - mapping.tp_rank, - is_bias=True, - num_kv_heads=num_attention_heads) - weights.update( - get_tllm_linear_weight(qkv_w, tllm_prex + 'attention.qkv.', qkv_b, - use_weight_only, - plugin_weight_only_quant_type)) - - attn_dense_weight, attn_dense_bias = get_weight_and_bias( - model_params, prefix + 'attention.dense', dtype) - split_v = split_matrix_tp(attn_dense_weight, - tensor_parallel, - rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.dense.', - attn_dense_bias, use_weight_only, - plugin_weight_only_quant_type)) - - mlp_fc_weight, mlp_fc_bias = get_weight_and_bias( - model_params, prefix + 'mlp.dense_h_to_4h', dtype) - split_v = split_matrix_tp(mlp_fc_weight, tensor_parallel, rank, dim=0) - bias = split_matrix_tp(mlp_fc_bias, tensor_parallel, rank, dim=0) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.fc.', bias, - use_weight_only, - plugin_weight_only_quant_type)) - - mlp_proj_weight, mlp_proj_bias = get_weight_and_bias( - model_params, prefix + 'mlp.dense_4h_to_h', dtype) - split_v = split_matrix_tp(mlp_proj_weight, tensor_parallel, rank, dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.proj.', - mlp_proj_bias, use_weight_only, - plugin_weight_only_quant_type)) - - # Layer norms do not use tensor parallelism - input_ln_weight, input_ln_bias = get_weight_and_bias( - model_params, prefix + 'input_layernorm', dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - weights[tllm_prex + 'input_layernorm.bias'] = input_ln_bias - - post_ln_weight, post_ln_bias = get_weight_and_bias( - model_params, prefix + 'post_attention_layernorm', dtype) - weights[tllm_prex + 'post_attention_layernorm.weight'] = post_ln_weight - weights[tllm_prex + 'post_attention_layernorm.bias'] = post_ln_bias - - embed_w = get_weight(model_params, 'gpt_neox.embed_in', dtype) - lm_head_w = get_weight(model_params, 'embed_out', dtype) - - weights['lm_head.weight'] = split_matrix_tp(lm_head_w, - tensor_parallel, - rank, - dim=0) - - if not use_parallel_embedding: - weights['transformer.vocab_embedding.weight'] = embed_w - else: - assert hf_model.config.vocab_size % tensor_parallel == 0 - weights['transformer.vocab_embedding.weight'] = split_matrix_tp( - embed_w, tensor_parallel, rank, dim=sharding_dim) - - ln_f_w, ln_f_b = get_weight_and_bias(model_params, - 'gpt_neox.final_layer_norm', dtype) - weights['transformer.ln_f.weight'] = ln_f_w - weights['transformer.ln_f.bias'] = ln_f_b - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -if __name__ == '__main__': - emit_engine_arch_deprecation("convert_checkpoint.py") - # TODO(qijun): Currently, the convert script depends on a torch op: - # torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix, - # which is included in tensorrt_llm Python package. Otherwise, the convert - # script does not need to import tensorrt_llm. Will remove it after reimplementing - # the op with PyTorch. - print(tensorrt_llm.__version__) - args = parse_arguments() - world_size = args.tp_size * args.pp_size - assert args.pp_size == 1, "Pipeline parallelism is not supported." - - tensorrt_llm.logger.set_level('info') - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - quant_algo = None - plugin_weight_only_quant_type = None - if args.use_weight_only and args.weight_only_precision == 'int8': - plugin_weight_only_quant_type = torch.int8 - quant_algo = QuantAlgo.W8A16 - elif args.use_weight_only and args.weight_only_precision == 'int4': - plugin_weight_only_quant_type = torch.quint4x2 - quant_algo = QuantAlgo.W4A16 - elif args.use_weight_only and args.weight_only_precision == 'int4_gptq': - quant_algo = QuantAlgo.W4A16_GPTQ - - hf_config = AutoConfig.from_pretrained(args.model_dir) - hf_model = AutoModelForCausalLM.from_pretrained(args.model_dir, - dtype="auto") - - config = { - 'architecture': hf_config.architectures[0], - 'dtype': args.dtype, - 'num_hidden_layers': hf_config.num_hidden_layers, - 'num_attention_heads': hf_config.num_attention_heads, - 'hidden_size': hf_config.hidden_size, - 'vocab_size': hf_config.vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'max_position_embeddings': hf_config.max_position_embeddings, - 'rotary_emb_base': hf_config.rotary_emb_base, - 'rotary_pct': hf_config.rotary_pct, - 'hidden_act': hf_config.hidden_act, - 'quantization': { - 'quant_algo': quant_algo, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - } - if args.use_weight_only and args.weight_only_precision == 'int4_gptq': - config['quantization'].update({ - 'has_zero_point': - True, - 'group_size': - get_gptq_gptneox_group_size(args.quant_ckpt_path, hf_config) - }) - - with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: - json.dump(config, f, indent=4) - - def covert_and_save(rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - - if args.use_weight_only and args.weight_only_precision == 'int4_gptq': - weights = load_from_gptq_gptneox( - args.quant_ckpt_path, - hf_config, - use_parallel_embedding=args.use_parallel_embedding, - sharding_dim=args.embedding_sharding_dim, - mapping=mapping, - dtype=args.dtype) - else: - weights = convert_hf_gptneox( - hf_model, - mapping, - dtype=args.dtype, - use_weight_only=args.use_weight_only, - plugin_weight_only_quant_type=plugin_weight_only_quant_type, - use_parallel_embedding=args.use_parallel_embedding, - sharding_dim=args.embedding_sharding_dim) - safe_save_path = os.path.join(args.output_dir, - f'rank{rank}.safetensors') - tensorrt_llm.logger.info(f'Saving safetensors to: {safe_save_path}') - safetensors.torch.save_file(weights, safe_save_path) - tensorrt_llm.logger.info(f'Saved safetensors to: {safe_save_path}') - return True - - if args.workers == 1: - for rank in range(world_size): - passed = covert_and_save(rank) - assert passed, "Convert checkpoint failed" - else: - with ThreadPoolExecutor(max_workers=args.workers) as p: - futures = [ - p.submit(covert_and_save, rank) for rank in range(world_size) - ] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') diff --git a/examples/models/contrib/gptneox/gptq_convert.sh b/examples/models/contrib/gptneox/gptq_convert.sh deleted file mode 100644 index b0a89462cd1a..000000000000 --- a/examples/models/contrib/gptneox/gptq_convert.sh +++ /dev/null @@ -1,9 +0,0 @@ -git clone https://github.com/qwopqwop200/GPTQ-for-LLaMa.git GPTQ-for-LLaMa - -pip install -r ./GPTQ-for-LLaMa/requirements.txt - -CUDA_VISIBLE_DEVICES=0 python3 GPTQ-for-LLaMa/neox.py ./gptneox_model \ -wikitext2 \ ---wbits 4 \ ---groupsize 128 \ ---save_safetensors ./gptneox_model/gptneox-20b-4bit-gs128.safetensors diff --git a/examples/models/contrib/gptneox/requirements.txt b/examples/models/contrib/gptneox/requirements.txt deleted file mode 100644 index cc77e27f78ee..000000000000 --- a/examples/models/contrib/gptneox/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -rouge_score -evaluate diff --git a/examples/models/contrib/grok/README.md b/examples/models/contrib/grok/README.md deleted file mode 100644 index 41bcd8505b62..000000000000 --- a/examples/models/contrib/grok/README.md +++ /dev/null @@ -1,92 +0,0 @@ -# Grok-1 - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run grok-1 model in TensorRT LLM on both single GPU, single node multi-GPU and multi-node multi-GPU. - -- [Grok1](#Grok-1) - - [Prerequisite](#prerequisite) - - [Hardware](#hardware) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [Build TensorRT engine(s)](#build-tensorrt-engines) - -## Prerequisite -First of all, please clone the official grok-1 code repo with below commands and install the dependencies. -```bash -git clone https://github.com/xai-org/grok-1.git /path/to/folder -``` -And then downloading the weights per [instructions](https://github.com/xai-org/grok-1?tab=readme-ov-file#downloading-the-weights). - -## Hardware -The grok-1 model requires a node with 8x80GB GPU memory(at least). - -## Overview - -The TensorRT LLM Grok-1 implementation can be found in [tensorrt_llm/models/grok/model.py](../../../../tensorrt_llm/models/grok/model.py). The TensorRT LLM Grok-1 example code is located in [`examples/models/contrib/grok`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert the Grok-1 model into TensorRT LLM checkpoint format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * INT8 Weight-Only - * Tensor Parallel - * STRONGLY TYPED - -## Usage - -The TensorRT LLM Grok-1 example code locates at [examples/models/contrib/grok](./). It takes xai weights as input, and builds the corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Build TensorRT engine(s) - -Please install required packages first to make sure the example uses matched `tensorrt_llm` version: - -```bash -pip install -r requirements.txt -``` - -Need to prepare the Grok-1 checkpoint by following the guides here https://github.com/xai-org/grok-1. - -TensorRT LLM Grok-1 builds TensorRT engine(s) from Xai's checkpoints. - -Normally `trtllm-build` only requires single GPU, but if you've already got all the GPUs needed for inference, you could enable parallel building to make the engine building process faster by adding `--workers` argument. Please note that currently `workers` feature only supports single node. - - -Below is the step-by-step to run Grok-1 with TensorRT LLM. - -```bash -# Build the bfloat16 engine from xai official weights. -python convert_checkpoint.py --model_dir ./tmp/grok-1/ \ - --output_dir ./tllm_checkpoint_8gpus_bf16 \ - --dtype bfloat16 \ - --use_weight_only \ - --tp_size 8 \ - --workers 8 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_8gpus_bf16 \ - --output_dir ./tmp/grok-1/trt_engines/bf16/8-gpus \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 \ - --moe_plugin bfloat16 \ - --paged_kv_cache enable \ - --remove_input_padding enable \ - --workers 8 - - -# Run Grok-1 with 8 GPUs -mpirun -n 8 --allow-run-as-root \ - python ../../../run.py \ - --input_text "The answer to life the universe and everything is of course" \ - --engine_dir ./tmp/grok-1/trt_engines/bf16/8-gpus \ - --max_output_len 50 --top_p 1 --top_k 8 --temperature 0.3 \ - --vocab_file ./tmp/grok-1/tokenizer.model -``` diff --git a/examples/models/contrib/grok/convert_checkpoint.py b/examples/models/contrib/grok/convert_checkpoint.py deleted file mode 100644 index 18fb429eb04b..000000000000 --- a/examples/models/contrib/grok/convert_checkpoint.py +++ /dev/null @@ -1,366 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import argparse -import json -import os -import sys -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import numpy as np - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.layers import MoeConfig -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import GrokForCausalLM -from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument('--weights_dir', type=str, default=None) - - 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('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument('--vocab_size', type=int, default=32000) - parser.add_argument('--n_positions', type=int, default=2048) - parser.add_argument('--n_layer', type=int, default=32) - parser.add_argument('--n_head', type=int, default=32) - parser.add_argument('--n_kv_head', type=int, default=None) - parser.add_argument('--n_embd', type=int, default=4096) - parser.add_argument('--inter_size', type=int, default=11008) - parser.add_argument('--rms_norm_eps', type=float, default=1e-06) - - 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( - '--disable_weight_only_quant_plugin', - default=False, - action="store_true", - help= - 'By default, using plugin implementation for weight quantization. Enabling disable_weight_only_quant_plugin flag will use ootb implementation instead of plugin.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8'], - 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('--load_by_shard', - action='store_true', - help='Load a pretrained model shard-by-shard.') - parser.add_argument('--hidden_act', type=str, default='silu') - - parser.add_argument('--rotary_base', type=float, default=10000.0) - - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument( - '--moe_num_experts', - default=0, - type=int, - help='Specify the number of experts to use for MOE layers') - parser.add_argument( - '--moe_top_k', - default=0, - type=int, - help= - 'Specify the top_k value to use for MOE layers. Default to 1 if --moe_num_experts is set' - ) - parser.add_argument( - '--moe_tp_size', - type=int, - default=-1, - help= - 'N-way tensor parallelism size for MOE, default is tp_size, which will do tp-only for MoE' - ) - parser.add_argument( - '--moe_ep_size', - type=int, - default=-1, - help= - 'N-way expert parallelism size for MOE, default is 1, which will do tp-only for MoE' - ) - parser.add_argument( - '--moe_renorm_mode', - default=MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE, - type=int, - help= - 'Controls renormalization after gate logits. Check layers/moe.py for accepted values', - ) - parser.add_argument( - '--save_config_only', - action="store_true", - default=False, - help= - 'Only save the model config w/o read and converting weights, be careful, this is for debug only' - ) - - args = parser.parse_args() - # changing the default to be consistent as the cli help said. - if args.moe_num_experts and args.moe_top_k == 0: - args.moe_top_k = 1 - return args - - -def args_to_quantization(args: argparse.Namespace) -> QuantConfig: - '''return config dict with quantization info based on the command line args - ''' - quant_config = QuantConfig() - if args.use_weight_only: - if args.weight_only_precision == 'int8': - quant_config.quant_algo = QuantAlgo.W8A16 - - return quant_config - - -def args_to_build_options(args): - return { - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'disable_weight_only_quant_plugin': - args.disable_weight_only_quant_plugin - } - - -def from_cli_args(args): - n_kv_head = args.n_kv_head if args.n_kv_head is not None else args.n_head - config = { - 'architecture': "LlamaForCausalLM", - 'dtype': args.dtype, - 'logits_dtype': 'float32', - 'num_hidden_layers': args.n_layer, - 'num_attention_heads': args.n_head, - 'hidden_size': args.n_embd, - 'intermediate_size': args.inter_size, - 'num_key_value_heads': n_kv_head, - 'vocab_size': args.vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'max_position_embeddings': args.n_positions, - 'hidden_act': args.hidden_act, - 'rotary_base': args.rotary_base, - 'norm_epsilon': args.rms_norm_eps, - 'moe_num_experts': args.moe_num_experts, - 'moe_top_k': args.moe_top_k, - 'moe_normalization_mode': args.moe_renorm_mode, - 'mapping': { - 'world_size': args.tp_size * args.pp_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - 'moe_tp_size': args.moe_tp_size, - 'moe_ep_size': args.moe_ep_size, - }, - 'quantization': args_to_quantization(args).asdict() - } - config.update(args_to_build_options(args)) - return config - - -def preload_model(model_dir, weights_dir=None): - sys.path.append(model_dir) - from model import LanguageModelConfig, TransformerConfig - from runners import ModelRunner - if weights_dir and os.path.exists(weights_dir): - CKPT_PATH = weights_dir - else: - CKPT_PATH = os.path.join(model_dir, "checkpoints") - - grok_1_model = LanguageModelConfig( - vocab_size=128 * 1024, - pad_token=0, - eos_token=2, - sequence_len=8192, - embedding_init_scale=1.0, - output_multiplier_scale=0.5773502691896257, - embedding_multiplier_scale=78.38367176906169, - model=TransformerConfig( - emb_size=48 * 128, - widening_factor=8, - key_size=128, - num_q_heads=48, - num_kv_heads=8, - num_layers=64, - attn_output_multiplier=0.08838834764831845, - shard_activations=True, - # MoE. - num_experts=8, - num_selected_experts=2, - # Activation sharding. - data_axis="data", - model_axis="model", - ), - ) - - runner = ModelRunner( - model=grok_1_model, - bs_per_device=0.125, - checkpoint_path=CKPT_PATH, - ) - dummy_data = dict( - inputs=np.zeros((1, 256), dtype=np.int32), - targets=np.zeros((1, 256), dtype=np.int32), - ) - runner.transform_forward = True - runner.initialize(dummy_data, (1, 8), (1, 1)) - - params = runner.load_or_init(dummy_data) - - return params - - -def convert_and_save_xai(args): - model_dir = args.model_dir - load_by_shard = args.load_by_shard - world_size = args.tp_size * args.pp_size - if (args.moe_tp_size == -1 and args.moe_ep_size == -1): - # moe default to tp-only - args.moe_tp_size = args.tp_size - args.moe_ep_size = 1 - elif (args.moe_tp_size == -1): - args.moe_tp_size = args.tp_size // args.moe_ep_size - elif (args.moe_ep_size == -1): - args.moe_ep_size = args.tp_size // args.moe_tp_size - assert (args.moe_tp_size * args.moe_ep_size == args.tp_size - ), "moe_tp_size * moe_ep_size must equal to tp_size" - # Need to convert the cli args to the kay-value pairs and override them in the generate config dict. - # Ideally these fields will be moved out of the config and pass them into build API, keep them here for compatibility purpose for now, - # before the refactor is done. - override_fields = {} - quantization = args_to_quantization(args) - override_fields.update(args_to_build_options(args)) - - # When not loading by shard, preload one complete model and then slice per rank weights from this - # this saves the disk reloading time - xai_model = preload_model( - model_dir, args.weights_dir) if not args.load_by_shard else None - - def convert_and_save_rank(args, rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size, - moe_tp_size=args.moe_tp_size, - moe_ep_size=args.moe_ep_size) - grok = GrokForCausalLM.from_hugging_face( - model_dir, - args.dtype, - mapping=mapping, - quantization=quantization, - load_by_shard=load_by_shard, - override_fields=override_fields, - preloaded_model=xai_model, - ) - grok.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del grok - - execute(args.workers, [convert_and_save_rank] * world_size, args) - release_gc() - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - - args.tp_size * args.pp_size - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - if args.model_dir is None: # generate fake config.json - config = from_cli_args(args) - with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: - json.dump(config, f, indent=4) - else: # all other non-gptq paths from hf model - assert args.model_dir is not None - convert_and_save_xai(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/grok/requirements.txt b/examples/models/contrib/grok/requirements.txt deleted file mode 100644 index 3ab319c0efc5..000000000000 --- a/examples/models/contrib/grok/requirements.txt +++ /dev/null @@ -1,10 +0,0 @@ --f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score -sentencepiece==0.2.0 -jax[cuda12-pip]==0.4.29 -jaxlib[cuda12-pip]==0.4.29 -dm_haiku==0.0.12 diff --git a/examples/models/contrib/mmdit/README.md b/examples/models/contrib/mmdit/README.md deleted file mode 100644 index 4a3d80986566..000000000000 --- a/examples/models/contrib/mmdit/README.md +++ /dev/null @@ -1,95 +0,0 @@ -# MMDiT in SD 3 & SD 3.5 - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a [MMDiT](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sd3.py) in Stable Diffusion 3/3.5 with TensorRT-LLM. - -## Overview - -The TensorRT LLM implementation of MMDiT can be found in [tensorrt_llm/models/sd3/model.py](../../../../tensorrt_llm/models/mmdit_sd3/model.py). The TensorRT LLM MMDiT (SD 3/3.5) example code is located in [`examples/models/contrib/mmdit`](./). There are main files to build and run MMDiT with TensorRT-LLM: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert the MMDiT model into TensorRT LLM checkpoint format. -* [`sample.py`](./sample.py) to run the [diffusers](https://huggingface.co/docs/diffusers/index) pipeline with TensorRT engine(s) to generate images. - -## Support Matrix - -- [x] TP -- [x] CP -- [ ] FP8 - -## Usage - -The TensorRT LLM MMDiT example code locates at [examples/models/contrib/mmdit](./). It takes HuggingFace checkpoint as input, and builds the corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Build MMDiT TensorRT engine(s) - -This checkpoint will be converted to the TensorRT LLM checkpoint format by [`convert_checkpoint.py`](./convert_checkpoint.py). After that, we can build TensorRT engine(s) with the TensorRT LLM checkpoint. - -``` -# Convert to TRT-LLM -python convert_checkpoint.py --model_path='stabilityai/stable-diffusion-3.5-medium' -trtllm-build --checkpoint_dir=./tllm_checkpoint/ \ - --max_batch_size=2 \ - --remove_input_padding=disable \ - --bert_attention_plugin=auto -``` - -Set `--max_batch_size` to tell how many images at most you would like to generate. We disable `--remove_input_padding` since we don't need to padding MMDiT's patches. - -After build, we can find a `./engine_output` directory, it is ready for running MMDiT model with TensorRT LLM now. - -### Generate images - -A [`sample.py`](./sample.py) is provided to generated images with the optimized TensorRT engines. - -If using `float16` for inference, `FusedRMSNorm` from `Apex` used by T5-encoder should be disabled in the [huggingface/transformers](https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/t5/modeling_t5.py#L259) or just uninstall the `apex`: -```python -try: - from apex.normalization import FusedRMSNorm - - # [NOTE] Avoid using `FusedRMSNorm` for T5 encoder. - # T5LayerNorm = FusedRMSNorm # noqa - - logger.info("Discovered apex.normalization.FusedRMSNorm - will use it instead of T5LayerNorm") -except ImportError: - # using the normal T5LayerNorm - pass -except Exception: - logger.warning("discovered apex but it failed to load, falling back to T5LayerNorm") - pass - -ALL_LAYERNORM_LAYERS.append(T5LayerNorm) -``` - -Just run `python sample.py` and we can see an image named `sd3.5-mmdit.png` will be generated: -![sd3.5-mmdit.png](./assets/sd3.5-mmdit.png). - -### Tensor Parallel - -``` -# Convert to TRT-LLM -python convert_checkpoint.py --tp_size=2 --model_path='stabilityai/stable-diffusion-3.5-medium' -trtllm-build --checkpoint_dir=./tllm_checkpoint/ \ - --max_batch_size=2 \ - --remove_input_padding=disable \ - --bert_attention_plugin=auto -mpirun -n 2 --allow-run-as-root python sample.py "A capybara holding a sign that reads 'Hello World' in the forrest." -``` - -### Context Parallel - -Pipeline with CP is similar to that with TP, but it doesn't support `BertAttention` plugin. And make sure `tensorrt>=10.8.0.43`. - -``` -# Convert to TRT-LLM -python convert_checkpoint.py --tp_size=2 --model_path='stabilityai/stable-diffusion-3.5-medium' -trtllm-build --checkpoint_dir=./tllm_checkpoint/ \ - --max_batch_size=2 \ - --remove_input_padding=disable \ - --bert_attention_plugin=disable -mpirun -n 2 --allow-run-as-root python sample.py "A capybara holding a sign that reads 'Hello World' in the forrest." -``` diff --git a/examples/models/contrib/mmdit/assets/sd3.5-mmdit.png b/examples/models/contrib/mmdit/assets/sd3.5-mmdit.png deleted file mode 100644 index 9ccee43d85f5..000000000000 Binary files a/examples/models/contrib/mmdit/assets/sd3.5-mmdit.png and /dev/null differ diff --git a/examples/models/contrib/mmdit/convert_checkpoint.py b/examples/models/contrib/mmdit/convert_checkpoint.py deleted file mode 100644 index 3584ab309288..000000000000 --- a/examples/models/contrib/mmdit/convert_checkpoint.py +++ /dev/null @@ -1,123 +0,0 @@ -import argparse -import json -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import SD3Transformer2DModel - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_path', - type=str, - default="stabilityai/stable-diffusion-3.5-medium") - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--cp_size', - type=int, - default=1, - help='N-way context parallelism size') - parser.add_argument( - '--dtype', - type=str, - default='float16', - choices=['auto', 'float16', 'bfloat16', 'float32'], - help= - "The data type for the model weights and activations if not quantized. " - "If 'auto', the data type is automatically inferred from the source model; " - "however, if the source dtype is float32, it is converted to float16.") - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument('--log_level', type=str, default='info') - - args = parser.parse_args() - return args - - -def convert_and_save_hf(args): - world_size = args.tp_size * args.cp_size - - def convert_and_save_rank(args, rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - cp_size=args.cp_size) - tik = time.time() - mmdit = SD3Transformer2DModel.from_pretrained(args.model_path, - args.dtype, - mapping=mapping) - print(f'Total time of reading and converting: {time.time()-tik:.3f} s') - tik = time.time() - mmdit.save_checkpoint(args.output_dir, save_config=(rank == 0)) - # post-process checkpoint config - with open(f"{args.output_dir}/config.json", 'r') as f: - config = json.load(f) - config["model_path"] = args.model_path - config["use_pretrained_pos_emb"] = True if "pos_embed.pos_embed" in [ - name for name, _ in mmdit.named_parameters() - ] else False - with open(f"{args.output_dir}/config.json", 'w') as f: - json.dump(config, f, indent=4) - del mmdit - print(f'Total time of saving checkpoint: {time.time()-tik:.3f} s') - - execute(args.workers, [convert_and_save_rank] * world_size, args) - release_gc() - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - logger.set_level(args.log_level) - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - assert args.model_path is not None - convert_and_save_hf(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/mmdit/requirements.txt b/examples/models/contrib/mmdit/requirements.txt deleted file mode 100644 index 876bf841255d..000000000000 --- a/examples/models/contrib/mmdit/requirements.txt +++ /dev/null @@ -1,2 +0,0 @@ -transformers>=4.56.0 -diffusers>=0.32.2 diff --git a/examples/models/contrib/mmdit/sample.py b/examples/models/contrib/mmdit/sample.py deleted file mode 100644 index 31da097bd3fd..000000000000 --- a/examples/models/contrib/mmdit/sample.py +++ /dev/null @@ -1,266 +0,0 @@ -import argparse -import json -import os -from functools import wraps -from typing import List, Optional - -import numpy as np -import tensorrt as trt -import torch -from cuda import cudart -from diffusers import StableDiffusion3Pipeline -from diffusers.models.modeling_outputs import Transformer2DModelOutput - -import tensorrt_llm -from tensorrt_llm._utils import str_dtype_to_torch, trt_dtype_to_torch -from tensorrt_llm.logger import logger -from tensorrt_llm.models.mmdit_sd3.config import SD3Transformer2DModelConfig -from tensorrt_llm.runtime.session import Session - - -def CUASSERT(cuda_ret): - err = cuda_ret[0] - if err != cudart.cudaError_t.cudaSuccess: - raise RuntimeError( - f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t" - ) - if len(cuda_ret) > 1: - return cuda_ret[1:] - return None - - -class TllmMMDiT(object): - - def __init__(self, - config, - skip_layers: Optional[List[int]] = None, - debug_mode: bool = True, - stream: torch.cuda.Stream = None, - pos_embedder: torch.nn.Module = None): - self.dtype = config['pretrained_config']['dtype'] - # Update config with `skip_layers` - config['pretrained_config']['skip_layers'] = skip_layers - self.config = SD3Transformer2DModelConfig.from_dict( - config['pretrained_config']) - - rank = tensorrt_llm.mpi_rank() - world_size = config['pretrained_config']['mapping']['world_size'] - cp_size = config['pretrained_config']['mapping']['cp_size'] - tp_size = config['pretrained_config']['mapping']['tp_size'] - pp_size = config['pretrained_config']['mapping']['pp_size'] - gpus_per_node = config['pretrained_config']['mapping']['gpus_per_node'] - assert pp_size == 1 - self.mapping = tensorrt_llm.Mapping(world_size=world_size, - rank=rank, - cp_size=cp_size, - tp_size=tp_size, - pp_size=1, - gpus_per_node=gpus_per_node) - - local_rank = rank % self.mapping.gpus_per_node - self.device = torch.device(f'cuda:{local_rank}') - torch.cuda.set_device(self.device) - CUASSERT(cudart.cudaSetDevice(local_rank)) - - self.stream = stream - if self.stream is None: - self.stream = torch.cuda.Stream(self.device) - torch.cuda.set_stream(self.stream) - - engine_file = os.path.join(args.tllm_model_dir, f"rank{rank}.engine") - logger.info(f'Loading engine from {engine_file}') - with open(engine_file, "rb") as f: - engine_buffer = f.read() - - assert engine_buffer is not None - - self.session = Session.from_serialized_engine(engine_buffer) - - self.debug_mode = debug_mode - - self.inputs = {} - self.outputs = {} - self.buffer_allocated = False - - expected_tensor_names = [ - 'hidden_states', 'encoder_hidden_states', 'pooled_projections', - 'timestep', 'output' - ] - - found_tensor_names = [ - self.session.engine.get_tensor_name(i) - for i in range(self.session.engine.num_io_tensors) - ] - if not self.debug_mode and set(expected_tensor_names) != set( - found_tensor_names): - logger.error( - f"The following expected tensors are not found: {set(expected_tensor_names).difference(set(found_tensor_names))}" - ) - logger.error( - f"Those tensors in engine are not expected: {set(found_tensor_names).difference(set(expected_tensor_names))}" - ) - logger.error(f"Expected tensor names: {expected_tensor_names}") - logger.error(f"Found tensor names: {found_tensor_names}") - raise RuntimeError( - "Tensor names in engine are not the same as expected.") - if self.debug_mode: - self.debug_tensors = list( - set(found_tensor_names) - set(expected_tensor_names)) - - def _tensor_dtype(self, name): - # return torch dtype given tensor name for convenience - dtype = trt_dtype_to_torch(self.session.engine.get_tensor_dtype(name)) - return dtype - - def _setup(self, outputs_shape): - for i in range(self.session.engine.num_io_tensors): - name = self.session.engine.get_tensor_name(i) - if self.session.engine.get_tensor_mode( - name) == trt.TensorIOMode.OUTPUT: - shape = list(self.session.engine.get_tensor_shape(name)) - if self.debug_mode: - shape = list(self.session.engine.get_tensor_shape(name)) - if shape[0] == -1: - shape[0] = outputs_shape['output'][0] - else: - shape = outputs_shape[name] - self.outputs[name] = torch.empty(shape, - dtype=self._tensor_dtype(name), - device=self.device) - - self.buffer_allocated = True - - def cuda_stream_guard(func): - """Sync external stream and set current stream to the one bound to the session. Reset on exit. - """ - - @wraps(func) - def wrapper(self, *args, **kwargs): - external_stream = torch.cuda.current_stream() - if external_stream != self.stream: - external_stream.synchronize() - torch.cuda.set_stream(self.stream) - ret = func(self, *args, **kwargs) - if external_stream != self.stream: - self.stream.synchronize() - torch.cuda.set_stream(external_stream) - return ret - - return wrapper - - def __call__(self, *args, **kwargs): - return self.forward(*args, **kwargs) - - @cuda_stream_guard - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor = None, - pooled_projections: torch.Tensor = None, - timestep: torch.LongTensor = None, - block_controlnet_hidden_states: List = None, - joint_attention_kwargs=None, - return_dict: bool = True, - skip_layers: Optional[List[int]] = None, - ): - if block_controlnet_hidden_states is not None: - raise NotImplementedError("ControlNet is not supported yet.") - if joint_attention_kwargs is not None: - if joint_attention_kwargs.get("scale", None) is not None: - raise NotImplementedError("LoRA is not supported yet.") - if "ip_adapter_image_embeds" in joint_attention_kwargs: - raise NotImplementedError("IP-Adapter is not supported yet.") - if skip_layers is not None: - raise RuntimeWarning( - "Please initialize model with `skip_layers` instead of passing via `forward`." - ) - - self._setup(outputs_shape={'output': hidden_states.shape}) - if not self.buffer_allocated: - raise RuntimeError('Buffer not allocated, please call setup first!') - - inputs = { - 'hidden_states': - hidden_states.to(str_dtype_to_torch(self.dtype)), - 'encoder_hidden_states': - encoder_hidden_states.to(str_dtype_to_torch(self.dtype)), - 'pooled_projections': - pooled_projections.to(str_dtype_to_torch(self.dtype)), - 'timestep': - timestep.to(str_dtype_to_torch(self.dtype)), - } - for k, v in inputs.items(): - inputs[k] = v.cuda().contiguous() - self.inputs.update(**inputs) - self.session.set_shapes(self.inputs) - ok = self.session.run(self.inputs, self.outputs, - self.stream.cuda_stream) - - if not ok: - raise RuntimeError('Executing TRT engine failed!') - output = self.outputs['output'].to(hidden_states.device) - - if self.debug_mode: - torch.cuda.synchronize() - output_np = { - k: v.cpu().float().numpy() - for k, v in self.outputs.items() - } - np.savez("tllm_output.npz", **output_np) - - if not return_dict: - return (output, ) - else: - return Transformer2DModelOutput(sample=output) - - -def main(args): - tensorrt_llm.logger.set_level(args.log_level) - assert torch.cuda.is_available() - - config_file = os.path.join(args.tllm_model_dir, 'config.json') - with open(config_file) as f: - config = json.load(f) - - pipe = StableDiffusion3Pipeline.from_pretrained( - config['pretrained_config']['model_path'], torch_dtype=torch.float16) - pipe.to("cuda") - - # replace sd3.5 transformer with TRTLLM model - torch_pos_embed = pipe.transformer.pos_embed - del pipe.transformer - torch.cuda.empty_cache() - - # Load model: - model = TllmMMDiT(config, - debug_mode=args.debug_mode, - pos_embedder=torch_pos_embed) - - pipe.transformer = model - - image = pipe(args.prompt, - height=1024, - width=1024, - guidance_scale=4.5, - num_inference_steps=40, - generator=torch.Generator("cpu").manual_seed(0)).images[0] - if tensorrt_llm.mpi_rank() == 0: - image.save("sd3.5-mmdit.png") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - 'prompt', - nargs='*', - default= - "A capybara holding a sign that reads 'Hello World' in the forrest.", - help="Text prompt(s) to guide image generation") - parser.add_argument("--tllm_model_dir", - type=str, - default='./engine_outputs/') - parser.add_argument("--gpus_per_node", type=int, default=8) - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument("--debug_mode", action='store_true') - args = parser.parse_args() - main(args) diff --git a/examples/models/contrib/mpt/.gitignore b/examples/models/contrib/mpt/.gitignore deleted file mode 100644 index a6c46e9436fc..000000000000 --- a/examples/models/contrib/mpt/.gitignore +++ /dev/null @@ -1,3 +0,0 @@ -mpt* -*.log -c-model diff --git a/examples/models/contrib/mpt/README.md b/examples/models/contrib/mpt/README.md deleted file mode 100644 index 60aacbc5d1c2..000000000000 --- a/examples/models/contrib/mpt/README.md +++ /dev/null @@ -1,182 +0,0 @@ -# MPT - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document explains how to build the [MPT](https://huggingface.co/mosaicml/mpt-7b) model using TensorRT LLM and run on a single GPU and a single node with multiple GPUs. - -- [MPT](#mpt) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [MPT 7B](#mpt-7b) - - [1.1 Convert from HF Transformers in FP](#11-convert-from-hf-transformers-in-fp) - - [1.2 Convert from HF Transformers with weight-only quantization](#12-convert-from-hf-transformers-with-weight-only-quantization) - - [1.3 Convert from HF Transformers with SmoothQuant quantization](#13-convert-from-hf-transformers-with-smoothquant-quantization) - - [1.4 Convert from HF Transformers with INT8 KV cache quantization](#14-convert-from-hf-transformers-with-int8-kv-cache-quantization) - - [1.5 AWQ weight-only quantization with Modelopt](#15-awq-weight-only-quantization-with-modelopt) - - [1.6 FP8 Post-Training Quantization with Modelopt](#16-fp8-post-training-quantization-with-modelopt) - - [1.6 Weight-only quantization with Modelopt](#16-weight-only-quantization-with-modelopt) - - [1.7 SmoothQuant and INT8 KV cache with Modelopt](#17-smoothquant-and-int8-kv-cache-with-modelopt) - - [2.1 Build TensorRT engine(s)](#21-build-tensorrt-engines) - - [MPT 30B](#mpt-30b) - - [1. Convert weights from HF Transformers to TRTLLM format](#1-convert-weights-from-hf-transformers-to-trtllm-format) - - [2. Build TensorRT engine(s)](#2-build-tensorrt-engines) - - [3. Run TRT engine to check if the build was correct](#3-run-trt-engine-to-check-if-the-build-was-correct) - -## Overview - -The TensorRT LLM MPT implementation can be found in [`tensorrt_llm/models/mpt/model.py`](../../tensorrt_llm/models/mpt/model.py). The TensorRT LLM MPT example code is located in [`examples/models/contrib/mpt`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 - * FP8 (with FP8 KV Cache) - * INT8 & INT4 Weight-Only - * INT8 Smooth Quant - * INT4 AWQ - * Tensor Parallel - * MHA, MQA & GQA - * STRONGLY TYPED - -### MPT 7B - -Please install required packages first: - -```bash -pip install -r requirements.txt -``` - -The [`convert_checkpoint.py`](./convert_checkpoint.py) script allows you to convert weights from HF Transformers format to TRTLLM checkpoints. - -#### 1.1 Convert from HF Transformers in FP - -```bash -# Generate FP16 checkpoints. -python convert_checkpoint.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/fp16/ --dtype float16 - -# Generate FP32 checkpoints with TP=4. -python convert_checkpoint.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/fp32_tp4/ --dtype float32 --tp_size 4 -``` - -#### 1.2 Convert from HF Transformers with weight-only quantization - -```bash -# Use int8 weight-only quantization. -python convert_checkpoint.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/int8_wo/ --use_weight_only - -# Use int4 weight-only quantization. -python convert_checkpoint.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/int4_wo/ --use_weight_only --weight_only_precision int4 -``` - -#### 1.3 Convert from HF Transformers with SmoothQuant quantization - -```bash -# Use int8 smoothquant (weight and activation) quantization. -python convert_checkpoint.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/int8_sq/ --smoothquant 0.5 -``` - -#### 1.4 Convert from HF Transformers with INT8 KV cache quantization - -```bash -# Use int8 kv cache quantization. -python convert_checkpoint.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/fp16_int8kv/ --dtype float16 --calibrate_kv_cache -``` -***INT8-KV-cache can be used with SQ and Weight-only at the same time*** - - -***We now introduce Modelopt to do all quantization*** -First make sure Modelopt toolkit is installed (see [examples/quantization/README.md](/examples/quantization/README.md#preparation)) - -#### 1.5 AWQ weight-only quantization with Modelopt - -```bash -# INT4 AWQ quantization using Modelopt. -python ../../../quantization/quantize.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/int4_awq/ --qformat int4_awq -``` - -#### 1.6 FP8 Post-Training Quantization with Modelopt - -```bash -# FP8 quantization using Modelopt. -python ../../../quantization/quantize.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/fp8/ --qformat fp8 --kv_cache_dtype fp8 -``` - -#### 1.6 Weight-only quantization with Modelopt - -```bash -# INT8 Weight-only quantization using Modelopt with TP=2. -python ../../../quantization/quantize.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/int8_wo/ --qformat int8_wo --tp_size 2 - -# INT4 Weight-only quantization using Modelopt. -python ../../../quantization/quantize.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/int4_wo/ --qformat int4_wo -``` - -#### 1.7 SmoothQuant and INT8 KV cache with Modelopt - -```bash -# Use int4 awq quantization. -python ../../../quantization/quantize.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/sq_int8kv/ --qformat int8_sq --kv_cache_dtype int8 -``` -***INT8-KV-cache can also be used with Weight-only at the same time*** - - -### 2.1 Build TensorRT engine(s) - -All of the checkpoint generated by `convert_checkpoint.py` or `quantize.py` (Modelopt) can share the same building commands. - -```bash -# Build a single-GPU float16 engine using TRTLLM checkpoints. -trtllm-build --checkpoint_dir=./ckpts/mpt-7b/fp16 \ - --max_batch_size 32 \ - --max_input_len 1024 \ - --max_seq_len 1536 \ - --gemm_plugin float16 \ - --workers 1 \ - --output_dir ./trt_engines/mpt-7b/fp16 -``` - -### MPT 30B - -Same commands can be changed to convert MPT 30B to TRT LLM format. Below is an example to build MPT30B fp16 4-way tensor parallelized TRT engine - -#### 1. Convert weights from HF Transformers to TRTLLM format - -The [`convert_checkpoint.py`](./convert_checkpoint.py) script allows you to convert weights from HF Transformers format to TRTLLM format. - -```bash -python convert_checkpoint.py --model_dir mosaicml/mpt-30b --output_dir ./ckpts/mpt-30b/fp16_tp4/ --tp_szie 4 --dtype float16 -``` - -#### 2. Build TensorRT engine(s) - -Examples of build invocations: - -```bash -# Build 4-GPU MPT-30B float16 engines -trtllm-build --checkpoint_dir ./ckpts/mpt-30b/fp16_tp4 \ - --max_batch_size 32 \ - --max_input_len 1024 \ - --max_seq_len 1536 \ - --gemm_plugin float16 \ - --workers 4 \ - --output_dir ./trt_engines/mpt-30b/fp16_tp4 -``` - -#### 3. Run TRT engine to check if the build was correct - -```bash -# Run 4-GPU MPT-30B TRT engine on a sample input prompt -mpirun -n 4 --allow-run-as-root \ - python ../../../run.py --max_output_len 10 \ - --engine_dir ./trt_engines/mpt-30b/fp16/4-gpu/ \ - --tokenizer_dir mosaicml/mpt-30b -``` diff --git a/examples/models/contrib/mpt/convert_checkpoint.py b/examples/models/contrib/mpt/convert_checkpoint.py deleted file mode 100644 index d8af794a93d0..000000000000 --- a/examples/models/contrib/mpt/convert_checkpoint.py +++ /dev/null @@ -1,923 +0,0 @@ -import argparse -import copy -import functools -import json -import os -import time -import traceback -from collections import defaultdict -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Dict, Optional - -import numpy as np -import safetensors -import torch -import torch.nn as nn -from tqdm import tqdm -from transformers import AutoTokenizer, MptConfig, MptForCausalLM -from transformers.pytorch_utils import Conv1D - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.convert_utils import (generate_int8, get_weight, - load_calib_dataset, smooth_gemm, - split) -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - 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('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument('--logits_dtype', - type=str, - default='float32', - choices=['float16', 'float32']) - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - - parser.add_argument( - '--calib_dataset', - type=str, - default='ccdv/cnn_dailymail', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - 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( - '--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( - "--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("--dataset_cache_dir", - type=str, - default=None, - help="cache dir to load the hugging face dataset") - 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('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - args = parser.parse_args() - - return args - - -@torch.no_grad() -def capture_activation_range(model, - tokenizer, - dataset, - num_samples=1, - seq_len=512): - model.eval() - device = next(model.parameters()).device - act_scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) - - tokenizer.pad_token = tokenizer.eos_token - - 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].float() - - 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))) - - for i in tqdm(range(num_samples), desc="calibrating model"): - datapoint = dataset[i:i + 1] - line = copy.copy(datapoint) - line[0] = line[0] + ' TL;DR: ' - line[0] = line[0].strip() - line[0] = line[0].replace(" n't", "n't") - 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() - - return act_scales - - -@torch.no_grad() -def smooth_mpt_model(model, scales, alpha, mpt_qkv_para, mpt_smoother): - # Smooth the activation and weights with smoother = $\diag{s}$ - for name, module in model.named_modules(): - if not isinstance(module, type(model.transformer.blocks[0])): - continue - # qkv_proj - layer_name_qkv = name + ".attn.Wqkv" - weight = module.attn.Wqkv.weight - smoother = smooth_gemm(weight, scales[layer_name_qkv]["x"], - module.norm_1.weight, module.norm_1.bias, alpha) - scales[layer_name_qkv]["x"] = scales[layer_name_qkv]["x"] / smoother - scales[layer_name_qkv]["w"] = weight.abs().max(dim=1)[0] - # see transpose_weights function - mpt_qkv_para[layer_name_qkv] = weight.transpose(0, 1) - - # ================================================================= - layer_name = name + ".attn.out_proj" - smoother = smooth_gemm(module.attn.out_proj.weight, - scales[layer_name]["x"], None, None, alpha) - mpt_smoother[layer_name] = smoother.float() - - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.attn.out_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - fc1_layer_name = name + ".ffn.up_proj" - - smoother = smooth_gemm(module.ffn.up_proj.weight, - scales[fc1_layer_name]["x"], - module.norm_2.weight, module.norm_2.bias, alpha) - - scales[fc1_layer_name]["x"] = scales[fc1_layer_name]["x"] / smoother - scales[fc1_layer_name]["w"] = module.ffn.up_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - layer_name = name + ".ffn.down_proj" - smoother = smooth_gemm(module.ffn.down_proj.weight, - scales[layer_name]["x"], None, None, alpha) - mpt_smoother[layer_name] = smoother.float() - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.ffn.down_proj.weight.abs().max( - dim=1)[0] - - -def get_tllm_linear_sq_weight(vals, - prefix, - shape, - tensor_parallel, - is_qkv=False, - per_token=False, - per_channel=False, - last_prefix=None, - bias=None, - smoother_value=None, - smoother_shape=None, - rank=0, - cat_dim=0, - multi_query_mode=False): - results = {} - - def multi_query_split(data, local_dim, head_size, tp_size, cur_rank): - q, k, v = np.split(data, [local_dim, local_dim + head_size], axis=-1) - q_split = np.split(q, tp_size, axis=-1) - k_split = np.split(k, tp_size, axis=-1) - v_split = np.split(v, tp_size, axis=-1) - return [ - np.concatenate((q_split[ii], k_split[ii], v_split[ii]), axis=-1) - for ii in range(tp_size) - ][cur_rank] - - col_shape = shape if (is_qkv or per_channel) else [1, 1] - - if per_token: - if per_channel: - original_weights = np.array(vals["weight.int8.col"]) - else: - original_weights = np.array(vals["weight.int8"]) - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = np.split(original_weights, - tensor_parallel, - axis=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + - 'weight'] = torch.from_numpy(cur_weights).t().contiguous() - if smoother_value is None: - results[last_prefix] = torch.from_numpy( - np.array([1.0], dtype=np.float32)) - - if per_channel: - cur_per_channel_value = vals["scale_w_quant_orig.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_w_quant_orig.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_w_quant_orig"] - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split(vals["scale_w_quant_orig"], - tensor_parallel, - axis=cat_dim)[rank] - - results[prefix + 'per_channel_scale'] = torch.from_numpy( - np.array(cur_per_channel_value, - dtype=np.float32).reshape(col_shape)).contiguous() - else: - if per_channel: - original_weights = np.array(vals["weight.int8.col"]) - else: - original_weights = np.array(vals["weight.int8"]) - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = np.split(original_weights, - tensor_parallel, - axis=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + - 'weight'] = torch.from_numpy(cur_weights).t().contiguous() - - if per_channel: - cur_per_channel_value = vals["scale_y_accum_quant.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_y_accum_quant"] - # QKV is always per_channel - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant"], - tensor_parallel, - axis=cat_dim)[rank] - - results[prefix + 'per_channel_scale'] = torch.from_numpy( - np.array([cur_per_channel_value], - dtype=np.float32).reshape(col_shape)).contiguous() - - results[last_prefix] = torch.from_numpy( - np.array([vals['scale_x_orig_quant']], - dtype=np.float32)).contiguous() - - results[prefix + 'act_scale'] = torch.from_numpy( - np.array([[vals["scale_y_quant_orig"]]], - dtype=np.float32)).contiguous() - - if smoother_value is not None: - cur_smoother_value = np.split(smoother_value, - tensor_parallel, - axis=cat_dim)[rank] - results[prefix + 'smoother'] = cur_smoother_value.reshape( - smoother_shape).contiguous().to(torch.float32) - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def split_qkv_tp(qkv, n_head, n_kv_heads, n_hidden, tensor_parallel, rank): - """ - Splits the QKV matrix according to tensor parallelism - """ - kv_head_size = n_kv_heads * (n_hidden // n_head) - q, k, v = torch.split(qkv, [n_hidden, kv_head_size, kv_head_size], dim=0) - q = split(q, tensor_parallel, rank, dim=0) - k = split(k, tensor_parallel, rank, dim=0) - v = split(v, tensor_parallel, rank, dim=0) - return torch.concatenate([q, k, v], dim=0).contiguous() - - -def split_matrix(weight: torch.Tensor, tp_size: int, rank: int, - dim: int) -> torch.Tensor: - return split(weight, tp_size, rank, dim=dim) - - -def get_tllm_linear_weight( - weight: torch.Tensor, - prefix: str, - bias: Optional[torch.Tensor] = None, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[f'{prefix}.weight'] = processed_torch_weights - results[f'{prefix}.per_channel_scale'] = torch_weight_scales - else: - results[f'{prefix}.weight'] = weight.contiguous() - - if bias is not None: - results[f'{prefix}.bias'] = bias - - return results - - -def get_tllm_param( - param: torch.Tensor, - name: str, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if name.endswith('.weight') and use_weight_only: - v = param.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[name] = processed_torch_weights - results[name.replace('weight', - 'per_channel_scale')] = torch_weight_scales - else: - results[name] = param - - return results - - -def convert_hf_mpt_legacy(hf_model, - hf_config, - mapping, - rank=0, - dtype='float32', - use_parallel_embedding: bool = False, - sharding_dim: int = 0, - use_weight_only=False, - plugin_weight_only_quant_type='int8', - use_smooth_quant=False, - per_channel=False, - per_token=False, - int8_kv_cache=False, - act_range=[], - qkv_para=[], - smoother=[]): - weights = {} - tik = time.time() - tensor_parallel = mapping.tp_size - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - num_attention_heads = hf_model.config.n_heads - hidden_size = hf_model.config.d_model - vocab_size = hf_model.config.vocab_size - num_key_value_heads = hf_config.attn_config['kv_n_heads'] if 'kv_n_heads' in hf_config.attn_config \ - else hf_config.n_heads - multi_query_mode = (num_key_value_heads != num_attention_heads) - - for l in range(hf_model.config.n_layers): - prefix = f'transformer.blocks.{l}.' - tllm_prex = f'transformer.layers.{l}.' - - # attn.Wqkv -> attention.qkv - qkv_weight = get_weight(model_params, prefix + 'attn.Wqkv', dtype) - - if use_smooth_quant: - qkv_out_dim = qkv_weight.shape[0] - qkv_weight = qkv_weight.t().numpy() - if not multi_query_mode: - qkv_weight = qkv_weight.reshape(hidden_size, 3, hidden_size) - int8_weights = generate_int8(qkv_weight, - act_range.get(prefix + 'attn.Wqkv'), - is_qkv=True, - multi_query_mode=multi_query_mode) - weights.update( - get_tllm_linear_sq_weight(int8_weights, - tllm_prex + 'attention.qkv.', - [1, qkv_out_dim // tensor_parallel], - tensor_parallel, - is_qkv=True, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'input_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=rank, - cat_dim=-1, - multi_query_mode=multi_query_mode)) - else: - qkv_weight = split_qkv_tp(qkv_weight, num_attention_heads, - num_key_value_heads, hidden_size, - mapping.tp_size, mapping.tp_rank) - weights.update( - get_tllm_linear_weight(qkv_weight, tllm_prex + 'attention.qkv', - None, use_weight_only, - plugin_weight_only_quant_type)) - - if int8_kv_cache: - qkv_weight = get_weight(model_params, prefix + 'attn.Wqkv', dtype) - qkv_weight = qkv_weight.t().numpy() - if not multi_query_mode: - qkv_weight = qkv_weight.reshape(hidden_size, 3, hidden_size) - int8_weights = generate_int8(qkv_weight, - act_range.get(prefix + 'attn.Wqkv'), - is_qkv=True, - multi_query_mode=multi_query_mode) - weights[tllm_prex + - 'attention.kv_cache_scaling_factor'] = torch.from_numpy( - np.array([int8_weights['scale_y_quant_orig']], - dtype=np.float32)).contiguous() - - # attn.out_proj -> attention.dense - attn_dense_weight = get_weight(model_params, prefix + 'attn.out_proj', - dtype) - if use_smooth_quant: - attn_dense_weight = attn_dense_weight.t().numpy() - int8_weights = generate_int8( - attn_dense_weight, act_range.get(prefix + 'attn.out_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'attention.dense.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'attention.quantization_scaling_factor', - smoother_value=smoother[(prefix + 'attn.out_proj')], - smoother_shape=[1, hidden_size // tensor_parallel], - rank=rank, - cat_dim=0)) - else: - attn_dense_w = split_matrix(attn_dense_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(attn_dense_w, - tllm_prex + 'attention.dense', None, - use_weight_only, - plugin_weight_only_quant_type)) - - # ffn.up_proj -> mlp.fc - mlp_fc_weight = get_weight(model_params, prefix + 'ffn.up_proj', dtype) - if use_smooth_quant: - mlp_fc_weight = mlp_fc_weight.t().numpy() - int8_weights = generate_int8(mlp_fc_weight, - act_range.get(prefix + 'ffn.up_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.fc.', - [1, 4 * hidden_size // tensor_parallel], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=rank, - cat_dim=-1)) - else: - mlp_fc_weight = split_matrix(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(mlp_fc_weight, tllm_prex + 'mlp.fc', - None, use_weight_only, - plugin_weight_only_quant_type)) - - # ffn.down_proj -> mlp.proj - mlp_proj_weight = get_weight(model_params, prefix + 'ffn.down_proj', - dtype) - if use_smooth_quant: - mlp_proj_weight = mlp_proj_weight.t().numpy() - int8_weights = generate_int8( - mlp_proj_weight, act_range.get(prefix + 'ffn.down_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.proj.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'mlp.quantization_scaling_factor', - smoother_value=smoother[prefix + 'ffn.down_proj'], - smoother_shape=[1, 4 * hidden_size // tensor_parallel], - rank=rank, - cat_dim=0)) - else: - mlp_proj_weight = split_matrix(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(mlp_proj_weight, tllm_prex + 'mlp.proj', - None, use_weight_only, - plugin_weight_only_quant_type)) - - # input layer_norm - input_ln_weight = get_weight(model_params, prefix + 'norm_1', dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - - # post layer_norm - post_ln_weight = get_weight(model_params, prefix + 'norm_2', dtype) - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - - embed_w = get_weight(model_params, 'transformer.wte', dtype) - if mapping.is_first_pp_rank(): - # Embedding - if not use_parallel_embedding: - weights['transformer.vocab_embedding.weight'] = embed_w - else: - if sharding_dim == 0: - assert vocab_size % mapping.tp_size == 0 - else: - assert hidden_size % mapping.tp_size == 0 - weights['transformer.vocab_embedding.weight'] = split_matrix( - embed_w, mapping.tp_size, mapping.tp_rank, sharding_dim) - if mapping.is_last_pp_rank(): - # lm_head weight and bias - weights['lm_head.weight'] = split_matrix(embed_w.clone(), - mapping.tp_size, - mapping.tp_rank, - dim=0) - ln_f_w = get_weight(model_params, 'transformer.norm_f', dtype) - # ln_f weight and bias - weights['transformer.ln_f.weight'] = ln_f_w - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def convert_hf_mpt(hf_model: MptForCausalLM, - hf_config: MptConfig, - mapping: Mapping, - dtype: str = 'float32', - use_parallel_embedding: bool = False, - sharding_dim: int = 0, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8): - - weights = {} - tik = time.time() - - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - num_hidden_layers = hf_config.n_layers - num_head = hf_config.n_heads - num_kv_heads = getattr(hf_config.attn_config, 'kv_n_heads', - hf_config.n_heads) - hidden_size = hf_config.d_model - vocab_size = hf_config.vocab_size - - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - prefix = f'transformer.blocks.{l}' - tllm_prex = f'transformer.layers.{l-layers_range[0]}' - # Attention QKV (no bias) - qkv_w = get_weight(model_params, f'{prefix}.attn.Wqkv', dtype) - qkv_w = split_qkv_tp(qkv_w, num_head, num_kv_heads, hidden_size, - mapping.tp_size, mapping.tp_rank) - weights.update( - get_tllm_linear_weight(qkv_w, f'{tllm_prex}.attention.qkv', None, - use_weight_only, - plugin_weight_only_quant_type)) - # Attention dense (no bias) - attn_dense_weight = get_weight(model_params, f'{prefix}.attn.out_proj', - dtype) - attn_dense_w = split_matrix(attn_dense_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(attn_dense_w, f'{tllm_prex}.attention.dense', - None, use_weight_only, - plugin_weight_only_quant_type)) - # MLP fc_in (no bias) - mlp_fc_weight = get_weight(model_params, f'{prefix}.ffn.up_proj', dtype) - mlp_fc_w = split_matrix(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(mlp_fc_w, f'{tllm_prex}.mlp.fc', None, - use_weight_only, - plugin_weight_only_quant_type)) - # MLP fc_out (no bias) - mlp_proj_weight = get_weight(model_params, f'{prefix}.ffn.down_proj', - dtype) - mlp_proj_w = split_matrix(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(mlp_proj_w, f'{tllm_prex}.mlp.proj', None, - use_weight_only, - plugin_weight_only_quant_type)) - # input layer_norm - input_ln_weight = get_weight(model_params, f'{prefix}.norm_1', dtype) - weights[f'{tllm_prex}.input_layernorm.weight'] = input_ln_weight - - # post layer_norm - post_ln_weight = get_weight(model_params, f'{prefix}.norm_2', dtype) - weights[f'{tllm_prex}.post_layernorm.weight'] = post_ln_weight - - embed_w = get_weight(model_params, 'transformer.wte', dtype) - if mapping.is_first_pp_rank(): - # Embedding - if not use_parallel_embedding: - weights['transformer.vocab_embedding.weight'] = embed_w - else: - if sharding_dim == 0: - assert vocab_size % mapping.tp_size == 0 - else: - assert hidden_size % mapping.tp_size == 0 - weights['transformer.vocab_embedding.weight'] = split_matrix( - embed_w, mapping.tp_size, mapping.tp_rank, sharding_dim) - if mapping.is_last_pp_rank(): - # lm_head weight and bias - weights['lm_head.weight'] = split_matrix(embed_w.clone(), - mapping.tp_size, - mapping.tp_rank, - dim=0) - ln_f_w = get_weight(model_params, 'transformer.norm_f', dtype) - # ln_f weight and bias - weights['transformer.ln_f.weight'] = ln_f_w - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -if __name__ == '__main__': - emit_engine_arch_deprecation("convert_checkpoint.py") - # TODO(qijun): Currently, the convert script depends on a torch op: - # torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix, - # which is included in tensorrt_llm Python package. Otherwise, the convert - # script does not need to import tensorrt_llm. Will remove it after reimplementing - # the op with PyTorch. - print(tensorrt_llm.__version__) - args = parse_arguments() - world_size = args.tp_size * args.pp_size - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - world_size = args.tp_size * args.pp_size - quant_algo = None - plugin_weight_only_quant_type = None - if args.use_weight_only and args.weight_only_precision == 'int8': - plugin_weight_only_quant_type = torch.int8 - quant_algo = QuantAlgo.W8A16 - elif args.use_weight_only and args.weight_only_precision == 'int4': - plugin_weight_only_quant_type = torch.quint4x2 - quant_algo = QuantAlgo.W4A16 - - if args.smoothquant: - if args.per_token and args.per_channel: - quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN - elif not args.per_token and not args.per_channel: - quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN - elif not args.per_token and args.per_channel: - quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN - elif args.per_token and not args.per_channel: - quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN - - if args.calibrate_kv_cache: - kv_cache_quant_algo = QuantAlgo.INT8 - else: - kv_cache_quant_algo = None - - hf_config = MptConfig.from_pretrained(args.model_dir, - trust_remote_code=True) - num_kv_heads = getattr(hf_config.attn_config, 'kv_n_heads', - hf_config.n_heads) - config = { - 'architecture': hf_config.architectures[0], - 'dtype': args.dtype, - 'logits_dtype': args.logits_dtype, - 'vocab_size': hf_config.vocab_size, - 'hidden_size': hf_config.d_model, - 'intermediate_size': hf_config.d_model * 4, - 'num_hidden_layers': hf_config.n_layers, - 'num_attention_heads': hf_config.n_heads, - 'num_key_value_heads': num_kv_heads, - 'position_embedding_type': 'alibi', - 'hidden_act': 'gelu', - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'bias': (not hf_config.no_bias), - 'clip_qkv': hf_config.attn_config.clip_qkv, - 'alibi_bias_max': hf_config.attn_config.alibi_bias_max - } - - with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: - json.dump(config, f, indent=4) - - hf_model = MptForCausalLM.from_pretrained(args.model_dir, - device_map="auto", - dtype=getattr(torch, args.dtype)) - - act_range = {} - mpt_qkv_para = {} - # smoother for inputs of self_attn.o_proj and mlp.down_proj - mpt_smoother = {} - if args.smoothquant is not None or args.calibrate_kv_cache: - tokenizer = AutoTokenizer.from_pretrained(args.model_dir, - padding_side='left') - dataset = load_calib_dataset(args.calib_dataset, - cache_dir=args.dataset_cache_dir) - - act_range = capture_activation_range(hf_model, tokenizer, dataset) - if args.smoothquant is not None: - smooth_mpt_model(hf_model, act_range, args.smoothquant, - mpt_qkv_para, mpt_smoother) - - def covert_and_save(rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - - if args.smoothquant is not None or args.calibrate_kv_cache: - weights = convert_hf_mpt_legacy( - hf_model, - hf_config, - mapping, - rank, - dtype=args.dtype, - use_parallel_embedding=args.use_parallel_embedding, - sharding_dim=args.embedding_sharding_dim, - use_weight_only=args.use_weight_only, - plugin_weight_only_quant_type=plugin_weight_only_quant_type, - use_smooth_quant=(args.smoothquant is not None), - per_channel=args.per_channel, - per_token=args.per_token, - int8_kv_cache=args.calibrate_kv_cache, - act_range=act_range, - qkv_para=mpt_qkv_para, - smoother=mpt_smoother) - else: - weights = convert_hf_mpt( - hf_model, - hf_config, - mapping, - dtype=args.dtype, - use_parallel_embedding=args.use_parallel_embedding, - sharding_dim=args.embedding_sharding_dim, - use_weight_only=args.use_weight_only, - plugin_weight_only_quant_type=plugin_weight_only_quant_type) - - safetensors.torch.save_file( - weights, os.path.join(args.output_dir, f'rank{rank}.safetensors')) - - if args.workers == 1: - for rank in range(world_size): - covert_and_save(rank) - else: - with ThreadPoolExecutor(max_workers=args.workers) as p: - futures = [ - p.submit(covert_and_save, rank) for rank in range(world_size) - ] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - del hf_model - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') diff --git a/examples/models/contrib/mpt/requirements.txt b/examples/models/contrib/mpt/requirements.txt deleted file mode 100644 index 2f8713d8658a..000000000000 --- a/examples/models/contrib/mpt/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score diff --git a/examples/models/contrib/opt/README.md b/examples/models/contrib/opt/README.md deleted file mode 100644 index 609c7ebfb910..000000000000 --- a/examples/models/contrib/opt/README.md +++ /dev/null @@ -1,252 +0,0 @@ -# OPT - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document explains how to build the [OPT](https://huggingface.co/docs/transformers/model_doc/opt) model using TensorRT LLM and run on a single GPU, a single node with -multiple GPUs or multiple nodes with multiple GPUs. - -- [OPT](#opt) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [1. Download weights from HuggingFace Transformers](#1-download-weights-from-huggingface-transformers) - - [2. Convert weights from HF Transformers to TensorRT LLM format](#2-convert-weights-from-hf-transformers-to-tensorrt-llm-format) - - [3. Build TensorRT engine(s)](#3-build-tensorrt-engines) - - [4. Summarization using the OPT model](#4-summarization-using-the-opt-model) - - [Fused MultiHead Attention (FMHA)](#fused-multihead-attention-fmha) - - [Tensor Parallelism for Embedding Lookup Table.](#tensor-parallelism-for-embedding-lookup-table) - - [1. Enable this feature](#1-enable-this-feature) - - [2. Choose the dimension for tensor parallelism](#2-choose-the-dimension-for-tensor-parallelism) - -## 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/models/contrib/opt`](./). There is one file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format - -In addition, there are two shared files in the parent folder [`examples`](../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 - * INT8 & INT4 Weight-Only - * Tensor Parallel - -## Usage - -The next two sections describe how to convert the weights from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) -format to the TensorRT LLM format. - -### 1. Download weights from HuggingFace Transformers - -You have to make sure `git-lfs` is properly installed to load the checkpoints. - -```bash -pip install -r requirements.txt && sudo apt-get install git-lfs -``` - -There are four different checkpoints available. Use one of the following commands to fetch the checkpoint you are interested in. - -```bash -# OPT-125M -git-lfs clone https://huggingface.co/facebook/opt-125m - -# OPT-350M -git-lfs clone https://huggingface.co/facebook/opt-350m - -# OPT-2.7B -git-lfs clone https://huggingface.co/facebook/opt-2.7b - -# OPT-66B -git-lfs clone https://huggingface.co/facebook/opt-66b -``` - -### 2. Convert weights from HF Transformers to TensorRT LLM format - -```bash -# OPT-125M -python3 convert_checkpoint.py --model_dir ./opt-125m \ - --dtype float16 \ - --output_dir ./opt/125M/trt_ckpt/fp16/1-gpu/ - -# OPT-350M -python3 convert_checkpoint.py --model_dir ./opt-350m \ - --dtype float16 \ - --output_dir ./opt/350M/trt_ckpt/fp16/1-gpu/ - -# OPT-2.7B -python3 convert_checkpoint.py --model_dir ./opt-2.7b \ - --dtype float16 \ - --output_dir ./opt/2.7B/trt_ckpt/fp16/1-gpu/ - -# OPT-66B -python3 convert_checkpoint.py --model_dir ./opt-66b \ - --dtype float16 \ - --tp_size 4 \ - --output_dir ./opt/66B/trt_ckpt/fp16/4-gpu/ \ - --workers 2 -``` - -### 3. Build TensorRT engine(s) - -```bash -# OPT-125M -trtllm-build --checkpoint_dir ./opt/125M/trt_ckpt/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./opt/125M/trt_engines/fp16/1-gpu/ - -# OPT-350M -trtllm-build --checkpoint_dir ./opt/350M/trt_ckpt/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./opt/350M/trt_engines/fp16/1-gpu/ - -# OPT-2.7B -trtllm-build --checkpoint_dir ./opt/2.7B/trt_ckpt/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./opt/2.7B/trt_engines/fp16/1-gpu/ - -# OPT-66B -trtllm-build --checkpoint_dir ./opt/66B/trt_ckpt/fp16/4-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./opt/66B/trt_engines/fp16/4-gpu/ \ - --workers 2 -``` - -### 4. Summarization using the OPT model - -The following section describes how to run a TensorRT LLM OPT model to summarize the articles from the -[cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. For each summary, the script can compute the -[ROUGE](https://en.wikipedia.org/wiki/ROUGE_(metric)) scores and use the `ROUGE-1` score to validate the implementation. -The script can also perform the same summarization using the HF OPT model. - -```bash -# OPT-125M -python3 ../../../summarize.py --engine_dir ./opt/125M/trt_engines/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 - -# OPT-350M -python3 ../../../summarize.py --engine_dir ./opt/350M/trt_engines/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 - -# OPT-2.7B -python3 ../../../summarize.py --engine_dir ./opt/2.7B/trt_engines/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=20 - -# OPT-66B -mpirun -n 4 --allow-run-as-root \ - python3 ../../../summarize.py --engine_dir ./opt/66B/trt_engines/fp16/4-gpu/ \ - --batch_size 1 \ - --test_trt_llm \ - --hf_model_dir opt-66b \ - --data_type fp16 \ - --check_accuracy \ - --tensorrt_llm_rouge1_threshold=20 -``` - -#### Fused MultiHead Attention (FMHA) - -You can enable the FMHA kernels for OPT by adding `--enable_context_fmha` to the invocation of `trtllm-build`. Note that it is disabled by default because of possible accuracy issues due to the use of Flash Attention. - -If you find that the default fp16 accumulation (`--context_fmha`) cannot meet the requirement, you can try to enable fp32 accumulation by adding `--enable_context_fmha_fp32_acc` to the inference command (`run.py` or `summarize.py`). However, it is expected to see performance drop. - -Note `--context_fmha` has to be used together with `--gpt_attention_plugin float16`. - -## Tensor Parallelism for Embedding Lookup Table. -Since the embedding lookup table can be several gigabytes in size. We can distribute this weight across multiple GPUs in order to reduce the memory consumption per GPU. - -### 1. Enable this feature -To enable this feature, add the flag `--use_parallel_embedding` to `trtllm-build`. - -### 2. Choose the dimension for tensor parallelism - -Assume the size of embedding lookup table is (vocab\_size \* hidden\_size), we can shard it along the vocab\_size (`--embedding_sharding_dim 0`) or hidden\_size (`--embedding_sharding_dim 1`) dimension. - -2.1 To shard the embedding lookup table along the hidden\_size dimension, set the flag `--use_parallel_embedding --embedding_sharding_dim 1`. Here is an example: - -```Bash -python3 convert_checkpoint.py --model_dir ./opt-125m \ - --dtype float16 \ - --output_dir ./opt/125M/trt_ckpt/fp16/2-gpu/ \ - --tp_size 2 \ - --use_parallel_embedding \ - --embedding_sharding_dim 1 -``` -2.2 To shard the embedding lookup table along the vocab\_size dimension, set the flag `--use_parallel_embedding --embedding_sharding_dim 0`. - -Meanwhile, we provide a lookup plugin to support tensor parallelism on vocab\_size dimension. - -- An example of sharing along vocab\_size dimension with lookup plugin: - -```Bash -python3 convert_checkpoint.py --model_dir ./opt-125m \ - --dtype float16 \ - --output_dir ./opt/125M/trt_ckpt/fp16/2-gpu/ \ - --tp_size 2 \ - --use_parallel_embedding \ - --embedding_sharding_dim 0 - -trtllm-build --checkpoint_dir ./opt/125M/trt_ckpt/fp16/2-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./opt/125M/trt_engines/fp16/2-gpu/ \ - --workers 2 - -mpirun -n 2 --allow-run-as-root \ - python3 ../../../summarize.py --engine_dir ./opt/125M/trt_engines/fp16/2-gpu/ \ - --batch_size 1 \ - --test_trt_llm \ - --hf_model_dir opt-125m \ - --data_type fp16 \ - --check_accuracy \ - --tensorrt_llm_rouge1_threshold=14 -``` - -- An example of sharing along vocab\_size dimension without lookup plugin: - -```Bash -trtllm-build --checkpoint_dir ./opt/125M/trt_ckpt/fp16/2-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./opt/125M/trt_engines/fp16/2-gpu/ \ - --workers 2 -``` diff --git a/examples/models/contrib/opt/convert_checkpoint.py b/examples/models/contrib/opt/convert_checkpoint.py deleted file mode 100644 index 37b3b610f24e..000000000000 --- a/examples/models/contrib/opt/convert_checkpoint.py +++ /dev/null @@ -1,361 +0,0 @@ -import argparse -import json -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import safetensors -import torch -from transformers import AutoModelForCausalLM, Blip2ForConditionalGeneration - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import pad_vocab_size -from tensorrt_llm.models.convert_utils import (get_weight, get_weight_and_bias, - split, split_matrix_tp, - split_qkv_bias_tp, split_qkv_tp) -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument( - '--model_type', - type=str, - default='opt', - choices=['opt', 'blip2'], - help= - 'Multimodal type when this script is used for multimodal conversion.') - 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('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - 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_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - args = parser.parse_args() - - return args - - -def split_embedding( - param: torch.Tensor, - tp_size: int, - tp_rank: int, - use_parallel_embedding: bool = False, - sharding_dim: int = 0, -) -> torch.Tensor: - if param is None: - return None - if not use_parallel_embedding: - return param - - vocab_size, hidden_size = param.size() - if sharding_dim == 0: - if vocab_size % tp_size != 0: - vocab_size_padded = pad_vocab_size(vocab_size, tp_size) - pad_width = vocab_size_padded - vocab_size - param = torch.nn.functional.pad(param, (0, 0, 0, pad_width), - value=0) - else: - assert hidden_size % tp_size == 0 - return split(param, tp_size, tp_rank, dim=sharding_dim) - - -def get_tllm_linear_weight(weight, - prefix, - bias=None, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8): - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[prefix + 'weight'] = processed_torch_weights - results[prefix + 'per_channel_scale'] = torch_weight_scales - else: - results[prefix + 'weight'] = weight.contiguous() - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def convert_hf_opt(hf_model, - rank=0, - tensor_parallel=1, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8): - - weights = {} - tik = time.time() - - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - do_layer_norm_before = hf_model.config.do_layer_norm_before - num_attention_heads = hf_model.config.num_attention_heads - hidden_size = hf_model.config.hidden_size - - for l in range(hf_model.config.num_hidden_layers): - prefix = f'model.decoder.layers.{l}.' - tllm_prex = f'transformer.layers.{l}.' - - q_weight, q_bias = get_weight_and_bias(model_params, - prefix + 'self_attn.q_proj', - dtype) - k_weight, k_bias = get_weight_and_bias(model_params, - prefix + 'self_attn.k_proj', - dtype) - v_weight, v_bias = get_weight_and_bias(model_params, - prefix + 'self_attn.v_proj', - dtype) - qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0) - split_v = split_qkv_tp(qkv_weight, num_attention_heads, hidden_size, - tensor_parallel, rank) - qkv_bias = torch.cat([q_bias, k_bias, v_bias], dim=0) - bias = split_qkv_bias_tp(qkv_bias, num_attention_heads, hidden_size, - tensor_parallel, rank) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.qkv.', bias, - use_weight_only, - plugin_weight_only_quant_type)) - - attn_dense_weight, attn_dense_bias = get_weight_and_bias( - model_params, prefix + 'self_attn.out_proj', dtype) - split_v = split_matrix_tp(attn_dense_weight, - tensor_parallel, - rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.dense.', - attn_dense_bias, use_weight_only, - plugin_weight_only_quant_type)) - - mlp_fc_weight, mlp_fc_bias = get_weight_and_bias( - model_params, prefix + 'fc1', dtype) - split_v = split_matrix_tp(mlp_fc_weight, tensor_parallel, rank, dim=0) - bias = split_matrix_tp(mlp_fc_bias, tensor_parallel, rank, dim=0) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.fc.', bias, - use_weight_only, - plugin_weight_only_quant_type)) - - mlp_proj_weight, mlp_proj_bias = get_weight_and_bias( - model_params, prefix + 'fc2', dtype) - split_v = split_matrix_tp(mlp_proj_weight, tensor_parallel, rank, dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.proj.', - mlp_proj_bias, use_weight_only, - plugin_weight_only_quant_type)) - - # Layer norms do not use tensor parallelism - input_ln_weight, input_ln_bias = get_weight_and_bias( - model_params, prefix + 'self_attn_layer_norm', dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - weights[tllm_prex + 'input_layernorm.bias'] = input_ln_bias - - post_ln_weight, post_ln_bias = get_weight_and_bias( - model_params, prefix + 'final_layer_norm', dtype) - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - weights[tllm_prex + 'post_layernorm.bias'] = post_ln_bias - - embed_w = get_weight(model_params, 'model.decoder.embed_tokens', dtype) - if 'model.decoder.project_in.weight' in model_params.keys(): - project_in = get_weight(model_params, 'model.decoder.project_in', dtype) - project_out = get_weight(model_params, 'model.decoder.project_out', - dtype) - lm_head_w = torch.matmul(embed_w.float(), project_out.float()).to(dtype) - embed_w = torch.matmul(embed_w.float(), - project_in.t().float()).to(dtype) - elif 'lm_head.weight' in model_params.keys(): - lm_head_w = get_weight(model_params, 'lm_head', dtype) - else: - lm_head_w = embed_w.clone() - - weights['lm_head.weight'] = split_matrix_tp(lm_head_w, - tensor_parallel, - rank, - dim=0) - - weights['transformer.vocab_embedding.weight'] = split_embedding( - embed_w, - tp_size=tensor_parallel, - tp_rank=rank, - use_parallel_embedding=use_parallel_embedding, - sharding_dim=sharding_dim) - - embed_p = get_weight(model_params, 'model.decoder.embed_positions', dtype) - weights['transformer.position_embedding.weight'] = split_embedding( - embed_p[2:, :], - tp_size=tensor_parallel, - tp_rank=rank, - use_parallel_embedding=use_parallel_embedding, - sharding_dim=sharding_dim) - - if do_layer_norm_before: - ln_f_w, ln_f_b = get_weight_and_bias(model_params, - 'model.decoder.final_layer_norm', - dtype) - weights['transformer.ln_f.weight'] = ln_f_w - weights['transformer.ln_f.bias'] = ln_f_b - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -if __name__ == '__main__': - emit_engine_arch_deprecation("convert_checkpoint.py") - # TODO(qijun): Currently, the convert script depends on a torch op: - # torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix, - # which is included in tensorrt_llm Python package. Otherwise, the convert - # script does not need to import tensorrt_llm. Will remove it after reimplementing - # the op with PyTorch. - print(tensorrt_llm.__version__) - args = parse_arguments() - world_size = args.tp_size * args.pp_size - assert args.pp_size == 1, "Pipeline parallelism is not supported." - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - if args.model_type == 'opt': - hf_model = AutoModelForCausalLM.from_pretrained(args.model_dir, - dtype="auto") - elif args.model_type == 'blip2': - hf_model = Blip2ForConditionalGeneration.from_pretrained( - args.model_dir, dtype="auto").language_model - - hf_config = hf_model.config - if hf_config.hidden_size != hf_config.word_embed_proj_dim: - args.use_parallel_embedding = False - - quant_algo = None - plugin_weight_only_quant_type = None - if args.use_weight_only and args.weight_only_precision == 'int8': - plugin_weight_only_quant_type = torch.int8 - quant_algo = QuantAlgo.W8A16 - elif args.use_weight_only and args.weight_only_precision == 'int4': - plugin_weight_only_quant_type = torch.quint4x2 - quant_algo = QuantAlgo.W4A16 - - config = { - 'architecture': hf_config.architectures[0], - 'dtype': args.dtype, - 'num_hidden_layers': hf_config.num_hidden_layers, - 'num_attention_heads': hf_config.num_attention_heads, - 'hidden_size': hf_config.hidden_size, - 'vocab_size': hf_config.vocab_size, - 'position_embedding_type': 'learned_absolute', - 'max_position_embeddings': hf_config.max_position_embeddings, - 'hidden_act': hf_config.activation_function, - 'quantization': { - 'quant_algo': quant_algo - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'do_layer_norm_before': hf_config.do_layer_norm_before, - } - - with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: - json.dump(config, f, indent=4) - - def covert_and_save(rank): - weights = convert_hf_opt( - hf_model, - rank, - world_size, - dtype=args.dtype, - use_weight_only=args.use_weight_only, - plugin_weight_only_quant_type=plugin_weight_only_quant_type, - use_parallel_embedding=args.use_parallel_embedding, - sharding_dim=args.embedding_sharding_dim) - safetensors.torch.save_file( - weights, os.path.join(args.output_dir, f'rank{rank}.safetensors')) - - if args.workers == 1: - for rank in range(world_size): - covert_and_save(rank) - else: - with ThreadPoolExecutor(max_workers=args.workers) as p: - futures = [ - p.submit(covert_and_save, rank) for rank in range(world_size) - ] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') diff --git a/examples/models/contrib/opt/requirements.txt b/examples/models/contrib/opt/requirements.txt deleted file mode 100644 index 2f8713d8658a..000000000000 --- a/examples/models/contrib/opt/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score diff --git a/examples/models/contrib/stdit/README.md b/examples/models/contrib/stdit/README.md deleted file mode 100644 index 42116be7428e..000000000000 --- a/examples/models/contrib/stdit/README.md +++ /dev/null @@ -1,94 +0,0 @@ -# STDiT in OpenSoRA - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a STDiT in [OpenSoRA](https://github.com/hpcaitech/Open-Sora/tree/main) with TensorRT-LLM. - -## Overview - -The TensorRT LLM implementation of STDiT can be found in [tensorrt_llm/models/stdit/model.py](../../../../tensorrt_llm/models/stdit/model.py). The TensorRT LLM STDiT (OpenSoRA) example code is located in [`examples/models/contrib/stdit`](./). There are main files to build and run STDiT with TensorRT-LLM: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert the STDiT model into TensorRT LLM checkpoint format. -* [`sample.py`](./sample.py) to run the pipeline with TensorRT engine(s) to generate videos. - -## Support Matrix - -- [x] TP -- [ ] CP -- [ ] FP8 - -## Usage - -The TensorRT LLM STDiT example code locates at [examples/models/contrib/stdit](./). It takes HuggingFace checkpoint as input, and builds the corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Requirements - -Please install required packages first: - -```bash -pip install -r requirements.txt -# ColossalAI is also needed for text encoder. -pip install colossalai --no-deps -``` - -### Build STDiT TensorRT engine(s) - -This checkpoint will be converted to the TensorRT LLM checkpoint format by [`convert_checkpoint.py`](./convert_checkpoint.py). After that, we can build TensorRT engine(s) with the TensorRT LLM checkpoint. The pretrained checkpoint can be downloaded from [here](https://huggingface.co/hpcai-tech/OpenSora-STDiT-v3). - -```bash -# Convert to TRT-LLM -python convert_checkpoint.py --timm_ckpt= -# Build engine -trtllm-build --checkpoint_dir=tllm_checkpoint/ \ - --max_batch_size=2 \ - --gemm_plugin=float16 \ - --kv_cache_type=disabled \ - --remove_input_padding=enable \ - --gpt_attention_plugin=auto \ - --bert_attention_plugin=auto \ - --context_fmha=enable -``` - -After build, we can find a `./engine_output` directory, it is ready for running STDiT model with TensorRT LLM now. - -### Generate videos - -A [`sample.py`](./sample.py) is provided to generated videos with the optimized TensorRT engines. - -```bash -python sample.py "a beautiful waterfall" -``` - -And we can see a video named `sample_outputs/sample_0000.mp4` will be generated: - - - -### Tensor Parallel - -We can levaerage tensor parallel to further reduce latency and memory consumption on each GPU. - -```bash -# Convert to TRT-LLM -python convert_checkpoint.py --tp_size=2 --timm_ckpt= -# Build engines -trtllm-build --checkpoint_dir=tllm_checkpoint/ \ - --max_batch_size=2 \ - --gemm_plugin=float16 \ - --kv_cache_type=disabled \ - --remove_input_padding=enable \ - --gpt_attention_plugin=auto \ - --bert_attention_plugin=auto \ - --context_fmha=enable -# Run example -mpirun -n 2 --allow-run-as-root python sample.py "a beautiful waterfall" -``` - -### Context Parallel - -Not supported yet. diff --git a/examples/models/contrib/stdit/aspect.py b/examples/models/contrib/stdit/aspect.py deleted file mode 100644 index fe629a5da6ec..000000000000 --- a/examples/models/contrib/stdit/aspect.py +++ /dev/null @@ -1,526 +0,0 @@ -# Copyright 2024 HPC-AI Technology Inc. - -# 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. - -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# reference: https://github.com/hpcaitech/Open-Sora/blob/main/opensora/datasets/aspect.py - -import math - - -# computation -def get_h_w(a, ts, eps=1e-4): - h = (ts * a)**0.5 - h = h + eps - h = math.ceil(h) if math.ceil(h) % 2 == 0 else math.floor(h) - w = h / a - w = w + eps - w = math.ceil(w) if math.ceil(w) % 2 == 0 else math.floor(w) - return h, w - - -def get_aspect_ratios_dict(ars, ts=360 * 640): - est = {f"{a:.2f}": get_h_w(a, ts) for a in ars} - return est - - -def get_ar(ratio): - h, w = ratio.split(":") - return int(h) / int(w) - - -# H:W -ASPECT_RATIO_MAP = { - "3:8": "0.38", - "9:21": "0.43", - "12:25": "0.48", - "1:2": "0.50", - "9:17": "0.53", - "27:50": "0.54", - "9:16": "0.56", - "5:8": "0.62", - "2:3": "0.67", - "3:4": "0.75", - "1:1": "1.00", - "4:3": "1.33", - "3:2": "1.50", - "16:9": "1.78", - "17:9": "1.89", - "2:1": "2.00", - "50:27": "2.08", -} - -AR = [get_ar(ratio) for ratio in ASPECT_RATIO_MAP.keys()] - -# computed from above code -# S = 8294400 -ASPECT_RATIO_4K = { - "0.38": (1764, 4704), - "0.43": (1886, 4400), - "0.48": (1996, 4158), - "0.50": (2036, 4072), - "0.53": (2096, 3960), - "0.54": (2118, 3918), - "0.62": (2276, 3642), - "0.56": (2160, 3840), # base - "0.67": (2352, 3528), - "0.75": (2494, 3326), - "1.00": (2880, 2880), - "1.33": (3326, 2494), - "1.50": (3528, 2352), - "1.78": (3840, 2160), - "1.89": (3958, 2096), - "2.00": (4072, 2036), - "2.08": (4156, 1994), -} - -# S = 3686400 -ASPECT_RATIO_2K = { - "0.38": (1176, 3136), - "0.43": (1256, 2930), - "0.48": (1330, 2770), - "0.50": (1358, 2716), - "0.53": (1398, 2640), - "0.54": (1412, 2612), - "0.56": (1440, 2560), # base - "0.62": (1518, 2428), - "0.67": (1568, 2352), - "0.75": (1662, 2216), - "1.00": (1920, 1920), - "1.33": (2218, 1664), - "1.50": (2352, 1568), - "1.78": (2560, 1440), - "1.89": (2638, 1396), - "2.00": (2716, 1358), - "2.08": (2772, 1330), -} - -# S = 2073600 -ASPECT_RATIO_1080P = { - "0.38": (882, 2352), - "0.43": (942, 2198), - "0.48": (998, 2080), - "0.50": (1018, 2036), - "0.53": (1048, 1980), - "0.54": (1058, 1958), - "0.56": (1080, 1920), # base - "0.62": (1138, 1820), - "0.67": (1176, 1764), - "0.75": (1248, 1664), - "1.00": (1440, 1440), - "1.33": (1662, 1246), - "1.50": (1764, 1176), - "1.78": (1920, 1080), - "1.89": (1980, 1048), - "2.00": (2036, 1018), - "2.08": (2078, 998), -} - -# S = 921600 -ASPECT_RATIO_720P = { - "0.38": (588, 1568), - "0.43": (628, 1466), - "0.48": (666, 1388), - "0.50": (678, 1356), - "0.53": (698, 1318), - "0.54": (706, 1306), - "0.56": (720, 1280), # base - "0.62": (758, 1212), - "0.67": (784, 1176), - "0.75": (832, 1110), - "1.00": (960, 960), - "1.33": (1108, 832), - "1.50": (1176, 784), - "1.78": (1280, 720), - "1.89": (1320, 698), - "2.00": (1358, 680), - "2.08": (1386, 666), -} - -# S = 409920 -ASPECT_RATIO_480P = { - "0.38": (392, 1046), - "0.43": (420, 980), - "0.48": (444, 925), - "0.50": (452, 904), - "0.53": (466, 880), - "0.54": (470, 870), - "0.56": (480, 854), # base - "0.62": (506, 810), - "0.67": (522, 784), - "0.75": (554, 738), - "1.00": (640, 640), - "1.33": (740, 555), - "1.50": (784, 522), - "1.78": (854, 480), - "1.89": (880, 466), - "2.00": (906, 454), - "2.08": (924, 444), -} - -# S = 230400 -ASPECT_RATIO_360P = { - "0.38": (294, 784), - "0.43": (314, 732), - "0.48": (332, 692), - "0.50": (340, 680), - "0.53": (350, 662), - "0.54": (352, 652), - "0.56": (360, 640), # base - "0.62": (380, 608), - "0.67": (392, 588), - "0.75": (416, 554), - "1.00": (480, 480), - "1.33": (554, 416), - "1.50": (588, 392), - "1.78": (640, 360), - "1.89": (660, 350), - "2.00": (678, 340), - "2.08": (692, 332), -} - -# S = 102240 -ASPECT_RATIO_240P = { - "0.38": (196, 522), - "0.43": (210, 490), - "0.48": (222, 462), - "0.50": (226, 452), - "0.53": (232, 438), - "0.54": (236, 436), - "0.56": (240, 426), # base - "0.62": (252, 404), - "0.67": (262, 393), - "0.75": (276, 368), - "1.00": (320, 320), - "1.33": (370, 278), - "1.50": (392, 262), - "1.78": (426, 240), - "1.89": (440, 232), - "2.00": (452, 226), - "2.08": (462, 222), -} - -# S = 36864 -ASPECT_RATIO_144P = { - "0.38": (117, 312), - "0.43": (125, 291), - "0.48": (133, 277), - "0.50": (135, 270), - "0.53": (139, 262), - "0.54": (141, 260), - "0.56": (144, 256), # base - "0.62": (151, 241), - "0.67": (156, 234), - "0.75": (166, 221), - "1.00": (192, 192), - "1.33": (221, 165), - "1.50": (235, 156), - "1.78": (256, 144), - "1.89": (263, 139), - "2.00": (271, 135), - "2.08": (277, 132), -} - -# from PixArt -# S = 8294400 -ASPECT_RATIO_2880 = { - "0.25": (1408, 5760), - "0.26": (1408, 5568), - "0.27": (1408, 5376), - "0.28": (1408, 5184), - "0.32": (1600, 4992), - "0.33": (1600, 4800), - "0.34": (1600, 4672), - "0.40": (1792, 4480), - "0.42": (1792, 4288), - "0.47": (1920, 4096), - "0.49": (1920, 3904), - "0.51": (1920, 3776), - "0.55": (2112, 3840), - "0.59": (2112, 3584), - "0.68": (2304, 3392), - "0.72": (2304, 3200), - "0.78": (2496, 3200), - "0.83": (2496, 3008), - "0.89": (2688, 3008), - "0.93": (2688, 2880), - "1.00": (2880, 2880), - "1.07": (2880, 2688), - "1.12": (3008, 2688), - "1.21": (3008, 2496), - "1.28": (3200, 2496), - "1.39": (3200, 2304), - "1.47": (3392, 2304), - "1.70": (3584, 2112), - "1.82": (3840, 2112), - "2.03": (3904, 1920), - "2.13": (4096, 1920), - "2.39": (4288, 1792), - "2.50": (4480, 1792), - "2.92": (4672, 1600), - "3.00": (4800, 1600), - "3.12": (4992, 1600), - "3.68": (5184, 1408), - "3.82": (5376, 1408), - "3.95": (5568, 1408), - "4.00": (5760, 1408), -} - -# S = 4194304 -ASPECT_RATIO_2048 = { - "0.25": (1024, 4096), - "0.26": (1024, 3968), - "0.27": (1024, 3840), - "0.28": (1024, 3712), - "0.32": (1152, 3584), - "0.33": (1152, 3456), - "0.35": (1152, 3328), - "0.40": (1280, 3200), - "0.42": (1280, 3072), - "0.48": (1408, 2944), - "0.50": (1408, 2816), - "0.52": (1408, 2688), - "0.57": (1536, 2688), - "0.60": (1536, 2560), - "0.68": (1664, 2432), - "0.72": (1664, 2304), - "0.78": (1792, 2304), - "0.82": (1792, 2176), - "0.88": (1920, 2176), - "0.94": (1920, 2048), - "1.00": (2048, 2048), - "1.07": (2048, 1920), - "1.13": (2176, 1920), - "1.21": (2176, 1792), - "1.29": (2304, 1792), - "1.38": (2304, 1664), - "1.46": (2432, 1664), - "1.67": (2560, 1536), - "1.75": (2688, 1536), - "2.00": (2816, 1408), - "2.09": (2944, 1408), - "2.40": (3072, 1280), - "2.50": (3200, 1280), - "2.89": (3328, 1152), - "3.00": (3456, 1152), - "3.11": (3584, 1152), - "3.62": (3712, 1024), - "3.75": (3840, 1024), - "3.88": (3968, 1024), - "4.00": (4096, 1024), -} - -# S = 1048576 -ASPECT_RATIO_1024 = { - "0.25": (512, 2048), - "0.26": (512, 1984), - "0.27": (512, 1920), - "0.28": (512, 1856), - "0.32": (576, 1792), - "0.33": (576, 1728), - "0.35": (576, 1664), - "0.40": (640, 1600), - "0.42": (640, 1536), - "0.48": (704, 1472), - "0.50": (704, 1408), - "0.52": (704, 1344), - "0.57": (768, 1344), - "0.60": (768, 1280), - "0.68": (832, 1216), - "0.72": (832, 1152), - "0.78": (896, 1152), - "0.82": (896, 1088), - "0.88": (960, 1088), - "0.94": (960, 1024), - "1.00": (1024, 1024), - "1.07": (1024, 960), - "1.13": (1088, 960), - "1.21": (1088, 896), - "1.29": (1152, 896), - "1.38": (1152, 832), - "1.46": (1216, 832), - "1.67": (1280, 768), - "1.75": (1344, 768), - "2.00": (1408, 704), - "2.09": (1472, 704), - "2.40": (1536, 640), - "2.50": (1600, 640), - "2.89": (1664, 576), - "3.00": (1728, 576), - "3.11": (1792, 576), - "3.62": (1856, 512), - "3.75": (1920, 512), - "3.88": (1984, 512), - "4.00": (2048, 512), -} - -# S = 262144 -ASPECT_RATIO_512 = { - "0.25": (256, 1024), - "0.26": (256, 992), - "0.27": (256, 960), - "0.28": (256, 928), - "0.32": (288, 896), - "0.33": (288, 864), - "0.35": (288, 832), - "0.40": (320, 800), - "0.42": (320, 768), - "0.48": (352, 736), - "0.50": (352, 704), - "0.52": (352, 672), - "0.57": (384, 672), - "0.60": (384, 640), - "0.68": (416, 608), - "0.72": (416, 576), - "0.78": (448, 576), - "0.82": (448, 544), - "0.88": (480, 544), - "0.94": (480, 512), - "1.00": (512, 512), - "1.07": (512, 480), - "1.13": (544, 480), - "1.21": (544, 448), - "1.29": (576, 448), - "1.38": (576, 416), - "1.46": (608, 416), - "1.67": (640, 384), - "1.75": (672, 384), - "2.00": (704, 352), - "2.09": (736, 352), - "2.40": (768, 320), - "2.50": (800, 320), - "2.89": (832, 288), - "3.00": (864, 288), - "3.11": (896, 288), - "3.62": (928, 256), - "3.75": (960, 256), - "3.88": (992, 256), - "4.00": (1024, 256), -} - -# S = 65536 -ASPECT_RATIO_256 = { - "0.25": (128, 512), - "0.26": (128, 496), - "0.27": (128, 480), - "0.28": (128, 464), - "0.32": (144, 448), - "0.33": (144, 432), - "0.35": (144, 416), - "0.40": (160, 400), - "0.42": (160, 384), - "0.48": (176, 368), - "0.50": (176, 352), - "0.52": (176, 336), - "0.57": (192, 336), - "0.60": (192, 320), - "0.68": (208, 304), - "0.72": (208, 288), - "0.78": (224, 288), - "0.82": (224, 272), - "0.88": (240, 272), - "0.94": (240, 256), - "1.00": (256, 256), - "1.07": (256, 240), - "1.13": (272, 240), - "1.21": (272, 224), - "1.29": (288, 224), - "1.38": (288, 208), - "1.46": (304, 208), - "1.67": (320, 192), - "1.75": (336, 192), - "2.00": (352, 176), - "2.09": (368, 176), - "2.40": (384, 160), - "2.50": (400, 160), - "2.89": (416, 144), - "3.00": (432, 144), - "3.11": (448, 144), - "3.62": (464, 128), - "3.75": (480, 128), - "3.88": (496, 128), - "4.00": (512, 128), -} - - -def get_closest_ratio(height: float, width: float, ratios: dict): - aspect_ratio = height / width - closest_ratio = min(ratios.keys(), - key=lambda ratio: abs(float(ratio) - aspect_ratio)) - return closest_ratio - - -ASPECT_RATIOS = { - "144p": (36864, ASPECT_RATIO_144P), - "256": (65536, ASPECT_RATIO_256), - "240p": (102240, ASPECT_RATIO_240P), - "360p": (230400, ASPECT_RATIO_360P), - "512": (262144, ASPECT_RATIO_512), - "480p": (409920, ASPECT_RATIO_480P), - "720p": (921600, ASPECT_RATIO_720P), - "1024": (1048576, ASPECT_RATIO_1024), - "1080p": (2073600, ASPECT_RATIO_1080P), - "2k": (3686400, ASPECT_RATIO_2K), - "2048": (4194304, ASPECT_RATIO_2048), - "2880": (8294400, ASPECT_RATIO_2880), - "4k": (8294400, ASPECT_RATIO_4K), -} - - -def get_num_pixels(name): - return ASPECT_RATIOS[name][0] - - -def get_image_size(resolution, ar_ratio): - if ar_ratio in ASPECT_RATIO_MAP: - ar_key = ASPECT_RATIO_MAP[ar_ratio] - else: - ar_key = ar_ratio - rs_dict = ASPECT_RATIOS[resolution][1] - assert ar_key in rs_dict, f"Aspect ratio {ar_ratio} not found for resolution {resolution}" - return rs_dict[ar_key] - - -NUM_FRAMES_MAP = { - "1x": 51, - "2x": 102, - "4x": 204, - "8x": 408, - "16x": 816, - "2s": 51, - "4s": 102, - "8s": 204, - "16s": 408, - "32s": 816, -} - - -def get_num_frames(num_frames): - if num_frames in NUM_FRAMES_MAP: - return NUM_FRAMES_MAP[num_frames] - else: - return int(num_frames) diff --git a/examples/models/contrib/stdit/assets/a_beautiful_waterfall.mp4 b/examples/models/contrib/stdit/assets/a_beautiful_waterfall.mp4 deleted file mode 100644 index f423c02ad73e..000000000000 Binary files a/examples/models/contrib/stdit/assets/a_beautiful_waterfall.mp4 and /dev/null differ diff --git a/examples/models/contrib/stdit/convert_checkpoint.py b/examples/models/contrib/stdit/convert_checkpoint.py deleted file mode 100644 index 242dd4a9a141..000000000000 --- a/examples/models/contrib/stdit/convert_checkpoint.py +++ /dev/null @@ -1,252 +0,0 @@ -import argparse -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -from vae import get_vae - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import STDiT3Model - -PRETRAINED_STDIT_PATH = "hpcai-tech/OpenSora-STDiT-v3" - - -def pixel_size_to_latent_size(args): - vae = get_vae( - from_pretrained=args.vae_type, - micro_frame_size=args.vae_micro_frame_size, - micro_batch_size=args.vae_micro_batch_size, - ).eval() - spatial_patch_size = vae.spatial_vae.patch_size - temporal_patch_size = vae.temporal_vae.patch_size - vae_out_channels = vae.out_channels - pixel_size = (args.num_frames, args.height, args.width) - latent_size = vae.get_latent_size(pixel_size) - return { - 'in_channels': vae_out_channels, - 'latent_size': latent_size, - 'spatial_patch_size': spatial_patch_size, - 'temporal_patch_size': temporal_patch_size, - } - - -def size_str_to_list(repr): - return [int(it) for it in repr.split('x')] if 'x' in repr else [int(repr)] - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--timm_ckpt', type=str, default=None) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument('--caption_channels', - type=int, - default=4096, - help='The channel of input of caption embedder') - parser.add_argument('--depth', - type=int, - default=28, - help='The number of STDiT blocks') - parser.add_argument('--input_sq_size', - type=int, - default=512, - help='Base spatial position embedding size') - parser.add_argument('--stdit_type', - type=str, - default="STDiT3", - choices=["STDiT3"]) - parser.add_argument('--stdit_patch_size', - type=str, - default='1x2x2', - help='The patch size of stdit for patchify') - parser.add_argument('--width', - type=int, - default=640, - help='The width of image size') - parser.add_argument('--height', - type=int, - default=360, - help='The height of image size') - parser.add_argument('--num_frames', - type=int, - default=102, - help='The frames of generated video') - parser.add_argument('--vae_type', - type=str, - default="hpcai-tech/OpenSora-VAE-v1.2", - choices=["hpcai-tech/OpenSora-VAE-v1.2"]) - parser.add_argument('--vae_micro_frame_size', - type=int, - default=17, - help='The micro_frame_size for vae') - parser.add_argument('--vae_micro_batch_size', - type=int, - default=4, - help='The micro_batch_size for vae') - parser.add_argument('--hidden_size', - type=int, - default=1152, - help='The hidden size of STDiT') - parser.add_argument('--num_heads', - type=int, - default=16, - help='The number of heads of attention module') - parser.add_argument( - '--mlp_ratio', - type=float, - default=4.0, - help= - 'The ratio of hidden size compared to input hidden size in MLP layer') - parser.add_argument( - '--class_dropout_prob', - type=float, - default=0.1, - help='The probability to drop class token when training') - parser.add_argument('--model_max_length', - type=int, - default=300, - help='The max number of tokens (default: 300)') - parser.add_argument('--text_encoder_type', - type=str, - default="DeepFloyd/t5-v1_1-xxl", - choices=["DeepFloyd/t5-v1_1-xxl"]) - parser.add_argument('--learn_sigma', - type=bool, - default=True, - help='Whether the model learn sigma') - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--cp_size', - type=int, - default=1, - help='Context parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument('--disable_qk_norm', - action='store_true', - help='Disable norm for qk in attention') - parser.add_argument('--fp8', - action='store_true', - help='Whether use FP8 for layers') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument('--log_level', type=str, default='info') - args = parser.parse_args() - return args - - -def convert_and_save_model(args): - # [NOTE] PP is not supported yet. - world_size = args.tp_size * args.cp_size - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - def convert_and_save_rank(args, rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - cp_size=args.cp_size) - # Process args - runtime_config = { - 'architecture': "STDiT3", - 'checkpoint_path': os.path.abspath(args.timm_ckpt), - 'caption_channels': args.caption_channels, - 'num_hidden_layers': args.depth, - 'width': args.width, - 'height': args.height, - 'num_frames': args.num_frames, - 'hidden_size': args.hidden_size, - 'stdit_patch_size': size_str_to_list(args.stdit_patch_size), - 'input_sq_size': args.input_sq_size, - 'num_attention_heads': args.num_heads, - 'model_max_length': args.model_max_length, - 'mlp_ratio': args.mlp_ratio, - 'class_dropout_prob': args.class_dropout_prob, - 'learn_sigma': args.learn_sigma, - 'qk_norm': (not args.disable_qk_norm), - 'stdit_type': args.stdit_type, - 'vae_type': args.vae_type, - 'text_encoder_type': args.text_encoder_type, - } - runtime_config.update(pixel_size_to_latent_size(args)) - tik = time.time() - stdit = STDiT3Model.from_pretrained(os.path.dirname(args.timm_ckpt), - args.dtype, - mapping=mapping, - **runtime_config) - stdit.save_checkpoint(args.output_dir, save_config=True) - print(f'Total time of reading and converting: {time.time()-tik:.3f} s') - tik = time.time() - del stdit - print(f'Total time of saving checkpoint: {time.time()-tik:.3f} s') - - execute(args.workers, [convert_and_save_rank] * world_size, args) - release_gc() - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - logger.set_level(args.log_level) - - assert args.pp_size == 1, "PP is not supported yet." - - tik = time.time() - - if args.timm_ckpt is None: - print( - f"No pretrained checkpoint provided, use default checkpoint from Huggingface instead." - ) - args.timm_ckpt = "./pretrained_ckpt/model.safetensors" - if not os.path.exists(args.timm_ckpt): - from huggingface_hub import snapshot_download - snapshot_download(repo_id=PRETRAINED_STDIT_PATH, - local_dir=os.path.dirname(args.timm_ckpt)) - - convert_and_save_model(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/stdit/pipeline_tllm.py b/examples/models/contrib/stdit/pipeline_tllm.py deleted file mode 100644 index 2e3ff8e346a3..000000000000 --- a/examples/models/contrib/stdit/pipeline_tllm.py +++ /dev/null @@ -1,607 +0,0 @@ -import os -import time -from functools import wraps -from typing import List - -import numpy as np -import torch -import torch.distributed -from cuda import cudart -from scheduler import timestep_transform -from utils import DataProcessor, print_progress_bar - -import tensorrt_llm -from tensorrt_llm._utils import (str_dtype_to_torch, str_dtype_to_trt, - torch_dtype_to_trt, trt_dtype_to_torch) -from tensorrt_llm.logger import logger -from tensorrt_llm.plugin.plugin import CustomAllReduceHelper -from tensorrt_llm.runtime.session import Session, TensorInfo - - -def CUASSERT(cuda_ret): - err = cuda_ret[0] - if err != cudart.cudaError_t.cudaSuccess: - raise RuntimeError( - f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t" - ) - if len(cuda_ret) > 1: - return cuda_ret[1:] - return None - - -class TllmSTDiT(): - - def __init__(self, - config, - tllm_model_dir, - debug_mode=True, - stream: torch.cuda.Stream = None): - self.dtype = config['pretrained_config']['dtype'] - - self.depth = config['pretrained_config']['num_hidden_layers'] - - rank = tensorrt_llm.mpi_rank() - world_size = config['pretrained_config']['mapping']['world_size'] - cp_size = config['pretrained_config']['mapping']['cp_size'] - tp_size = config['pretrained_config']['mapping']['tp_size'] - pp_size = config['pretrained_config']['mapping']['pp_size'] - gpus_per_node = config['pretrained_config']['mapping']['gpus_per_node'] - assert pp_size == 1 - self.mapping = tensorrt_llm.Mapping(world_size=world_size, - rank=rank, - cp_size=cp_size, - tp_size=tp_size, - pp_size=1, - gpus_per_node=gpus_per_node) - - local_rank = rank % self.mapping.gpus_per_node - self.device = torch.device(f'cuda:{local_rank}') - torch.cuda.set_device(self.device) - CUASSERT(cudart.cudaSetDevice(local_rank)) - - self.stream = stream - if self.stream is None: - self.stream = torch.cuda.Stream(self.device) - torch.cuda.set_stream(self.stream) - - engine_file = os.path.join(tllm_model_dir, f"rank{rank}.engine") - logger.info(f'Loading engine from {engine_file}') - with open(engine_file, "rb") as f: - engine_buffer = f.read() - - assert engine_buffer is not None - - self.session = Session.from_serialized_engine(engine_buffer) - - self.debug_mode = debug_mode - - self.inputs = {} - self.outputs = {} - self.buffer_allocated = False - - expected_tensor_names = [ - 'x', 'timestep', 'y', 'mask', 'x_mask', 'fps', 'height', 'width', - 'output' - ] - - if self.mapping.tp_size > 1: - self.buffer, self.all_reduce_workspace = CustomAllReduceHelper.allocate_workspace( - self.mapping, - CustomAllReduceHelper.max_workspace_size_auto( - self.mapping.tp_size)) - self.inputs['all_reduce_workspace'] = self.all_reduce_workspace - expected_tensor_names += ['all_reduce_workspace'] - - self.latent_size = config['pretrained_config']['latent_size'] - self.patch_size = config['pretrained_config']['stdit_patch_size'] - self.in_channels = config['pretrained_config']['in_channels'] - self.caption_channels = config['pretrained_config']['caption_channels'] - self.model_max_length = config['pretrained_config']['model_max_length'] - self.config = config['pretrained_config'] - - self.max_cattn_seq_len = int( - np.prod([ - np.ceil(d / p) - for d, p in zip(self.latent_size, self.patch_size) - ])) - - def _tensor_dtype(self, name): - # return torch dtype given tensor name for convenience - dtype = trt_dtype_to_torch(self.session.engine.get_tensor_dtype(name)) - return dtype - - def _get_extra_inputs_for_attention_plugin(self, batch_size, max_seq_len, - max_encoder_seq_len): - host_max_attention_window_sizes = torch.tensor([max_seq_len] * - self.depth, - dtype=torch.int32).cpu() - host_sink_token_length = torch.tensor([0], dtype=torch.int32).cpu() - context_lengths = torch.full([batch_size], - max_seq_len, - dtype=torch.int32).cuda() - host_context_lengths = torch.full([batch_size], - max_seq_len, - dtype=torch.int32).cpu() - host_request_types = torch.zeros([batch_size], dtype=torch.int32).cpu() - perf_knob_tensor_size = 16 - host_runtime_perf_knobs = torch.tensor([-1] * perf_knob_tensor_size, - dtype=torch.int64, - device='cpu') - host_context_progress = torch.tensor([0], - dtype=torch.int64, - device='cpu') - cross_encoder_input_lengths = torch.full([batch_size], - max_encoder_seq_len, - dtype=torch.int32, - device='cuda') - cross_encoder_max_input_length = torch.empty((max_encoder_seq_len, ), - dtype=torch.int32, - device='cuda') - - extra_inputs = { - 'host_max_attention_window_sizes': host_max_attention_window_sizes, - 'host_sink_token_length': host_sink_token_length, - 'context_lengths': context_lengths, - 'host_context_lengths': host_context_lengths, - 'encoder_input_lengths': cross_encoder_input_lengths, - 'encoder_max_input_length': cross_encoder_max_input_length, - 'host_request_types': host_request_types, - 'host_runtime_perf_knobs': host_runtime_perf_knobs, - 'host_context_progress': host_context_progress - } - return extra_inputs - - def _setup(self, batch_size, max_encoder_seq_len): - input_info = [ - TensorInfo(name='x', - dtype=str_dtype_to_trt(self.dtype), - shape=(batch_size, self.in_channels, *self.latent_size)), - TensorInfo(name='timestep', - dtype=str_dtype_to_trt(self.dtype), - shape=(batch_size, )), - TensorInfo(name='y', - dtype=str_dtype_to_trt(self.dtype), - shape=(batch_size, 1, self.model_max_length, - self.caption_channels)), - TensorInfo(name='mask', - dtype=str_dtype_to_trt('int32'), - shape=(1, self.model_max_length)), - TensorInfo(name='x_mask', - dtype=str_dtype_to_trt('bool'), - shape=(batch_size, self.latent_size[0])), - TensorInfo(name='fps', - dtype=str_dtype_to_trt(self.dtype), - shape=(1, )), - TensorInfo(name='height', - dtype=str_dtype_to_trt(self.dtype), - shape=(1, )), - TensorInfo(name='width', - dtype=str_dtype_to_trt(self.dtype), - shape=(1, )), - ] - - extra_inputs = self._get_extra_inputs_for_attention_plugin( - batch_size=batch_size, - max_seq_len=self.max_cattn_seq_len, - max_encoder_seq_len=max_encoder_seq_len, - ) - input_info += [ - tensorrt_llm.runtime.TensorInfo(name, - torch_dtype_to_trt(tensor.dtype), - tensor.shape) - for name, tensor in extra_inputs.items() - ] - - output_info = self.session.infer_shapes(input_info) - for t_info in output_info: - self.outputs[t_info.name] = torch.empty(tuple(t_info.shape), - dtype=trt_dtype_to_torch( - t_info.dtype), - device=self.device) - self.buffer_allocated = True - - def cuda_stream_guard(func): - """Sync external stream and set current stream to the one bound to the session. Reset on exit. - """ - - @wraps(func) - def wrapper(self, *args, **kwargs): - external_stream = torch.cuda.current_stream() - if external_stream != self.stream: - external_stream.synchronize() - torch.cuda.set_stream(self.stream) - ret = func(self, *args, **kwargs) - if external_stream != self.stream: - self.stream.synchronize() - torch.cuda.set_stream(external_stream) - return ret - - return wrapper - - @cuda_stream_guard - def forward(self, x: torch.Tensor, timestep: torch.Tensor, y: torch.Tensor, - mask: torch.Tensor, x_mask: torch.Tensor, fps: torch.Tensor, - height: torch.Tensor, width: torch.Tensor, y_lens: List[int]): - """ - Forward pass of STDiT. - x: (N, C, F, H, W) - timestep: (N,) - y: () - mask: () - x_mask: (N * 2, ) - fps: (1) - height: (1) - width: (1) - """ - self._setup(x.shape[0], max_encoder_seq_len=np.max(y_lens)) - if not self.buffer_allocated: - raise RuntimeError('Buffer not allocated, please call setup first!') - - inputs = { - 'x': x, - "timestep": timestep, - 'y': y, - "mask": mask, - "x_mask": x_mask, - "fps": fps, - "height": height, - "width": width, - } - self.inputs.update(**inputs) - extra_inputs = self._get_extra_inputs_for_attention_plugin( - batch_size=x.shape[0], - max_seq_len=self.max_cattn_seq_len, - max_encoder_seq_len=np.max(y_lens), - ) - self.inputs.update(**extra_inputs) - - self.session.set_shapes(self.inputs) - ok = self.session.run(self.inputs, self.outputs, - self.stream.cuda_stream) - - if not ok: - raise RuntimeError('Executing TRT engine failed!') - debug_tensors = {} - for name in list(self.outputs.keys()): - if name != 'output': - debug_tensors[name] = self.outputs.pop(name) - if len(debug_tensors) == 0: - return self.outputs['output'] - else: - return self.outputs['output'], debug_tensors - - -class TllmOpenSoraPipeline(): - - def __init__( - self, - stdit, - text_encoder, - vae, - scheduler, - num_sampling_steps=30, - num_timesteps=1000, - cfg_scale=4.0, - align=None, - aes=None, - flow=None, - camera_motion=None, - image_size=None, - resolution=None, - aspect_ratio=None, - num_frames=None, - fps=None, - save_fps=None, - diffusion_model_type=None, - condition_frame_length=None, - condition_frame_edit=None, - use_discrete_timesteps=False, - use_timestep_transform=False, - video_save_dir='./samples/', - dtype='float16', - device=torch.device('cuda'), - seed=None, - **kwargs, - ): - self.stdit = stdit - self.text_encoder = text_encoder - self.vae = vae - self.scheduler = scheduler - self.data_processor = DataProcessor(text_encoder=text_encoder, vae=vae) - - self.num_sampling_steps = num_sampling_steps - self.num_timesteps = num_timesteps - self.cfg_scale = cfg_scale - self.align = align - self.aes = aes - self.flow = flow - self.camera_motion = camera_motion - - if image_size is None: - assert ( - resolution is not None and aspect_ratio is not None - ), "resolution and aspect_ratio must be provided if image_size is not provided" - image_size = self.data_processor.get_image_size( - resolution, aspect_ratio) - self.image_size = image_size - self.num_frames = self.data_processor.get_num_frames(num_frames) - self.input_size = (num_frames, *image_size) - self.latent_size = self.data_processor.get_latent_size(self.input_size) - - self.fps = fps - self.save_fps = save_fps - self.diffusion_model_type = diffusion_model_type - self.condition_frame_length = condition_frame_length - self.condition_frame_edit = condition_frame_edit - - self.use_discrete_timesteps = use_discrete_timesteps - self.use_timestep_transform = use_timestep_transform - - self.video_save_dir = video_save_dir - self.seed = self.data_processor.set_random_seed(seed) - self.dtype = dtype - self.device = device - - self.prompt_as_path = kwargs.get("prompt_as_path", False) - self.watermark = kwargs.get("watermark", False) - - def __call__( - self, - prompts, - batch_size=1, - num_sample=1, - loop=1, - reference_path=None, - mask_strategy=None, - sample_name=None, - start_idx=0, - ): - reference_path = [""] * len( - prompts) if reference_path is None else reference_path - mask_strategy = [""] * len( - prompts) if mask_strategy is None else mask_strategy - assert len(reference_path) == len( - prompts), "Length of reference must be the same as prompts" - assert len(mask_strategy) == len( - prompts), "Length of mask_strategy must be the same as prompts" - - for i in range(0, len(prompts), batch_size): - # == prepare batch prompts == - batch_prompts = prompts[i:i + batch_size] - ms = mask_strategy[i:i + batch_size] - refs = reference_path[i:i + batch_size] - - # == get json from prompts == - batch_prompts, refs, ms = self.data_processor.extract_json_from_prompts( - batch_prompts, refs, ms) - original_batch_prompts = batch_prompts - - # == get reference for condition == - refs = self.data_processor.collect_references_batch( - refs, self.image_size) - - # == multi-resolution info == - model_args = self.data_processor.prepare_multi_resolution_info( - self.diffusion_model_type, len(batch_prompts), self.image_size, - self.num_frames, self.fps, self.device, - str_dtype_to_torch(self.dtype)) - - # == Iter over number of sampling for one prompt == - for k in range(num_sample): - # == prepare save paths == - save_paths = [ - self.data_processor.get_save_path_name( - self.video_save_dir, - sample_name=sample_name, - sample_idx=start_idx + idx, - prompt=original_batch_prompts[idx], - prompt_as_path=self.prompt_as_path, - num_sample=num_sample, - k=k, - ) for idx in range(len(batch_prompts)) - ] - - # NOTE: Skip if the sample already exists - # This is useful for resuming sampling VBench - if self.prompt_as_path and \ - all(os.path.exists(path) for path in save_paths): - continue - - # == process prompts step by step == - # 0. split prompt - # each element in the list is [prompt_segment_list, loop_idx_list] - batched_prompt_segment_list = [] - batched_loop_idx_list = [] - for prompt in batch_prompts: - prompt_segment_list, loop_idx_list = self.data_processor.split_prompt( - prompt) - batched_prompt_segment_list.append(prompt_segment_list) - batched_loop_idx_list.append(loop_idx_list) - - # [NOTE] Skip refining prompt by OpenAI - for idx, prompt_segment_list in enumerate( - batched_prompt_segment_list): - batched_prompt_segment_list[ - idx] = self.data_processor.append_score_to_prompts( - prompt_segment_list, - aes=self.aes, - flow=self.flow, - camera_motion=self.camera_motion, - ) - - # 3. clean prompt with T5 - for idx, prompt_segment_list in enumerate( - batched_prompt_segment_list): - batched_prompt_segment_list[idx] = [ - self.data_processor.text_preprocessing(prompt) - for prompt in prompt_segment_list - ] - - # 4. merge to obtain the final prompt - batch_prompts = [] - for prompt_segment_list, loop_idx_list in zip( - batched_prompt_segment_list, batched_loop_idx_list): - batch_prompts.append( - self.data_processor.merge_prompt( - prompt_segment_list, loop_idx_list)) - - # == Iter over loop generation == - video_clips = [] - for loop_i in range(loop): - # == get prompt for loop i == - batch_prompts_loop = self.data_processor.extract_prompts_loop( - batch_prompts, loop_i) - - # == add condition frames for loop == - if loop_i > 0: - refs, ms = self.data_processor.append_generated( - video_clips[-1], refs, ms, loop_i, - self.condition_frame_length, - self.condition_frame_edit) - - # == sampling == - torch.manual_seed(self.seed) - noise = torch.randn(len(batch_prompts), - self.vae.out_channels, - *self.latent_size, - device=self.device, - dtype=str_dtype_to_torch(self.dtype)) - masks = self.data_processor.apply_mask_strategy( - noise, refs, ms, loop_i, align=self.align) - samples = self.sample( - latent=noise, - prompts=batch_prompts_loop, - mask=masks, - additional_args=model_args, - ) - samples = self.vae.decode(samples.to( - str_dtype_to_torch(self.dtype)), - num_frames=self.num_frames) - video_clips.append(samples) - - self.save_video(video_clips, save_paths, len(batch_prompts), - loop) - start_idx += len(batch_prompts) - - def sample( - self, - latent, - prompts, - mask=None, - additional_args=None, - guidance_scale=None, - ): - # if no specific guidance scale is provided, use the default scale when initializing the scheduler - if guidance_scale is None: - guidance_scale = self.cfg_scale - - # text encoding - model_args = self.text_encoder.encode_with_null(prompts) - if additional_args is not None: - model_args.update(additional_args) - - # prepare timesteps - timesteps = [(1.0 - i / self.num_sampling_steps) * self.num_timesteps - for i in range(self.num_sampling_steps)] - if self.use_discrete_timesteps: - timesteps = [int(round(t)) for t in timesteps] - timesteps = [ - torch.tensor([t] * latent.shape[0], device=self.device) - for t in timesteps - ] - if self.use_timestep_transform: - timesteps = [ - timestep_transform(t, - additional_args, - num_timesteps=self.num_timesteps) - for t in timesteps - ] - - if mask is not None: - noise_added = torch.zeros_like(mask, dtype=torch.bool) - noise_added = noise_added | (mask == 1) - - for i, t in enumerate(timesteps): - print_progress_bar(i, self.num_sampling_steps) - # mask for adding noise - if mask is not None: - mask_t = mask * self.num_timesteps - x0 = latent.clone() - x_noise = self.scheduler.add_noise(x0, torch.randn_like(x0), t) - - mask_t_upper = mask_t >= t.unsqueeze(1) - model_args["x_mask"] = mask_t_upper.repeat(2, 1) - mask_add_noise = mask_t_upper & ~noise_added - - latent = torch.where(mask_add_noise[:, None, :, None, None], - x_noise, x0) - noise_added = mask_t_upper - - # classifier-free guidance - latent_in = torch.cat([latent, latent], 0) - t = torch.cat([t, t], 0) - for k in list(model_args.keys()): - if k not in ['y', 'mask', 'x_mask', 'fps', 'height', 'width']: - model_args.pop(k) - else: - if model_args[k].dtype == torch.float32: - model_args[k] = model_args[k].to(torch.float16) - latent_in = latent_in.to(torch.float16) - t = t.to(torch.float16) - model_args['mask'] = model_args['mask'].to(torch.int32) - model_args['y_lens'] = self._get_y_lens(model_args['y'], - model_args['mask']) - - pred = self.stdit.forward(x=latent_in, timestep=t, - **model_args).chunk(2, dim=1)[0] - pred_cond, pred_uncond = pred.chunk(2, dim=0) - v_pred = pred_uncond + guidance_scale * (pred_cond - pred_uncond) - - # update latent - dt = timesteps[i] - timesteps[i + 1] if i < len( - timesteps) - 1 else timesteps[i] - dt = dt / self.num_timesteps - latent = latent + v_pred * dt[:, None, None, None, None] - - if mask is not None: - latent = torch.where(mask_t_upper[:, None, :, None, None], - latent, x0) - print() - return latent - - def _get_y_lens(self, y, mask=None): - assert (len(y.shape) == 4) - if mask is not None: - if mask.shape[0] != y.shape[0]: - mask = mask.repeat(y.shape[0] // mask.shape[0], 1) - mask = mask.squeeze(1).squeeze(1) - y_lens = mask.sum(dim=1).tolist() - else: - y_lens = [y.shape[2]] * y.shape[0] - return y_lens - - def save_video( - self, - video_clips, - save_paths, - batch_size, - loop, - ): - os.makedirs(self.video_save_dir, exist_ok=True) - if tensorrt_llm.mpi_rank() == 0: - for idx in range(batch_size): - save_path = save_paths[idx] - video = [video_clips[i][idx] for i in range(loop)] - for i in range(1, loop): - video[i] = video[i][:, - self.data_processor.dframe_to_frame( - self.condition_frame_length):] - video = torch.cat(video, dim=1) - save_path = self.data_processor.save_sample( - video, - fps=self.save_fps, - save_path=save_path, - ) - if save_path.endswith(".mp4") and self.watermark: - time.sleep(1) # prevent loading previous generated video - self.data_processor.add_watermark(save_path) diff --git a/examples/models/contrib/stdit/requirements.txt b/examples/models/contrib/stdit/requirements.txt deleted file mode 100644 index bc9a3f3a24c8..000000000000 --- a/examples/models/contrib/stdit/requirements.txt +++ /dev/null @@ -1,13 +0,0 @@ -ftfy -av==12.3.0 -git+https://github.com/facebookresearch/xformers.git@v0.0.29#egg=xformers -rotary_embedding_torch==0.5.3 -git+https://github.com/hpcaitech/TensorNVMe.git -# For colossalai -fabric -contexttimer -ray -protobuf -bitsandbytes>=0.39.0 -rpyc==6.0.0 -galore_torch diff --git a/examples/models/contrib/stdit/sample.py b/examples/models/contrib/stdit/sample.py deleted file mode 100644 index 88890d8ee1d5..000000000000 --- a/examples/models/contrib/stdit/sample.py +++ /dev/null @@ -1,170 +0,0 @@ -import argparse -import glob -import json -import os - -import torch -from pipeline_tllm import TllmOpenSoraPipeline, TllmSTDiT -from safetensors.torch import load_file -from scheduler import RFlowScheduler -from text_encoder import CaptionEmbedder, T5Encoder -from utils import DataProcessor -from vae import get_vae - -import tensorrt_llm -from tensorrt_llm._utils import str_dtype_to_torch - - -class RuntimeConfig(dict): - - def __getattr__(self, key): - if key in self: - return self[key] - raise AttributeError( - f"'RuntimeConfig {self}' object has no attribute '{key}'") - - def __setattr__(self, key, value): - self[key] = value - - -def main(cfg): - tensorrt_llm.logger.set_level(cfg.log_level) - - torch.set_grad_enabled(False) - dtype = cfg.get("dtype", "float16") - device = torch.device('cuda') - - config_file = os.path.join(cfg.tllm_model_dir, 'config.json') - with open(config_file) as f: - config = json.load(f) - ## Build modules - model_config = config.get("pretrained_config") - stdit = TllmSTDiT(config, cfg.tllm_model_dir, debug_mode=cfg.debug_mode) - # HACK: for classifier-free guidance - ckpt_path = model_config['checkpoint_path'] - if os.path.isdir(ckpt_path): - ckpt_files = glob.glob(model_config['checkpoint_path'] + "/*") - else: - ckpt_files = [ckpt_path] - pretrained_weights = None - for file in ckpt_files: - if file.endswith('.safetensors'): - pretrained_weights = load_file(file) - elif file.endswith('.pt'): - pretrained_weights = dict(torch.load(file, weights_only=True)) - if pretrained_weights is not None: - break - if pretrained_weights is None: - raise FileNotFoundError - y_embedder = CaptionEmbedder( - in_channels=model_config.get('caption_channels'), - hidden_size=model_config.get('hidden_size'), - uncond_prob=model_config.get('class_dropout_prob'), - act_layer=torch.nn.GELU(approximate="tanh"), - token_num=model_config.get('model_max_length'), - ) - for name, param in y_embedder.named_parameters(): - param.data = pretrained_weights['y_embedder.' + name] - y_embedder.y_embedding = pretrained_weights['y_embedder.y_embedding'] - reuse_y_embedding = pretrained_weights['y_embedder.y_embedding'] - text_encoder = T5Encoder( - from_pretrained=cfg.text_encoder, - model_max_length=model_config.get('model_max_length'), - caption_channels=model_config.get('caption_channels'), - y_embedding=reuse_y_embedding, # HACK: for classifier-free guidance - hidden_size=model_config.get('hidden_size'), - device=device, - ) - vae = get_vae( - from_pretrained=cfg.vae, - micro_frame_size=17, - micro_batch_size=4, - ).to(dtype=str_dtype_to_torch(dtype), device=device).eval() - scheduler = RFlowScheduler( - use_timestep_transform=True, - num_timesteps=cfg.num_timesteps, - num_sampling_steps=cfg.num_sampling_steps, - ) - pipe = TllmOpenSoraPipeline( - stdit=stdit, - text_encoder=text_encoder, - vae=vae, - scheduler=scheduler, - num_sampling_steps=cfg.num_sampling_steps, - num_timesteps=cfg.num_timesteps, - cfg_scale=cfg.cfg_scale, - align=cfg.get("align", None), - aes=cfg.get("aes", None), - flow=cfg.get("flow", None), - camera_motion=cfg.get("camera_motion", None), - resolution=cfg.get("resolution", None), - aspect_ratio=cfg.get("aspect_ratio", None), - num_frames=cfg.get("num_frames", None), - fps=cfg.get("fps", 30), - save_fps=cfg.get("save_fps", - cfg.get("fps", 30) // cfg.get("frame_interval", 1)), - diffusion_model_type=cfg.get("diffusion_model_type", None), - condition_frame_length=cfg.get("condition_frame_length", 5), - condition_frame_edit=cfg.get("condition_frame_edit", 0.0), - use_timestep_transform=True, - video_save_dir=cfg.get("video_save_dir", "sample_outputs"), - dtype=dtype, - device=device, - seed=cfg.get('seed', 0), - ) - - # load prompts - prompts = cfg.get("prompt", None) - start_idx = cfg.get("start_index", 0) - if prompts is None: - if cfg.get("prompt_path", None) is not None: - prompts = DataProcessor.load_prompts(cfg.prompt_path, start_idx, - cfg.get("end_index", None)) - else: - prompts = [cfg.get("prompt_generator", "") - ] * 1_000_000 # endless loop - - pipe(prompts) - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('prompt', - nargs='*', - help="Text prompt(s) to guide video generation") - parser.add_argument("--batch-size", type=int, default=1) - parser.add_argument("--resolution", type=str, default="360p") - parser.add_argument("--aspect-ratio", type=str, default="9:16") - parser.add_argument("--num-timesteps", type=int, default=1000) - parser.add_argument("--num-sampling-steps", type=int, default=30) - parser.add_argument("--seed", type=int, default=2946901) - parser.add_argument("--tllm_model_dir", - type=str, - default='./engine_outputs/') - parser.add_argument("--video_save_dir", - type=str, - default='./sample_outputs/') - parser.add_argument("--gpus_per_node", type=int, default=8) - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument("--debug_mode", action='store_true') - args = parser.parse_args() - - runtime_config = dict( - num_frames=102, - fps=30, - frame_interval=1, - save_fps=30, - diffusion_model_type="STDiT3", - text_encoder="DeepFloyd/t5-v1_1-xxl", - vae="hpcai-tech/OpenSora-VAE-v1.2", - cfg_scale=7.0, - dtype="float16", - condition_frame_length=5, - align=5, - aes=6.5, - flow=None, - prompt="A scene from disaster movie.", - ) - runtime_config.update(vars(args)) - - main(RuntimeConfig(runtime_config)) diff --git a/examples/models/contrib/stdit/scheduler.py b/examples/models/contrib/stdit/scheduler.py deleted file mode 100644 index acd0cac58bdd..000000000000 --- a/examples/models/contrib/stdit/scheduler.py +++ /dev/null @@ -1,101 +0,0 @@ -# Copyright 2024 HPC-AI Technology Inc. - -# 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. - -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# reference: https://github.com/hpcaitech/Open-Sora/blob/main/opensora/schedulers/rf/rectified_flow.py - -import torch -from torch.distributions import LogisticNormal - - -def timestep_transform( - t, - model_kwargs, - base_resolution=512 * 512, - base_num_frames=1, - scale=1.0, - num_timesteps=1, -): - # Force fp16 input to fp32 to avoid nan output - for key in ["height", "width", "num_frames"]: - if model_kwargs[key].dtype == torch.float16: - model_kwargs[key] = model_kwargs[key].float() - t = t / num_timesteps - resolution = model_kwargs["height"] * model_kwargs["width"] - ratio_space = (resolution / base_resolution).sqrt() - # NOTE: currently, we do not take fps into account - # NOTE: temporal_reduction is hardcoded, this should be equal to the temporal reduction factor of the vae - if model_kwargs["num_frames"][0] == 1: - num_frames = torch.ones_like(model_kwargs["num_frames"]) - else: - num_frames = model_kwargs["num_frames"] // 17 * 5 - ratio_time = (num_frames / base_num_frames).sqrt() - ratio = ratio_space * ratio_time * scale - new_t = ratio * t / (1 + (ratio - 1) * t) - new_t = new_t * num_timesteps - return new_t - - -class RFlowScheduler: - - def __init__( - self, - num_timesteps=1000, - num_sampling_steps=10, - sample_method="uniform", - loc=0.0, - scale=1.0, - use_timestep_transform=False, - ): - self.num_timesteps = num_timesteps - self.num_sampling_steps = num_sampling_steps - - assert sample_method in ["uniform", "logit-normal"] - self.sample_method = sample_method - if sample_method == "logit-normal": - self.distribution = LogisticNormal(torch.tensor([loc]), - torch.tensor([scale])) - self.sample_t = lambda x: self.distribution.sample( - (x.shape[0], ))[:, 0].to(x.device) - self.use_timestep_transform = use_timestep_transform - - def add_noise( - self, - original_samples: torch.FloatTensor, - noise: torch.FloatTensor, - timesteps: torch.IntTensor, - ) -> torch.FloatTensor: - timepoints = timesteps.float() / self.num_timesteps - timepoints = 1 - timepoints # [1,1/1000] - # timepoint (bsz) noise: (bsz, 4, frame, w ,h) - # expand timepoint to noise shape - timepoints = timepoints.unsqueeze(1).unsqueeze(1).unsqueeze( - 1).unsqueeze(1) - timepoints = timepoints.repeat(1, noise.shape[1], noise.shape[2], - noise.shape[3], noise.shape[4]) - return timepoints * original_samples + (1 - timepoints) * noise diff --git a/examples/models/contrib/stdit/text_encoder.py b/examples/models/contrib/stdit/text_encoder.py deleted file mode 100644 index c48ea5b1cfca..000000000000 --- a/examples/models/contrib/stdit/text_encoder.py +++ /dev/null @@ -1,445 +0,0 @@ -# Copyright (C) 2023 PixArt-alpha/PixArt-alpha -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. - -# Copyright 2024 HPC-AI Technology Inc. - -# 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. - -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# reference: https://github.com/hpcaitech/Open-Sora/blob/main/opensora/models/text_encoder/t5.py - -import collections.abc -from functools import partial -from itertools import repeat - -import torch -from colossalai.shardformer import ShardConfig, ShardFormer -from colossalai.shardformer.modeling.jit import get_jit_fused_dropout_add_func -from colossalai.shardformer.modeling.t5 import ( - get_jit_fused_T5_layer_ff_forward, get_T5_layer_self_attention_forward) -from colossalai.shardformer.policies.base_policy import ( - Policy, SubModuleReplacementDescription) -from transformers import AutoTokenizer, T5EncoderModel -from transformers.models.t5.modeling_t5 import (T5LayerFF, T5LayerSelfAttention, - T5Stack) - - -def default(var, default_var): - return default_var if var is None else var - - -def _ntuple(n): - - def parse(x): - if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): - return tuple(x) - return tuple(repeat(x, n)) - - return parse - - -to_1tuple = _ntuple(1) -to_2tuple = _ntuple(2) -to_3tuple = _ntuple(3) -to_4tuple = _ntuple(4) -to_ntuple = _ntuple - - -class T5LayerNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-6): - """ - Construct a layernorm module in the T5 style. No bias and no subtraction of mean. - """ - super().__init__() - self.weight = torch.nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - - def forward(self, hidden_states): - # T5 uses a layer_norm which only scales and doesn't shift, which is also known as Root Mean - # Square Layer Normalization https://arxiv.org/abs/1910.07467 thus variance is calculated - # w/o mean and there is no bias. Additionally we want to make sure that the accumulation for - # half-precision inputs is done in fp32 - variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + - self.variance_epsilon) - # convert into half-precision if necessary - if self.weight.dtype in [torch.float16, torch.bfloat16]: - hidden_states = hidden_states.to(self.weight.dtype) - return self.weight * hidden_states - - @staticmethod - def from_native_module(module, *args, **kwargs): - assert module.__class__.__name__ == "FusedRMSNorm", ( - "Recovering T5LayerNorm requires the original layer to be apex's Fused RMS Norm." - "Apex's fused norm is automatically used by Hugging Face Transformers https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/modeling_t5.py#L265C5-L265C48" - ) - layer_norm = T5LayerNorm(module.normalized_shape, eps=module.eps) - layer_norm.weight.data.copy_(module.weight.data) - layer_norm = layer_norm.to(module.weight.device) - return layer_norm - - -class T5EncoderPolicy(Policy): - - def config_sanity_check(self): - assert not self.shard_config.enable_tensor_parallelism - assert not self.shard_config.enable_flash_attention - - def preprocess(self): - return self.model - - def module_policy(self): - policy = {} - # check whether apex is installed - try: - # recover hf from fused rms norm to T5 norm which is faster - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="layer_norm", - target_module=T5LayerNorm, - ), - policy=policy, - target_key=T5LayerFF, - ) - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="layer_norm", target_module=T5LayerNorm), - policy=policy, - target_key=T5LayerSelfAttention, - ) - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="final_layer_norm", target_module=T5LayerNorm), - policy=policy, - target_key=T5Stack, - ) - except (ImportError, ModuleNotFoundError): - pass - # use jit operator - if self.shard_config.enable_jit_fused: - self.append_or_create_method_replacement( - description={ - "forward": get_jit_fused_T5_layer_ff_forward(), - "dropout_add": get_jit_fused_dropout_add_func(), - }, - policy=policy, - target_key=T5LayerFF, - ) - self.append_or_create_method_replacement( - description={ - "forward": get_T5_layer_self_attention_forward(), - "dropout_add": get_jit_fused_dropout_add_func(), - }, - policy=policy, - target_key=T5LayerSelfAttention, - ) - return policy - - def postprocess(self): - return self.model - - -class T5Embedder: - - def __init__( - self, - device, - from_pretrained=None, - *, - cache_dir=None, - hf_token=None, - use_text_preprocessing=True, - t5_model_kwargs=None, - torch_dtype=None, - use_offload_folder=None, - model_max_length=120, - local_files_only=False, - ): - self.device = torch.device(device) - self.torch_dtype = torch_dtype or torch.float16 - self.cache_dir = cache_dir - - if t5_model_kwargs is None: - t5_model_kwargs = { - "low_cpu_mem_usage": True, - "torch_dtype": self.torch_dtype, - } - - if use_offload_folder is not None: - t5_model_kwargs["offload_folder"] = use_offload_folder - t5_model_kwargs["device_map"] = { - "shared": self.device, - "encoder.embed_tokens": self.device, - "encoder.block.0": self.device, - "encoder.block.1": self.device, - "encoder.block.2": self.device, - "encoder.block.3": self.device, - "encoder.block.4": self.device, - "encoder.block.5": self.device, - "encoder.block.6": self.device, - "encoder.block.7": self.device, - "encoder.block.8": self.device, - "encoder.block.9": self.device, - "encoder.block.10": self.device, - "encoder.block.11": self.device, - "encoder.block.12": "disk", - "encoder.block.13": "disk", - "encoder.block.14": "disk", - "encoder.block.15": "disk", - "encoder.block.16": "disk", - "encoder.block.17": "disk", - "encoder.block.18": "disk", - "encoder.block.19": "disk", - "encoder.block.20": "disk", - "encoder.block.21": "disk", - "encoder.block.22": "disk", - "encoder.block.23": "disk", - "encoder.final_layer_norm": "disk", - "encoder.dropout": "disk", - } - else: - t5_model_kwargs["device_map"] = { - "shared": self.device, - "encoder": self.device, - } - - self.use_text_preprocessing = use_text_preprocessing - self.hf_token = hf_token - - self.tokenizer = AutoTokenizer.from_pretrained( - from_pretrained, - cache_dir=cache_dir, - local_files_only=local_files_only, - ) - self.model = T5EncoderModel.from_pretrained( - from_pretrained, - cache_dir=cache_dir, - local_files_only=local_files_only, - **t5_model_kwargs, - ).eval() - self.model_max_length = model_max_length - - def get_text_embeddings(self, texts): - text_tokens_and_mask = self.tokenizer( - texts, - max_length=self.model_max_length, - padding="max_length", - truncation=True, - return_attention_mask=True, - add_special_tokens=True, - return_tensors="pt", - ) - - input_ids = text_tokens_and_mask["input_ids"].to(self.device) - attention_mask = text_tokens_and_mask["attention_mask"].to(self.device) - with torch.no_grad(): - text_encoder_embs = self.model( - input_ids=input_ids, - attention_mask=attention_mask, - )["last_hidden_state"].detach() - return text_encoder_embs, attention_mask - - -class Mlp(torch.nn.Module): - - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - act_layer=torch.nn.GELU(), - norm_layer=None, - bias=True, - drop=0., - use_conv=False, - ): - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features - bias = to_2tuple(bias) - drop_probs = to_2tuple(drop) - linear_layer = partial(torch.nn.Conv2d, - kernel_size=1) if use_conv else torch.nn.Linear - - self.fc1 = linear_layer(in_features, hidden_features, bias=bias[0]) - self.act = act_layer - self.drop1 = torch.nn.Dropout(drop_probs[0]) - self.norm = norm_layer( - hidden_features) if norm_layer is not None else torch.nn.Identity() - self.fc2 = linear_layer(hidden_features, out_features, bias=bias[1]) - self.drop2 = torch.nn.Dropout(drop_probs[1]) - - def forward(self, x): - x = self.fc1(x) - x = self.act(x) - x = self.drop1(x) - x = self.norm(x) - x = self.fc2(x) - x = self.drop2(x) - return x - - -class CaptionEmbedder(torch.nn.Module): - - def __init__( - self, - in_channels, - hidden_size, - uncond_prob, - act_layer=torch.nn.GELU(approximate="tanh"), - token_num=120, - ): - super().__init__() - self.y_proj = Mlp( - in_features=in_channels, - hidden_features=hidden_size, - out_features=hidden_size, - act_layer=act_layer, - drop=0, - ) - self.register_buffer( - "y_embedding", - torch.randn(token_num, in_channels) / in_channels**0.5, - ) - self.uncond_prob = uncond_prob - - def token_drop(self, caption, force_drop_ids=None): - if force_drop_ids is None: - drop_ids = torch.rand(caption.shape[0]).cuda() < self.uncond_prob - else: - drop_ids = force_drop_ids == 1 - caption = torch.where(drop_ids[:, None, None, None], self.y_embedding, - caption) - return caption - - def forward(self, caption, train, force_drop_ids=None): - if train: - assert caption.shape[2:] == self.y_embedding.shape - use_dropout = self.uncond_prob > 0 - if (train and use_dropout) or (force_drop_ids is not None): - caption = self.token_drop(caption, force_drop_ids) - caption = self.y_proj(caption) - return caption - - -class T5Encoder: - - def __init__( - self, - from_pretrained=None, - model_max_length=120, - caption_channels=4096, - hidden_size=1152, - class_dropout_prob=0.1, - y_embedding=None, - device="cuda", - dtype=torch.float, - cache_dir=None, - shardformer=False, - local_files_only=False, - ): - assert from_pretrained is not None, "Please specify the path to the T5 model" - self.t5 = T5Embedder( - device=device, - torch_dtype=dtype, - from_pretrained=from_pretrained, - cache_dir=cache_dir, - model_max_length=model_max_length, - local_files_only=local_files_only, - ) - self.t5.model.to(dtype=dtype) - # [NOTE] disable y_embedder - if False: - self.y_embedder = CaptionEmbedder( - in_channels=caption_channels, - hidden_size=hidden_size, - uncond_prob=class_dropout_prob, - act_layer=torch.nn.GELU(approximate="tanh"), - token_num=model_max_length, - ) - else: - self.y_embedder = None - self.y_embedding = default( - y_embedding, - torch.randn(model_max_length, caption_channels) / - caption_channels**0.5).to(device) - - self.model_max_length = model_max_length - self.output_dim = self.t5.model.config.d_model - self.dtype = dtype - - if shardformer: - self.shardformer_t5() - - def shardformer_t5(self): - shard_config = ShardConfig( - tensor_parallel_process_group=None, - pipeline_stage_manager=None, - enable_tensor_parallelism=False, - enable_fused_normalization=False, - enable_flash_attention=False, - enable_jit_fused=True, - enable_sequence_parallelism=False, - enable_sequence_overlap=False, - ) - shard_former = ShardFormer(shard_config=shard_config) - optim_model, _ = shard_former.optimize(self.t5.model, - policy=T5EncoderPolicy()) - self.t5.model = optim_model.to(self.dtype) - - # ensure the weights are frozen - for p in self.t5.model.parameters(): - p.requires_grad = False - - def encode(self, text): - caption_embs, emb_masks = self.t5.get_text_embeddings(text) - caption_embs = caption_embs[:, None] - return dict(y=caption_embs, mask=emb_masks) - - def null(self, n): - if self.y_embedder is None: - null_y = self.y_embedding[None].repeat(n, 1, 1)[:, None] - else: - null_y = self.y_embedder.y_embedding[None].repeat(n, 1, 1)[:, None] - return null_y - - def encode_with_null(self, text): - batch_size = len(text) - encoded_outputs = self.encode(text) - y_null = self.null(batch_size) - encoded_outputs["y"] = torch.cat([encoded_outputs["y"], y_null], 0) - return encoded_outputs diff --git a/examples/models/contrib/stdit/utils.py b/examples/models/contrib/stdit/utils.py deleted file mode 100644 index 88a7269bc8e8..000000000000 --- a/examples/models/contrib/stdit/utils.py +++ /dev/null @@ -1,997 +0,0 @@ -# Copyright 2024 HPC-AI Technology Inc. - -# 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. - -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# reference: https://github.com/hpcaitech/Open-Sora/blob/main/opensora/utils/ckpt_utils.py - -import html -import json -import os -import random -import re -import sys -import urllib.parse as ul - -import ftfy -import numpy as np -import pandas as pd -import requests -import torch -import video_transforms -from aspect import get_image_size, get_num_frames -from bs4 import BeautifulSoup -from colossalai.checkpoint_io import GeneralCheckpointIO -from PIL import Image -from safetensors.torch import load_file -from torchvision import transforms -from torchvision.datasets.folder import IMG_EXTENSIONS, pil_loader -from torchvision.io import read_video, write_video -from torchvision.utils import save_image - -VID_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv") - -URL_REGEX = re.compile( - r"^(?:http|ftp)s?://" # http:// or https:// - r"(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+(?:[A-Z]{2,6}\.?|[A-Z0-9-]{2,}\.?)|" # domain... - r"localhost|" # localhost... - r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})" # ...or ip - r"(?::\d+)?" # optional port - r"(?:/?|[/?]\S+)$", - re.IGNORECASE, -) - -HF_ENDPOINT = os.environ.get("HF_ENDPOINT") -if HF_ENDPOINT is None: - HF_ENDPOINT = "https://huggingface.co" -PRETRAINED_MODELS = { - "DiT-XL-2-512x512.pt": - "https://dl.fbaipublicfiles.com/DiT/models/DiT-XL-2-512x512.pt", - "DiT-XL-2-256x256.pt": - "https://dl.fbaipublicfiles.com/DiT/models/DiT-XL-2-256x256.pt", - "Latte-XL-2-256x256-ucf101.pt": - HF_ENDPOINT + "/maxin-cn/Latte/resolve/main/ucf101.pt", - "PixArt-XL-2-256x256.pth": - HF_ENDPOINT + - "/PixArt-alpha/PixArt-alpha/resolve/main/PixArt-XL-2-256x256.pth", - "PixArt-XL-2-SAM-256x256.pth": - HF_ENDPOINT + - "/PixArt-alpha/PixArt-alpha/resolve/main/PixArt-XL-2-SAM-256x256.pth", - "PixArt-XL-2-512x512.pth": - HF_ENDPOINT + - "/PixArt-alpha/PixArt-alpha/resolve/main/PixArt-XL-2-512x512.pth", - "PixArt-XL-2-1024-MS.pth": - HF_ENDPOINT + - "/PixArt-alpha/PixArt-alpha/resolve/main/PixArt-XL-2-1024-MS.pth", - "OpenSora-v1-16x256x256.pth": - HF_ENDPOINT + - "/hpcai-tech/Open-Sora/resolve/main/OpenSora-v1-16x256x256.pth", - "OpenSora-v1-HQ-16x256x256.pth": - HF_ENDPOINT + - "/hpcai-tech/Open-Sora/resolve/main/OpenSora-v1-HQ-16x256x256.pth", - "OpenSora-v1-HQ-16x512x512.pth": - HF_ENDPOINT + - "/hpcai-tech/Open-Sora/resolve/main/OpenSora-v1-HQ-16x512x512.pth", - "PixArt-Sigma-XL-2-256x256.pth": - HF_ENDPOINT + - "/PixArt-alpha/PixArt-Sigma/resolve/main/PixArt-Sigma-XL-2-256x256.pth", - "PixArt-Sigma-XL-2-512-MS.pth": - HF_ENDPOINT + - "/PixArt-alpha/PixArt-Sigma/resolve/main/PixArt-Sigma-XL-2-512-MS.pth", - "PixArt-Sigma-XL-2-1024-MS.pth": - HF_ENDPOINT + - "/PixArt-alpha/PixArt-Sigma/resolve/main/PixArt-Sigma-XL-2-1024-MS.pth", - "PixArt-Sigma-XL-2-2K-MS.pth": - HF_ENDPOINT + - "/PixArt-alpha/PixArt-Sigma/resolve/main/PixArt-Sigma-XL-2-2K-MS.pth", -} - - -def is_img(path): - ext = os.path.splitext(path)[-1].lower() - return ext in IMG_EXTENSIONS - - -def is_vid(path): - ext = os.path.splitext(path)[-1].lower() - return ext in VID_EXTENSIONS - - -def is_url(url): - return re.match(URL_REGEX, url) is not None - - -def read_file(input_path): - if input_path.endswith(".csv"): - return pd.read_csv(input_path) - elif input_path.endswith(".parquet"): - return pd.read_parquet(input_path) - else: - raise NotImplementedError(f"Unsupported file format: {input_path}") - - -def download_url(input_path): - output_dir = "cache" - os.makedirs(output_dir, exist_ok=True) - base_name = os.path.basename(input_path) - output_path = os.path.join(output_dir, base_name) - try: - img_data = requests.get(input_path).content - except requests.exceptions.RequestException as e: - print(f"Error downloading URL {input_path}: {e}") - return None - with open(output_path, "wb") as handler: - handler.write(img_data) - print(f"URL {input_path} downloaded to {output_path}") - return output_path - - -def reparameter(ckpt, name=None, model=None): - model_name = name - name = os.path.basename(name) - if not torch.distributed.is_initialized() or torch.distributed.get_rank( - ) == 0: - print("loading pretrained model: %s", model_name) - if name in ["DiT-XL-2-512x512.pt", "DiT-XL-2-256x256.pt"]: - ckpt["x_embedder.proj.weight"] = ckpt[ - "x_embedder.proj.weight"].unsqueeze(2) - del ckpt["pos_embed"] - if name in ["Latte-XL-2-256x256-ucf101.pt"]: - ckpt = ckpt["ema"] - ckpt["x_embedder.proj.weight"] = ckpt[ - "x_embedder.proj.weight"].unsqueeze(2) - del ckpt["pos_embed"] - del ckpt["temp_embed"] - if name in [ - "PixArt-XL-2-256x256.pth", - "PixArt-XL-2-SAM-256x256.pth", - "PixArt-XL-2-512x512.pth", - "PixArt-XL-2-1024-MS.pth", - "PixArt-Sigma-XL-2-256x256.pth", - "PixArt-Sigma-XL-2-512-MS.pth", - "PixArt-Sigma-XL-2-1024-MS.pth", - "PixArt-Sigma-XL-2-2K-MS.pth", - ]: - ckpt = ckpt["state_dict"] - ckpt["x_embedder.proj.weight"] = ckpt[ - "x_embedder.proj.weight"].unsqueeze(2) - if "pos_embed" in ckpt: - del ckpt["pos_embed"] - if name in [ - "PixArt-1B-2.pth", - ]: - ckpt = ckpt["state_dict"] - if "pos_embed" in ckpt: - del ckpt["pos_embed"] - # no need pos_embed - if "pos_embed_temporal" in ckpt: - del ckpt["pos_embed_temporal"] - if "pos_embed" in ckpt: - del ckpt["pos_embed"] - # different text length - if "y_embedder.y_embedding" in ckpt: - if ckpt["y_embedder.y_embedding"].shape[ - 0] < model.y_embedder.y_embedding.shape[0]: - print( - "Extend y_embedding from %s to %s", - ckpt["y_embedder.y_embedding"].shape[0], - model.y_embedder.y_embedding.shape[0], - ) - additional_length = model.y_embedder.y_embedding.shape[0] - ckpt[ - "y_embedder.y_embedding"].shape[0] - new_y_embedding = torch.zeros(additional_length, - model.y_embedder.y_embedding.shape[1]) - new_y_embedding[:] = ckpt["y_embedder.y_embedding"][-1] - ckpt["y_embedder.y_embedding"] = torch.cat( - [ckpt["y_embedder.y_embedding"], new_y_embedding], dim=0) - elif ckpt["y_embedder.y_embedding"].shape[ - 0] > model.y_embedder.y_embedding.shape[0]: - print( - "Shrink y_embedding from %s to %s", - ckpt["y_embedder.y_embedding"].shape[0], - model.y_embedder.y_embedding.shape[0], - ) - ckpt["y_embedder.y_embedding"] = ckpt[ - "y_embedder.y_embedding"][:model.y_embedder.y_embedding. - shape[0]] - # stdit3 special case - if type(model).__name__ == "STDiT3" and "PixArt-Sigma" in name: - ckpt_keys = list(ckpt.keys()) - for key in ckpt_keys: - if "blocks." in key: - ckpt[key.replace("blocks.", "spatial_blocks.")] = ckpt[key] - del ckpt[key] - return ckpt - - -def download_model(model_name=None, local_path=None, url=None): - """ - Downloads a pre-trained DiT model from the web. - """ - if model_name is not None: - assert model_name in PRETRAINED_MODELS - local_path = f"pretrained_models/{model_name}" - web_path = PRETRAINED_MODELS[model_name] - else: - assert local_path is not None - assert url is not None - web_path = url - if not os.path.isfile(local_path): - os.makedirs("pretrained_models", exist_ok=True) - dir_name = os.path.dirname(local_path) - file_name = os.path.basename(local_path) - download_url(web_path, dir_name, file_name) - model = torch.load(local_path, map_location=lambda storage, loc: storage) - return model - - -def find_model(model_name, model=None): - """ - Finds a pre-trained DiT model, downloading it if necessary. Alternatively, loads a model from a local path. - """ - if model_name in PRETRAINED_MODELS: # Find/download our pre-trained DiT checkpoints - model_ckpt = download_model(model_name) - model_ckpt = reparameter(model_ckpt, model_name, model=model) - else: # Load a custom DiT checkpoint: - assert os.path.isfile( - model_name), f"Could not find DiT checkpoint at {model_name}" - model_ckpt = torch.load(model_name, - map_location=lambda storage, loc: storage) - model_ckpt = reparameter(model_ckpt, model_name, model=model) - return model_ckpt - - -def load_from_sharded_state_dict(model, - ckpt_path, - model_name="model", - strict=False): - ckpt_io = GeneralCheckpointIO() - ckpt_io.load_model(model, - os.path.join(ckpt_path, model_name), - strict=strict) - - -def load_checkpoint(model, - ckpt_path, - save_as_pt=False, - model_name="model", - strict=False): - if ckpt_path.endswith(".pt") or ckpt_path.endswith(".pth"): - state_dict = find_model(ckpt_path, model=model) - missing_keys, unexpected_keys = model.load_state_dict(state_dict, - strict=strict) - print("Missing keys: %s", missing_keys) - print("Unexpected keys: %s", unexpected_keys) - elif ckpt_path.endswith(".safetensors"): - state_dict = load_file(ckpt_path) - missing_keys, unexpected_keys = model.load_state_dict(state_dict, - strict=False) - print(f"Missing keys: {missing_keys}") - print(f"Unexpected keys: {unexpected_keys}") - elif os.path.isdir(ckpt_path): - load_from_sharded_state_dict(model, - ckpt_path, - model_name, - strict=strict) - print("Model checkpoint loaded from %s", ckpt_path) - if save_as_pt: - save_path = os.path.join(ckpt_path, model_name + "_ckpt.pt") - torch.save(model.state_dict(), save_path) - print("Model checkpoint saved to %s", save_path) - else: - raise ValueError(f"Invalid checkpoint path: {ckpt_path}") - - -def print_progress_bar(iteration, total, length=40): - iteration += 1 - filled_length = int(length * iteration // total) - bar = '█' * filled_length + '-' * (length - filled_length) - sys.stdout.write(f'\rDenoising steps: |{bar}| {iteration}/{total}') - sys.stdout.flush() - - -class PromptProcessor(): - - @staticmethod - def load_prompts(prompt_path, start_idx=None, end_idx=None): - with open(prompt_path, "r") as f: - prompts = [line.strip() for line in f.readlines()] - prompts = prompts[start_idx:end_idx] - return prompts - - @staticmethod - def extract_json_from_prompts(prompts, reference, mask_strategy): - ret_prompts = [] - for i, prompt in enumerate(prompts): - parts = re.split(r"(?=[{])", prompt) - assert len(parts) <= 2, f"Invalid prompt: {prompt}" - ret_prompts.append(parts[0]) - if len(parts) > 1: - additional_info = json.loads(parts[1]) - for key in additional_info: - assert key in ["reference_path", - "mask_strategy"], f"Invalid key: {key}" - if key == "reference_path": - reference[i] = additional_info[key] - elif key == "mask_strategy": - mask_strategy[i] = additional_info[key] - return ret_prompts, reference, mask_strategy - - @staticmethod - def split_prompt(prompt_text): - if prompt_text.startswith("|0|"): - # this is for prompts which look like - # |0| a beautiful day |1| a sunny day |2| a rainy day - # we want to parse it into a list of prompts with the loop index - prompt_list = prompt_text.split("|")[1:] - text_list = [] - loop_idx = [] - for i in range(0, len(prompt_list), 2): - start_loop = int(prompt_list[i]) - text = prompt_list[i + 1].strip() - text_list.append(text) - loop_idx.append(start_loop) - return text_list, loop_idx - else: - return [prompt_text], None - - @staticmethod - def merge_prompt(text_list, loop_idx_list=None): - if loop_idx_list is None: - return text_list[0] - else: - prompt = "" - for i, text in enumerate(text_list): - prompt += f"|{loop_idx_list[i]}|{text}" - return prompt - - @staticmethod - def extract_prompts_loop(prompts, num_loop): - ret_prompts = [] - for prompt in prompts: - if prompt.startswith("|0|"): - prompt_list = prompt.split("|")[1:] - text_list = [] - for i in range(0, len(prompt_list), 2): - start_loop = int(prompt_list[i]) - text = prompt_list[i + 1] - end_loop = int(prompt_list[i + 2]) if i + 2 < len( - prompt_list) else num_loop + 1 - text_list.extend([text] * (end_loop - start_loop)) - prompt = text_list[num_loop] - ret_prompts.append(prompt) - return ret_prompts - - @staticmethod - def append_score_to_prompts(prompts, - aes=None, - flow=None, - camera_motion=None): - new_prompts = [] - for prompt in prompts: - new_prompt = prompt - if aes is not None and "aesthetic score:" not in prompt: - new_prompt = f"{new_prompt} aesthetic score: {aes:.1f}." - if flow is not None and "motion score:" not in prompt: - new_prompt = f"{new_prompt} motion score: {flow:.1f}." - if camera_motion is not None and "camera motion:" not in prompt: - new_prompt = f"{new_prompt} camera motion: {camera_motion}." - new_prompts.append(new_prompt) - return new_prompts - - @classmethod - def basic_clean(cls, text): - text = ftfy.fix_text(text) - text = html.unescape(html.unescape(text)) - return text.strip() - - @classmethod - def clean_caption(cls, caption): - BAD_PUNCT_REGEX = re.compile(r"[" + "#®•©™&@·º½¾¿¡§~" + r"\)" + r"\(" + - r"\]" + r"\[" + r"\}" + r"\{" + r"\|" + - "\\" + "\/" + "\*" + r"]{1,}") # noqa - caption = str(caption) - caption = ul.unquote_plus(caption) - caption = caption.strip().lower() - caption = re.sub("", "person", caption) - # urls: - caption = re.sub( - r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - caption = re.sub( - r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - # html: - caption = BeautifulSoup(caption, features="html.parser").text - # @ - caption = re.sub(r"@[\w\d]+\b", "", caption) - # 31C0—31EF CJK Strokes - # 31F0—31FF Katakana Phonetic Extensions - # 3200—32FF Enclosed CJK Letters and Months - # 3300—33FF CJK Compatibility - # 3400—4DBF CJK Unified Ideographs Extension A - # 4DC0—4DFF Yijing Hexagram Symbols - # 4E00—9FFF CJK Unified Ideographs - caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) - caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) - caption = re.sub(r"[\u3200-\u32ff]+", "", caption) - caption = re.sub(r"[\u3300-\u33ff]+", "", caption) - caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) - caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) - caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) - ####################################################### - - # все виды тире / all types of dash --> "-" - caption = re.sub( - r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa - "-", - caption, - ) - # кавычки к одному стандарту - caption = re.sub(r"[`´«»“”¨]", '"', caption) - caption = re.sub(r"[‘’]", "'", caption) - # " - caption = re.sub(r""?", "", caption) - # & - caption = re.sub(r"&", "", caption) - # ip addresses: - caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) - # article ids: - caption = re.sub(r"\d:\d\d\s+$", "", caption) - # \n - caption = re.sub(r"\\n", " ", caption) - # "#123" - caption = re.sub(r"#\d{1,3}\b", "", caption) - # "#12345.." - caption = re.sub(r"#\d{5,}\b", "", caption) - # "123456.." - caption = re.sub(r"\b\d{6,}\b", "", caption) - # filenames: - caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", - "", caption) - # - caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" - caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" - caption = re.sub(BAD_PUNCT_REGEX, r" ", - caption) # ***AUSVERKAUFT***, #AUSVERKAUFT - caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " - # this-is-my-cute-cat / this_is_my_cute_cat - regex2 = re.compile(r"(?:\-|\_)") - if len(re.findall(regex2, caption)) > 3: - caption = re.sub(regex2, " ", caption) - caption = cls.basic_clean(caption) - caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 - caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc - caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 - caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) - caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) - caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) - caption = re.sub( - r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", - caption) - caption = re.sub(r"\bpage\s+\d+\b", "", caption) - caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", - caption) # j2d1a2a... - caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) - caption = re.sub(r"\b\s+\:\s+", r": ", caption) - caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) - caption = re.sub(r"\s+", " ", caption) - caption.strip() - caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) - caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) - caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) - caption = re.sub(r"^\.\S+$", "", caption) - return caption.strip() - - @classmethod - def text_preprocessing(cls, text, use_text_preprocessing: bool = True): - if use_text_preprocessing: - # The exact text cleaning as was in the training stage: - text = cls.clean_caption(text) - text = cls.clean_caption(text) - return text - else: - return text.lower().strip() - - -class VideoProcessor(): - MASK_DEFAULT = ["0", "0", "0", "0", "1", "0"] - - @staticmethod - def get_image_size(resolution, ar_ratio): - return get_image_size(resolution=resolution, ar_ratio=ar_ratio) - - @staticmethod - def get_num_frames(num_frames): - return get_num_frames(num_frames=num_frames) - - @staticmethod - def get_latent_size(vae, input_size): - return vae.get_latent_size(input_size=input_size) - - @staticmethod - def center_crop_arr(pil_image, image_size): - """ - Center cropping implementation from ADM. - https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126 - """ - while min(*pil_image.size) >= 2 * image_size: - pil_image = pil_image.resize(tuple(x // 2 for x in pil_image.size), - resample=Image.BOX) - - scale = image_size / min(*pil_image.size) - pil_image = pil_image.resize(tuple( - round(x * scale) for x in pil_image.size), - resample=Image.BICUBIC) - - arr = np.array(pil_image) - crop_y = (arr.shape[0] - image_size) // 2 - crop_x = (arr.shape[1] - image_size) // 2 - return Image.fromarray(arr[crop_y:crop_y + image_size, - crop_x:crop_x + image_size]) - - @staticmethod - def resize_crop_to_fill(pil_image, image_size): - w, h = pil_image.size # PIL is (W, H) - th, tw = image_size - rh, rw = th / h, tw / w - if rh > rw: - sh, sw = th, round(w * rh) - image = pil_image.resize((sw, sh), Image.BICUBIC) - i = 0 - j = int(round((sw - tw) / 2.0)) - else: - sh, sw = round(h * rw), tw - image = pil_image.resize((sw, sh), Image.BICUBIC) - i = int(round((sh - th) / 2.0)) - j = 0 - arr = np.array(image) - assert i + th <= arr.shape[0] and j + tw <= arr.shape[1] - return Image.fromarray(arr[i:i + th, j:j + tw]) - - @classmethod - def get_transforms_image(cls, name="center", image_size=(256, 256)): - if name is None: - return None - elif name == "center": - assert image_size[0] == image_size[ - 1], "Image size must be square for center crop" - transform = transforms.Compose([ - transforms.Lambda(lambda pil_image: cls.center_crop_arr( - pil_image, image_size[0])), - # transforms.RandomHorizontalFlip(), - transforms.ToTensor(), - transforms.Normalize(mean=[0.5, 0.5, 0.5], - std=[0.5, 0.5, 0.5], - inplace=True), - ]) - elif name == "resize_crop": - transform = transforms.Compose([ - transforms.Lambda(lambda pil_image: cls.resize_crop_to_fill( - pil_image, image_size)), - transforms.ToTensor(), - transforms.Normalize(mean=[0.5, 0.5, 0.5], - std=[0.5, 0.5, 0.5], - inplace=True), - ]) - else: - raise NotImplementedError(f"Transform {name} not implemented") - return transform - - @classmethod - def get_transforms_video(cls, name="center", image_size=(256, 256)): - if name is None: - return None - elif name == "center": - assert image_size[0] == image_size[ - 1], "image_size must be square for center crop" - transform_video = transforms.Compose([ - video_transforms.ToTensorVideo(), # TCHW - # video_transforms.RandomHorizontalFlipVideo(), - video_transforms.UCFCenterCropVideo(image_size[0]), - transforms.Normalize(mean=[0.5, 0.5, 0.5], - std=[0.5, 0.5, 0.5], - inplace=True), - ]) - elif name == "resize_crop": - transform_video = transforms.Compose([ - video_transforms.ToTensorVideo(), # TCHW - video_transforms.ResizeCrop(image_size), - transforms.Normalize(mean=[0.5, 0.5, 0.5], - std=[0.5, 0.5, 0.5], - inplace=True), - ]) - else: - raise NotImplementedError(f"Transform {name} not implemented") - return transform_video - - @classmethod - def read_image_from_path(cls, - path, - transform=None, - transform_name="center", - num_frames=1, - image_size=(256, 256)): - image = pil_loader(path) - if transform is None: - transform = cls.get_transforms_image(image_size=image_size, - name=transform_name) - image = transform(image) - video = image.unsqueeze(0).repeat(num_frames, 1, 1, 1) - video = video.permute(1, 0, 2, 3) - return video - - @classmethod - def read_video_from_path(cls, - path, - transform=None, - transform_name="center", - image_size=(256, 256)): - vframes, aframes, info = read_video(filename=path, - pts_unit="sec", - output_format="TCHW") - if transform is None: - transform = cls.get_transforms_video(image_size=image_size, - name=transform_name) - video = transform(vframes) # T C H W - video = video.permute(1, 0, 2, 3) - return video - - @classmethod - def read_from_path(cls, path, image_size, transform_name="center"): - if is_url(path): - path = download_url(path) - ext = os.path.splitext(path)[-1].lower() - if ext.lower() in VID_EXTENSIONS: - return cls.read_video_from_path(path, - image_size=image_size, - transform_name=transform_name) - else: - assert ext.lower( - ) in IMG_EXTENSIONS, f"Unsupported file format: {ext}" - return cls.read_image_from_path(path, - image_size=image_size, - transform_name=transform_name) - - @classmethod - def collect_references_batch(cls, vae, reference_paths, image_size): - refs_x = [] # refs_x: [batch, ref_num, C, T, H, W] - for reference_path in reference_paths: - if reference_path == "": - refs_x.append([]) - continue - ref_path = reference_path.split(";") - ref = [] - for r_path in ref_path: - r = cls.read_from_path(r_path, - image_size, - transform_name="resize_crop") - r_x = vae.encode(r.unsqueeze(0).to(vae.device, vae.dtype)) - r_x = r_x.squeeze(0) - ref.append(r_x) - refs_x.append(ref) - return refs_x - - @staticmethod - def append_generated(vae, generated_video, refs_x, mask_strategy, loop_i, - condition_frame_length, condition_frame_edit): - ref_x = vae.encode(generated_video) - for j, refs in enumerate(refs_x): - if refs is None: - refs_x[j] = [ref_x[j]] - else: - refs.append(ref_x[j]) - if mask_strategy[j] is None or mask_strategy[j] == "": - mask_strategy[j] = "" - else: - mask_strategy[j] += ";" - mask_strategy[ - j] += f"{loop_i},{len(refs)-1},-{condition_frame_length},0,{condition_frame_length},{condition_frame_edit}" - return refs_x, mask_strategy - - @classmethod - def parse_mask_strategy(cls, mask_strategy): - mask_batch = [] - if mask_strategy == "" or mask_strategy is None: - return mask_batch - mask_strategy = mask_strategy.split(";") - for mask in mask_strategy: - mask_group = mask.split(",") - num_group = len(mask_group) - assert num_group >= 1 and num_group <= 6, f"Invalid mask strategy: {mask}" - mask_group.extend(cls.MASK_DEFAULT[num_group:]) - for i in range(5): - mask_group[i] = int(mask_group[i]) - mask_group[5] = float(mask_group[5]) - mask_batch.append(mask_group) - return mask_batch - - @classmethod - def find_nearest_point(cls, value, point, max_value): - t = value // point - if value % point > point / 2 and t < max_value // point - 1: - t += 1 - return t * point - - @classmethod - def apply_mask_strategy(cls, z, refs_x, mask_strategys, loop_i, align=None): - masks = [] - no_mask = True - for i, mask_strategy in enumerate(mask_strategys): - no_mask = False - mask = torch.ones(z.shape[2], dtype=torch.float, device=z.device) - mask_strategy = cls.parse_mask_strategy(mask_strategy) - for mst in mask_strategy: - loop_id, m_id, m_ref_start, m_target_start, m_length, edit_ratio = mst - if loop_id != loop_i: - continue - ref = refs_x[i][m_id] - - if m_ref_start < 0: - # ref: [C, T, H, W] - m_ref_start = ref.shape[1] + m_ref_start - if m_target_start < 0: - # z: [B, C, T, H, W] - m_target_start = z.shape[2] + m_target_start - if align is not None: - m_ref_start = cls.find_nearest_point( - m_ref_start, align, ref.shape[1]) - m_target_start = cls.find_nearest_point( - m_target_start, align, z.shape[2]) - m_length = min(m_length, z.shape[2] - m_target_start, - ref.shape[1] - m_ref_start) - z[i, :, m_target_start:m_target_start + - m_length] = ref[:, m_ref_start:m_ref_start + m_length] - mask[m_target_start:m_target_start + m_length] = edit_ratio - masks.append(mask) - if no_mask: - return None - masks = torch.stack(masks) - return masks - - @staticmethod - def dframe_to_frame(num): - assert num % 5 == 0, f"Invalid num: {num}" - return num // 5 * 17 - - @staticmethod - def save_sample(x, - save_path=None, - fps=8, - normalize=True, - value_range=(-1, 1), - force_video=False, - verbose=True): - """ - Args: - x (Tensor): shape [C, T, H, W] - """ - assert x.ndim == 4 - if not force_video and x.shape[1] == 1: # T = 1: save as image - save_path += ".png" - x = x.squeeze(1) - save_image([x], - save_path, - normalize=normalize, - value_range=value_range) - else: - save_path += ".mp4" - if normalize: - low, high = value_range - x.clamp_(min=low, max=high) - x.sub_(low).div_(max(high - low, 1e-5)) - x = x.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 3, 0).to( - "cpu", torch.uint8) - write_video(save_path, x, fps=fps, video_codec="h264") - if verbose: - print(f"Saved to {save_path}") - return save_path - - @staticmethod - def add_watermark( - input_video_path, - watermark_image_path="./assets/images/watermark/watermark.png", - output_video_path=None): - # execute this command in terminal with subprocess - # return if the process is successful - if output_video_path is None: - output_video_path = input_video_path.replace( - ".mp4", "_watermark.mp4") - cmd = f'ffmpeg -y -i {input_video_path} -i {watermark_image_path}' - cmd += f'-filter_complex "[1][0]scale2ref=oh*mdar:ih*0.1[logo][video];[video][logo]overlay" {output_video_path}' - import subprocess - exit_code = subprocess.call(cmd, shell=True) - is_success = exit_code == 0 - return is_success - - -class DataProcessor(): - - def __init__(self, text_encoder, vae): - self._text_encoder = text_encoder - self._vae = vae - - @staticmethod - def set_random_seed(seed: int): - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed) - return seed - - @staticmethod - def prepare_multi_resolution_info(model_type, batch_size, image_size, - num_frames, fps, device, dtype): - IMG_FPS = 120 - ret = {} - if model_type in ["STDiT3", "OpenSora"]: - fps = fps if num_frames > 1 else IMG_FPS - fps = torch.tensor([fps], device=device, - dtype=dtype).repeat(batch_size) - height = torch.tensor([image_size[0]], device=device, - dtype=dtype).repeat(batch_size) - width = torch.tensor([image_size[1]], device=device, - dtype=dtype).repeat(batch_size) - num_frames = torch.tensor([num_frames], device=device, - dtype=dtype).repeat(batch_size) - ar = torch.tensor([image_size[0] / image_size[1]], - device=device, - dtype=dtype).repeat(batch_size) - ret = dict(height=height, - width=width, - num_frames=num_frames, - ar=ar, - fps=fps) - else: - raise NotImplementedError(f"Model type is {model_type}") - return ret - - @staticmethod - def get_save_path_name( - save_dir, - sample_name=None, # prefix - sample_idx=None, # sample index - prompt=None, # used prompt - prompt_as_path=False, # use prompt as path - num_sample=1, # number of samples to generate for one prompt - k=None, # kth sample - ): - if sample_name is None: - sample_name = "" if prompt_as_path else "sample" - sample_name_suffix = prompt if prompt_as_path else f"_{sample_idx:04d}" - save_path = os.path.join(save_dir, f"{sample_name}{sample_name_suffix}") - if num_sample != 1: - save_path = f"{save_path}-{k}" - return save_path - - @staticmethod - def load_prompts(prompt_path, start_idx=None, end_idx=None): - return PromptProcessor.load_prompts(prompt_path=prompt_path, - start_idx=start_idx, - end_idx=end_idx) - - @staticmethod - def extract_json_from_prompts(prompts, reference, mask_strategy): - return PromptProcessor.extract_json_from_prompts( - prompts=prompts, reference=reference, mask_strategy=mask_strategy) - - @staticmethod - def split_prompt(prompt_text): - return PromptProcessor.split_prompt(prompt_text=prompt_text) - - @staticmethod - def merge_prompt(text_list, loop_idx_list=None): - return PromptProcessor.merge_prompt(text_list=text_list, - loop_idx_list=loop_idx_list) - - @staticmethod - def extract_prompts_loop(prompts, num_loop): - return PromptProcessor.extract_prompts_loop(prompts=prompts, - num_loop=num_loop) - - @staticmethod - def append_score_to_prompts(prompts, - aes=None, - flow=None, - camera_motion=None): - return PromptProcessor.append_score_to_prompts( - prompts=prompts, aes=aes, flow=flow, camera_motion=camera_motion) - - @staticmethod - def text_preprocessing(text, use_text_preprocessing: bool = True): - return PromptProcessor.text_preprocessing( - text=text, use_text_preprocessing=use_text_preprocessing) - - @staticmethod - def get_image_size(resolution, ar_ratio): - return VideoProcessor.get_image_size(resolution=resolution, - ar_ratio=ar_ratio) - - @staticmethod - def get_num_frames(num_frames): - return VideoProcessor.get_num_frames(num_frames=num_frames) - - def get_latent_size(self, input_size): - return VideoProcessor.get_latent_size(self._vae, input_size=input_size) - - def collect_references_batch(self, reference_paths, image_size): - return VideoProcessor.collect_references_batch( - self._vae, reference_paths=reference_paths, image_size=image_size) - - def append_generated(self, generated_video, refs_x, mask_strategy, loop_i, - condition_frame_length, condition_frame_edit): - return VideoProcessor.append_generated( - self._vae, - generated_video=generated_video, - refs_x=refs_x, - mask_strategy=mask_strategy, - loop_i=loop_i, - condition_frame_length=condition_frame_length, - condition_frame_edit=condition_frame_edit) - - @staticmethod - def apply_mask_strategy(z, refs_x, mask_strategys, loop_i, align=None): - return VideoProcessor.apply_mask_strategy(z=z, - refs_x=refs_x, - mask_strategys=mask_strategys, - loop_i=loop_i, - align=align) - - @staticmethod - def dframe_to_frame(num): - return VideoProcessor.dframe_to_frame(num=num) - - @staticmethod - def save_sample(x, - save_path=None, - fps=8, - normalize=True, - value_range=(-1, 1), - force_video=False, - verbose=True): - return VideoProcessor.save_sample(x=x, - save_path=save_path, - fps=fps, - normalize=normalize, - value_range=value_range, - force_video=force_video, - verbose=verbose) - - @staticmethod - def add_watermark( - input_video_path, - watermark_image_path="./assets/images/watermark/watermark.png", - output_video_path=None): - return VideoProcessor.add_watermark( - input_video_path=input_video_path, - watermark_image_path=watermark_image_path, - output_video_path=output_video_path) diff --git a/examples/models/contrib/stdit/vae.py b/examples/models/contrib/stdit/vae.py deleted file mode 100644 index cd193b8eb75a..000000000000 --- a/examples/models/contrib/stdit/vae.py +++ /dev/null @@ -1,802 +0,0 @@ -# Copyright 2024 HPC-AI Technology Inc. - -# 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. - -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# reference: https://github.com/hpcaitech/Open-Sora/blob/main/opensora/models/vae/vae.py - -import os -from typing import Tuple, Union - -import torch -import torch.nn.functional as F -from diffusers.models import AutoencoderKL -from einops import rearrange -from transformers import PretrainedConfig, PreTrainedModel -from utils import load_checkpoint - - -class VideoAutoencoderKL(torch.nn.Module): - - def __init__( - self, - from_pretrained=None, - micro_batch_size=None, - cache_dir=None, - local_files_only=False, - subfolder=None, - scaling_factor=0.18215, - ): - super().__init__() - self.module = AutoencoderKL.from_pretrained( - from_pretrained, - cache_dir=cache_dir, - local_files_only=local_files_only, - subfolder=subfolder, - ) - self.out_channels = self.module.config.latent_channels - self.patch_size = (1, 8, 8) - self.micro_batch_size = micro_batch_size - self.scaling_factor = scaling_factor - - def encode(self, x): - # x: (B, C, T, H, W) - B = x.shape[0] - x = rearrange(x, "B C T H W -> (B T) C H W") - - if self.micro_batch_size is None: - x = self.module.encode(x).latent_dist.sample().mul_( - self.scaling_factor) - else: - # NOTE: cannot be used for training - bs = self.micro_batch_size - x_out = [] - for i in range(0, x.shape[0], bs): - x_bs = x[i:i + bs] - x_bs = self.module.encode(x_bs).latent_dist.sample().mul_( - self.scaling_factor) - x_out.append(x_bs) - x = torch.cat(x_out, dim=0) - x = rearrange(x, "(B T) C H W -> B C T H W", B=B) - return x - - def decode(self, x, **kwargs): - # x: (B, C, T, H, W) - B = x.shape[0] - x = rearrange(x, "B C T H W -> (B T) C H W") - if self.micro_batch_size is None: - x = self.module.decode(x / self.scaling_factor).sample - else: - # NOTE: cannot be used for training - bs = self.micro_batch_size - x_out = [] - for i in range(0, x.shape[0], bs): - x_bs = x[i:i + bs] - x_bs = self.module.decode(x_bs / self.scaling_factor).sample - x_out.append(x_bs) - x = torch.cat(x_out, dim=0) - x = rearrange(x, "(B T) C H W -> B C T H W", B=B) - return x - - def get_latent_size(self, input_size): - latent_size = [] - for i in range(3): - # assert ( - # input_size[i] is None or input_size[i] % self.patch_size[i] == 0 - # ), "Input size must be divisible by patch size" - latent_size.append( - input_size[i] // - self.patch_size[i] if input_size[i] is not None else None) - return latent_size - - @property - def device(self): - return next(self.parameters()).device - - @property - def dtype(self): - return next(self.parameters()).dtype - - -def cast_tuple(t, length=1): - return t if isinstance(t, tuple) else ((t, ) * length) - - -def divisible_by(num, den): - return (num % den) == 0 - - -def is_odd(n): - return not divisible_by(n, 2) - - -def pad_at_dim(t, pad, dim=-1): - dims_from_right = (-dim - 1) if dim < 0 else (t.ndim - dim - 1) - zeros = (0, 0) * dims_from_right - return F.pad(t, (*zeros, *pad), mode="constant") - - -def exists(v): - return v is not None - - -class DiagonalGaussianDistribution: - """Stripped version of https://github.com/richzhang/PerceptualSimilarity/tree/master/models""" - - def __init__( - self, - parameters, - deterministic=False, - ): - self.parameters = parameters - self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) - self.logvar = torch.clamp(self.logvar, -30.0, 20.0) - self.deterministic = deterministic - self.std = torch.exp(0.5 * self.logvar) - self.var = torch.exp(self.logvar) - if self.deterministic: - self.var = self.std = torch.zeros_like(self.mean).to( - device=self.parameters.device, dtype=self.mean.dtype) - - def sample(self): - # torch.randn: standard normal distribution - x = self.mean + self.std * torch.randn(self.mean.shape).to( - device=self.parameters.device, dtype=self.mean.dtype) - return x - - def kl(self, other=None): - if self.deterministic: - return torch.Tensor([0.0]) - else: - if other is None: # SCH: assumes other is a standard normal distribution - return 0.5 * torch.sum( - torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, - dim=[1, 2, 3, 4]) - else: - return 0.5 * torch.sum( - torch.pow(self.mean - other.mean, 2) / other.var + - self.var / other.var - 1.0 - self.logvar + other.logvar, - dim=[1, 2, 3, 4], - ) - - def nll(self, sample, dims=[1, 2, 3, 4]): - if self.deterministic: - return torch.Tensor([0.0]) - logtwopi = torch.log(torch.Tensor([2.0 * torch.pi])) - return 0.5 * torch.sum(logtwopi + self.logvar + - torch.pow(sample - self.mean, 2) / self.var, - dim=dims) - - def mode(self): - return self.mean - - -class CausalConv3d(torch.nn.Module): - - def __init__( - self, - chan_in, - chan_out, - kernel_size: Union[int, Tuple[int, int, int]], - pad_mode="constant", - strides=None, # allow custom stride - **kwargs, - ): - super().__init__() - kernel_size = cast_tuple(kernel_size, 3) - time_kernel_size, height_kernel_size, width_kernel_size = kernel_size - assert is_odd(height_kernel_size) and is_odd(width_kernel_size) - dilation = kwargs.pop("dilation", 1) - stride = strides[0] if strides is not None else kwargs.pop("stride", 1) - self.pad_mode = pad_mode - time_pad = dilation * (time_kernel_size - 1) + (1 - stride) - height_pad = height_kernel_size // 2 - width_pad = width_kernel_size // 2 - self.time_pad = time_pad - self.time_causal_padding = (width_pad, width_pad, height_pad, - height_pad, time_pad, 0) - stride = strides if strides is not None else (stride, 1, 1) - dilation = (dilation, 1, 1) - self.conv = torch.nn.Conv3d(chan_in, - chan_out, - kernel_size, - stride=stride, - dilation=dilation, - **kwargs) - - def forward(self, x): - x = F.pad(x, self.time_causal_padding, mode=self.pad_mode) - x = self.conv(x) - return x - - -class ResBlock(torch.nn.Module): - - def __init__( - self, - in_channels, # SCH: added - filters, - conv_fn, - activation_fn=torch.nn.SiLU, - use_conv_shortcut=False, - num_groups=32, - ): - super().__init__() - self.in_channels = in_channels - self.filters = filters - self.activate = activation_fn() - self.use_conv_shortcut = use_conv_shortcut - - # SCH: MAGVIT uses GroupNorm by default - self.norm1 = torch.nn.GroupNorm(num_groups, in_channels) - self.conv1 = conv_fn(in_channels, - self.filters, - kernel_size=(3, 3, 3), - bias=False) - self.norm2 = torch.nn.GroupNorm(num_groups, self.filters) - self.conv2 = conv_fn(self.filters, - self.filters, - kernel_size=(3, 3, 3), - bias=False) - if in_channels != filters: - if self.use_conv_shortcut: - self.conv3 = conv_fn(in_channels, - self.filters, - kernel_size=(3, 3, 3), - bias=False) - else: - self.conv3 = conv_fn(in_channels, - self.filters, - kernel_size=(1, 1, 1), - bias=False) - - def forward(self, x): - residual = x - x = self.norm1(x) - x = self.activate(x) - x = self.conv1(x) - x = self.norm2(x) - x = self.activate(x) - x = self.conv2(x) - if self.in_channels != self.filters: # SCH: ResBlock X->Y - residual = self.conv3(residual) - return x + residual - - -def get_activation_fn(activation): - if activation == "relu": - activation_fn = torch.nn.ReLU - elif activation == "swish": - activation_fn = torch.nn.SiLU - else: - raise NotImplementedError - return activation_fn - - -class Encoder(torch.nn.Module): - """Encoder Blocks.""" - - def __init__( - self, - in_out_channels=4, - latent_embed_dim=512, # num channels for latent vector - filters=128, - num_res_blocks=4, - channel_multipliers=(1, 2, 2, 4), - temporal_downsample=(False, True, True), - num_groups=32, # for nn.GroupNorm - activation_fn="swish", - ): - super().__init__() - self.filters = filters - self.num_res_blocks = num_res_blocks - self.num_blocks = len(channel_multipliers) - self.channel_multipliers = channel_multipliers - self.temporal_downsample = temporal_downsample - self.num_groups = num_groups - self.embedding_dim = latent_embed_dim - - self.activation_fn = get_activation_fn(activation_fn) - self.activate = self.activation_fn() - self.conv_fn = CausalConv3d - self.block_args = dict( - conv_fn=self.conv_fn, - activation_fn=self.activation_fn, - use_conv_shortcut=False, - num_groups=self.num_groups, - ) - - # first layer conv - self.conv_in = self.conv_fn( - in_out_channels, - filters, - kernel_size=(3, 3, 3), - bias=False, - ) - - # ResBlocks and conv downsample - self.block_res_blocks = torch.nn.ModuleList([]) - self.conv_blocks = torch.nn.ModuleList([]) - - filters = self.filters - prev_filters = filters # record for in_channels - for i in range(self.num_blocks): - filters = self.filters * self.channel_multipliers[i] - block_items = torch.nn.ModuleList([]) - for _ in range(self.num_res_blocks): - block_items.append( - ResBlock(prev_filters, filters, **self.block_args)) - prev_filters = filters # update in_channels - self.block_res_blocks.append(block_items) - - if i < self.num_blocks - 1: - if self.temporal_downsample[i]: - t_stride = 2 if self.temporal_downsample[i] else 1 - s_stride = 1 - self.conv_blocks.append( - self.conv_fn(prev_filters, - filters, - kernel_size=(3, 3, 3), - strides=(t_stride, s_stride, s_stride))) - prev_filters = filters # update in_channels - else: - # if no t downsample, don't add since this does nothing for pipeline models - self.conv_blocks.append( - torch.nn.Identity(prev_filters)) # Identity - prev_filters = filters # update in_channels - - # last layer res block - self.res_blocks = torch.nn.ModuleList([]) - for _ in range(self.num_res_blocks): - self.res_blocks.append( - ResBlock(prev_filters, filters, **self.block_args)) - prev_filters = filters # update in_channels - - # MAGVIT uses Group Normalization - self.norm1 = torch.nn.GroupNorm(self.num_groups, prev_filters) - - self.conv2 = self.conv_fn(prev_filters, - self.embedding_dim, - kernel_size=(1, 1, 1), - padding="same") - - def forward(self, x): - x = self.conv_in(x) - - for i in range(self.num_blocks): - for j in range(self.num_res_blocks): - x = self.block_res_blocks[i][j](x) - if i < self.num_blocks - 1: - x = self.conv_blocks[i](x) - for i in range(self.num_res_blocks): - x = self.res_blocks[i](x) - - x = self.norm1(x) - x = self.activate(x) - x = self.conv2(x) - return x - - -class Decoder(torch.nn.Module): - """Decoder Blocks.""" - - def __init__( - self, - in_out_channels=4, - latent_embed_dim=512, - filters=128, - num_res_blocks=4, - channel_multipliers=(1, 2, 2, 4), - temporal_downsample=(False, True, True), - num_groups=32, # for nn.GroupNorm - activation_fn="swish", - ): - super().__init__() - self.filters = filters - self.num_res_blocks = num_res_blocks - self.num_blocks = len(channel_multipliers) - self.channel_multipliers = channel_multipliers - self.temporal_downsample = temporal_downsample - self.num_groups = num_groups - self.embedding_dim = latent_embed_dim - self.s_stride = 1 - - self.activation_fn = get_activation_fn(activation_fn) - self.activate = self.activation_fn() - self.conv_fn = CausalConv3d - self.block_args = dict( - conv_fn=self.conv_fn, - activation_fn=self.activation_fn, - use_conv_shortcut=False, - num_groups=self.num_groups, - ) - - filters = self.filters * self.channel_multipliers[-1] - prev_filters = filters - - # last conv - self.conv1 = self.conv_fn(self.embedding_dim, - filters, - kernel_size=(3, 3, 3), - bias=True) - - # last layer res block - self.res_blocks = torch.nn.ModuleList([]) - for _ in range(self.num_res_blocks): - self.res_blocks.append(ResBlock(filters, filters, - **self.block_args)) - - # ResBlocks and conv upsample - self.block_res_blocks = torch.nn.ModuleList([]) - self.num_blocks = len(self.channel_multipliers) - self.conv_blocks = torch.nn.ModuleList([]) - # reverse to keep track of the in_channels, but append also in a reverse direction - for i in reversed(range(self.num_blocks)): - filters = self.filters * self.channel_multipliers[i] - # resblock handling - block_items = torch.nn.ModuleList([]) - for _ in range(self.num_res_blocks): - block_items.append( - ResBlock(prev_filters, filters, **self.block_args)) - prev_filters = filters # SCH: update in_channels - self.block_res_blocks.insert(0, block_items) # SCH: append in front - - # conv blocks with upsampling - if i > 0: - if self.temporal_downsample[i - 1]: - t_stride = 2 if self.temporal_downsample[i - 1] else 1 - # SCH: T-Causal Conv 3x3x3, f -> (t_stride * 2 * 2) * f, depth to space t_stride x 2 x 2 - self.conv_blocks.insert( - 0, - self.conv_fn(prev_filters, - prev_filters * t_stride * self.s_stride * - self.s_stride, - kernel_size=(3, 3, 3)), - ) - else: - self.conv_blocks.insert( - 0, - torch.nn.Identity(prev_filters), - ) - - self.norm1 = torch.nn.GroupNorm(self.num_groups, prev_filters) - - self.conv_out = self.conv_fn(filters, in_out_channels, 3) - - def forward(self, x): - x = self.conv1(x) - for i in range(self.num_res_blocks): - x = self.res_blocks[i](x) - for i in reversed(range(self.num_blocks)): - for j in range(self.num_res_blocks): - x = self.block_res_blocks[i][j](x) - if i > 0: - t_stride = 2 if self.temporal_downsample[i - 1] else 1 - x = self.conv_blocks[i - 1](x) - x = rearrange( - x, - "B (C ts hs ws) T H W -> B C (T ts) (H hs) (W ws)", - ts=t_stride, - hs=self.s_stride, - ws=self.s_stride, - ) - - x = self.norm1(x) - x = self.activate(x) - x = self.conv_out(x) - return x - - -class VAE_Temporal(torch.nn.Module): - - def __init__( - self, - in_out_channels=4, - latent_embed_dim=4, - embed_dim=4, - filters=128, - num_res_blocks=4, - channel_multipliers=(1, 2, 2, 4), - temporal_downsample=(True, True, False), - num_groups=32, # for nn.GroupNorm - activation_fn="swish", - ): - super().__init__() - - self.time_downsample_factor = 2**sum(temporal_downsample) - # self.time_padding = self.time_downsample_factor - 1 - self.patch_size = (self.time_downsample_factor, 1, 1) - self.out_channels = in_out_channels - - # NOTE: following MAGVIT, conv in bias=False in encoder first conv - self.encoder = Encoder( - in_out_channels=in_out_channels, - latent_embed_dim=latent_embed_dim * 2, - filters=filters, - num_res_blocks=num_res_blocks, - channel_multipliers=channel_multipliers, - temporal_downsample=temporal_downsample, - num_groups=num_groups, # for nn.GroupNorm - activation_fn=activation_fn, - ) - self.quant_conv = CausalConv3d(2 * latent_embed_dim, 2 * embed_dim, 1) - - self.post_quant_conv = CausalConv3d(embed_dim, latent_embed_dim, 1) - self.decoder = Decoder( - in_out_channels=in_out_channels, - latent_embed_dim=latent_embed_dim, - filters=filters, - num_res_blocks=num_res_blocks, - channel_multipliers=channel_multipliers, - temporal_downsample=temporal_downsample, - num_groups=num_groups, # for nn.GroupNorm - activation_fn=activation_fn, - ) - - def get_latent_size(self, input_size): - latent_size = [] - for i in range(3): - if input_size[i] is None: - lsize = None - elif i == 0: - time_padding = (0 if - (input_size[i] % self.time_downsample_factor - == 0) else self.time_downsample_factor - - input_size[i] % self.time_downsample_factor) - lsize = (input_size[i] + time_padding) // self.patch_size[i] - else: - lsize = input_size[i] // self.patch_size[i] - latent_size.append(lsize) - return latent_size - - def encode(self, x): - time_padding = x.shape[2] % self.time_downsample_factor - if time_padding != 0: - time_padding = self.time_downsample_factor - time_padding - x = pad_at_dim(x, (time_padding, 0), dim=2) - encoded_feature = self.encoder(x) - moments = self.quant_conv(encoded_feature).to(x.dtype) - posterior = DiagonalGaussianDistribution(moments) - return posterior - - def decode(self, z, num_frames=None): - time_padding = num_frames % self.time_downsample_factor - if time_padding != 0: - time_padding = self.time_downsample_factor - time_padding - z = self.post_quant_conv(z) - x = self.decoder(z) - x = x[:, :, time_padding:] - return x - - def forward(self, x, sample_posterior=True): - posterior = self.encode(x) - if sample_posterior: - z = posterior.sample() - else: - z = posterior.mode() - recon_video = self.decode(z, num_frames=x.shape[2]) - return recon_video, posterior, z - - -VAE_MODELS = { - "VideoAutoencoderKL": VideoAutoencoderKL, - "VAE_Temporal_SD": VAE_Temporal, -} - - -class VideoAutoencoderPipelineConfig(PretrainedConfig): - model_type = "VideoAutoencoderPipeline" - - def __init__( - self, - spatial_vae_config=None, - temporal_vae_config=None, - from_pretrained=None, - freeze_vae_2d=False, - cal_loss=False, - micro_frame_size=None, - shift=0.0, - scale=1.0, - **kwargs, - ): - self.spatial_vae_config = spatial_vae_config - self.temporal_vae_config = temporal_vae_config - self.from_pretrained = from_pretrained - self.freeze_vae_2d = freeze_vae_2d - self.cal_loss = cal_loss - self.micro_frame_size = micro_frame_size - self.shift = shift - self.scale = scale - super().__init__(**kwargs) - - -class VideoAutoencoderPipeline(PreTrainedModel): - config_class = VideoAutoencoderPipelineConfig - - def __init__(self, config: VideoAutoencoderPipelineConfig): - super().__init__(config=config) - vae_type = config.spatial_vae_config.pop('type') - self.spatial_vae = VAE_MODELS[vae_type](**config.spatial_vae_config) - - vae_type = config.temporal_vae_config.pop('type') - pretrained_path = config.temporal_vae_config.pop( - 'from_pretrained' - ) if 'from_pretrained' in config.temporal_vae_config else None - self.temporal_vae = VAE_MODELS[vae_type](**config.temporal_vae_config) - if pretrained_path is not None: - load_checkpoint(self.temporal_vae, pretrained_path) - - self.cal_loss = config.cal_loss - self.micro_frame_size = config.micro_frame_size - self.micro_z_frame_size = self.temporal_vae.get_latent_size( - [config.micro_frame_size, None, None])[0] - if config.freeze_vae_2d: - for param in self.spatial_vae.parameters(): - param.requires_grad = False - self.out_channels = self.temporal_vae.out_channels - # normalization parameters - scale = torch.tensor(config.scale) - shift = torch.tensor(config.shift) - if len(scale.shape) > 0: - scale = scale[None, :, None, None, None] - if len(shift.shape) > 0: - shift = shift[None, :, None, None, None] - self.register_buffer("scale", scale) - self.register_buffer("shift", shift) - - def encode(self, x): - x_z = self.spatial_vae.encode(x) - - if self.micro_frame_size is None: - posterior = self.temporal_vae.encode(x_z) - z = posterior.sample() - else: - z_list = [] - for i in range(0, x_z.shape[2], self.micro_frame_size): - x_z_bs = x_z[:, :, i:i + self.micro_frame_size] - posterior = self.temporal_vae.encode(x_z_bs) - z_list.append(posterior.sample()) - z = torch.cat(z_list, dim=2) - - if self.cal_loss: - return z, posterior, x_z - else: - return (z - self.shift) / self.scale - - def decode(self, z, num_frames=None): - if not self.cal_loss: - z = z * self.scale.to(z.dtype) + self.shift.to(z.dtype) - - if self.micro_frame_size is None: - x_z = self.temporal_vae.decode(z, num_frames=num_frames) - x = self.spatial_vae.decode(x_z) - else: - x_z_list = [] - for i in range(0, z.size(2), self.micro_z_frame_size): - z_bs = z[:, :, i:i + self.micro_z_frame_size] - x_z_bs = self.temporal_vae.decode(z_bs, - num_frames=min( - self.micro_frame_size, - num_frames)) - x_z_list.append(x_z_bs) - num_frames -= self.micro_frame_size - x_z = torch.cat(x_z_list, dim=2) - x = self.spatial_vae.decode(x_z) - - if self.cal_loss: - return x, x_z - else: - return x - - def forward(self, x): - assert self.cal_loss, "This method is only available when cal_loss is True" - z, posterior, x_z = self.encode(x) - x_rec, x_z_rec = self.decode(z, num_frames=x_z.shape[2]) - return x_rec, x_z_rec, z, posterior, x_z - - def get_latent_size(self, input_size): - if self.micro_frame_size is None or input_size[0] is None: - return self.temporal_vae.get_latent_size( - self.spatial_vae.get_latent_size(input_size)) - else: - sub_input_size = [ - self.micro_frame_size, input_size[1], input_size[2] - ] - sub_latent_size = self.temporal_vae.get_latent_size( - self.spatial_vae.get_latent_size(sub_input_size)) - sub_latent_size[0] = sub_latent_size[0] * (input_size[0] // - self.micro_frame_size) - remain_temporal_size = [ - input_size[0] % self.micro_frame_size, None, None - ] - if remain_temporal_size[0] > 0: - remain_size = self.temporal_vae.get_latent_size( - remain_temporal_size) - sub_latent_size[0] += remain_size[0] - return sub_latent_size - - def get_temporal_last_layer(self): - return self.temporal_vae.decoder.conv_out.conv.weight - - @property - def device(self): - return next(self.parameters()).device - - @property - def dtype(self): - return next(self.parameters()).dtype - - -def get_vae( - micro_batch_size=4, - micro_frame_size=17, - from_pretrained=None, - local_files_only=False, - freeze_vae_2d=False, - cal_loss=False, - force_huggingface=False, -): - spatial_vae_config = dict( - type="VideoAutoencoderKL", - from_pretrained="PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers", - subfolder="vae", - micro_batch_size=micro_batch_size, - local_files_only=local_files_only, - ) - temporal_vae_config = dict( - type="VAE_Temporal_SD", - from_pretrained=None, - in_out_channels=4, - latent_embed_dim=4, - embed_dim=4, - filters=128, - num_res_blocks=4, - channel_multipliers=(1, 2, 2, 4), - temporal_downsample=(False, True, True), - ) - shift = (-0.10, 0.34, 0.27, 0.98) - scale = (3.85, 2.32, 2.33, 3.06) - kwargs = dict( - spatial_vae_config=spatial_vae_config, - temporal_vae_config=temporal_vae_config, - freeze_vae_2d=freeze_vae_2d, - cal_loss=cal_loss, - micro_frame_size=micro_frame_size, - shift=shift, - scale=scale, - ) - - if force_huggingface or (from_pretrained is not None - and not os.path.exists(from_pretrained)): - model = VideoAutoencoderPipeline.from_pretrained( - from_pretrained, **kwargs) - else: - config = VideoAutoencoderPipelineConfig(**kwargs) - model = VideoAutoencoderPipeline(config) - if from_pretrained: - load_checkpoint(model, from_pretrained) - return model diff --git a/examples/models/contrib/stdit/video_transforms.py b/examples/models/contrib/stdit/video_transforms.py deleted file mode 100644 index 3cabd640c5c6..000000000000 --- a/examples/models/contrib/stdit/video_transforms.py +++ /dev/null @@ -1,581 +0,0 @@ -# Copyright 2024 Vchitect/Latte - -# 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.# Modified from Latte - -# Copyright 2024 HPC-AI Technology Inc. - -# 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. - -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# reference: https://github.com/hpcaitech/Open-Sora/blob/main/opensora/datasets/video_transforms.py - -import numbers -import random - -import numpy as np -import torch - - -def _is_tensor_video_clip(clip): - if not torch.is_tensor(clip): - raise TypeError("clip should be Tensor. Got %s" % type(clip)) - - if not clip.ndimension() == 4: - raise ValueError("clip should be 4D. Got %dD" % clip.dim()) - - return True - - -def crop(clip, i, j, h, w): - """ - Args: - clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W) - """ - if len(clip.size()) != 4: - raise ValueError("clip should be a 4D tensor") - return clip[..., i:i + h, j:j + w] - - -def resize(clip, target_size, interpolation_mode): - if len(target_size) != 2: - raise ValueError( - f"target size should be tuple (height, width), instead got {target_size}" - ) - return torch.nn.functional.interpolate(clip, - size=target_size, - mode=interpolation_mode, - align_corners=False) - - -def resize_scale(clip, target_size, interpolation_mode): - if len(target_size) != 2: - raise ValueError( - f"target size should be tuple (height, width), instead got {target_size}" - ) - H, W = clip.size(-2), clip.size(-1) - scale_ = target_size[0] / min(H, W) - return torch.nn.functional.interpolate(clip, - scale_factor=scale_, - mode=interpolation_mode, - align_corners=False) - - -def resized_crop(clip, i, j, h, w, size, interpolation_mode="bilinear"): - """ - Do spatial cropping and resizing to the video clip - Args: - clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W) - i (int): i in (i,j) i.e coordinates of the upper left corner. - j (int): j in (i,j) i.e coordinates of the upper left corner. - h (int): Height of the cropped region. - w (int): Width of the cropped region. - size (tuple(int, int)): height and width of resized clip - Returns: - clip (torch.tensor): Resized and cropped clip. Size is (T, C, H, W) - """ - if not _is_tensor_video_clip(clip): - raise ValueError("clip should be a 4D torch.tensor") - clip = crop(clip, i, j, h, w) - clip = resize(clip, size, interpolation_mode) - return clip - - -def center_crop(clip, crop_size): - if not _is_tensor_video_clip(clip): - raise ValueError("clip should be a 4D torch.tensor") - h, w = clip.size(-2), clip.size(-1) - th, tw = crop_size - if h < th or w < tw: - raise ValueError("height and width must be no smaller than crop_size") - - i = int(round((h - th) / 2.0)) - j = int(round((w - tw) / 2.0)) - return crop(clip, i, j, th, tw) - - -def center_crop_using_short_edge(clip): - if not _is_tensor_video_clip(clip): - raise ValueError("clip should be a 4D torch.tensor") - h, w = clip.size(-2), clip.size(-1) - if h < w: - th, tw = h, h - i = 0 - j = int(round((w - tw) / 2.0)) - else: - th, tw = w, w - i = int(round((h - th) / 2.0)) - j = 0 - return crop(clip, i, j, th, tw) - - -def resize_crop_to_fill(clip, target_size): - if not _is_tensor_video_clip(clip): - raise ValueError("clip should be a 4D torch.tensor") - h, w = clip.size(-2), clip.size(-1) - th, tw = target_size[0], target_size[1] - rh, rw = th / h, tw / w - if rh > rw: - sh, sw = th, round(w * rh) - clip = resize(clip, (sh, sw), "bilinear") - i = 0 - j = int(round(sw - tw) / 2.0) - else: - sh, sw = round(h * rw), tw - clip = resize(clip, (sh, sw), "bilinear") - i = int(round(sh - th) / 2.0) - j = 0 - assert i + th <= clip.size(-2) and j + tw <= clip.size(-1) - return crop(clip, i, j, th, tw) - - -def random_shift_crop(clip): - """ - Slide along the long edge, with the short edge as crop size - """ - if not _is_tensor_video_clip(clip): - raise ValueError("clip should be a 4D torch.tensor") - h, w = clip.size(-2), clip.size(-1) - - if h <= w: - short_edge = h - else: - short_edge = w - - th, tw = short_edge, short_edge - - i = torch.randint(0, h - th + 1, size=(1, )).item() - j = torch.randint(0, w - tw + 1, size=(1, )).item() - return crop(clip, i, j, th, tw) - - -def to_tensor(clip): - """ - Convert tensor data type from uint8 to float, divide value by 255.0 and - permute the dimensions of clip tensor - Args: - clip (torch.tensor, dtype=torch.uint8): Size is (T, C, H, W) - Return: - clip (torch.tensor, dtype=torch.float): Size is (T, C, H, W) - """ - _is_tensor_video_clip(clip) - if not clip.dtype == torch.uint8: - raise TypeError("clip tensor should have data type uint8. Got %s" % - str(clip.dtype)) - # return clip.float().permute(3, 0, 1, 2) / 255.0 - return clip.float() / 255.0 - - -def normalize(clip, mean, std, inplace=False): - """ - Args: - clip (torch.tensor): Video clip to be normalized. Size is (T, C, H, W) - mean (tuple): pixel RGB mean. Size is (3) - std (tuple): pixel standard deviation. Size is (3) - Returns: - normalized clip (torch.tensor): Size is (T, C, H, W) - """ - if not _is_tensor_video_clip(clip): - raise ValueError("clip should be a 4D torch.tensor") - if not inplace: - clip = clip.clone() - mean = torch.as_tensor(mean, dtype=clip.dtype, device=clip.device) - # print(mean) - std = torch.as_tensor(std, dtype=clip.dtype, device=clip.device) - clip.sub_(mean[:, None, None, None]).div_(std[:, None, None, None]) - return clip - - -def hflip(clip): - """ - Args: - clip (torch.tensor): Video clip to be normalized. Size is (T, C, H, W) - Returns: - flipped clip (torch.tensor): Size is (T, C, H, W) - """ - if not _is_tensor_video_clip(clip): - raise ValueError("clip should be a 4D torch.tensor") - return clip.flip(-1) - - -class ResizeCrop: - - def __init__(self, size): - if isinstance(size, numbers.Number): - self.size = (int(size), int(size)) - else: - self.size = size - - def __call__(self, clip): - clip = resize_crop_to_fill(clip, self.size) - return clip - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(size={self.size})" - - -class RandomCropVideo: - - def __init__(self, size): - if isinstance(size, numbers.Number): - self.size = (int(size), int(size)) - else: - self.size = size - - def __call__(self, clip): - """ - Args: - clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W) - Returns: - torch.tensor: randomly cropped video clip. - size is (T, C, OH, OW) - """ - i, j, h, w = self.get_params(clip) - return crop(clip, i, j, h, w) - - def get_params(self, clip): - h, w = clip.shape[-2:] - th, tw = self.size - - if h < th or w < tw: - raise ValueError( - f"Required crop size {(th, tw)} is larger than input image size {(h, w)}" - ) - - if w == tw and h == th: - return 0, 0, h, w - - i = torch.randint(0, h - th + 1, size=(1, )).item() - j = torch.randint(0, w - tw + 1, size=(1, )).item() - - return i, j, th, tw - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(size={self.size})" - - -class CenterCropResizeVideo: - """ - First use the short side for cropping length, - center crop video, then resize to the specified size - """ - - def __init__( - self, - size, - interpolation_mode="bilinear", - ): - if isinstance(size, tuple): - if len(size) != 2: - raise ValueError( - f"size should be tuple (height, width), instead got {size}") - self.size = size - else: - self.size = (size, size) - - self.interpolation_mode = interpolation_mode - - def __call__(self, clip): - """ - Args: - clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W) - Returns: - torch.tensor: scale resized / center cropped video clip. - size is (T, C, crop_size, crop_size) - """ - clip_center_crop = center_crop_using_short_edge(clip) - clip_center_crop_resize = resize( - clip_center_crop, - target_size=self.size, - interpolation_mode=self.interpolation_mode) - return clip_center_crop_resize - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}" - - -class UCFCenterCropVideo: - """ - First scale to the specified size in equal proportion to the short edge, - then center cropping - """ - - def __init__( - self, - size, - interpolation_mode="bilinear", - ): - if isinstance(size, tuple): - if len(size) != 2: - raise ValueError( - f"size should be tuple (height, width), instead got {size}") - self.size = size - else: - self.size = (size, size) - - self.interpolation_mode = interpolation_mode - - def __call__(self, clip): - """ - Args: - clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W) - Returns: - torch.tensor: scale resized / center cropped video clip. - size is (T, C, crop_size, crop_size) - """ - clip_resize = resize_scale(clip=clip, - target_size=self.size, - interpolation_mode=self.interpolation_mode) - clip_center_crop = center_crop(clip_resize, self.size) - return clip_center_crop - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}" - - -class KineticsRandomCropResizeVideo: - """ - Slide along the long edge, with the short edge as crop size. And resie to the desired size. - """ - - def __init__( - self, - size, - interpolation_mode="bilinear", - ): - if isinstance(size, tuple): - if len(size) != 2: - raise ValueError( - f"size should be tuple (height, width), instead got {size}") - self.size = size - else: - self.size = (size, size) - - self.interpolation_mode = interpolation_mode - - def __call__(self, clip): - clip_random_crop = random_shift_crop(clip) - clip_resize = resize(clip_random_crop, self.size, - self.interpolation_mode) - return clip_resize - - -class CenterCropVideo: - - def __init__( - self, - size, - interpolation_mode="bilinear", - ): - if isinstance(size, tuple): - if len(size) != 2: - raise ValueError( - f"size should be tuple (height, width), instead got {size}") - self.size = size - else: - self.size = (size, size) - - self.interpolation_mode = interpolation_mode - - def __call__(self, clip): - """ - Args: - clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W) - Returns: - torch.tensor: center cropped video clip. - size is (T, C, crop_size, crop_size) - """ - clip_center_crop = center_crop(clip, self.size) - return clip_center_crop - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}" - - -class NormalizeVideo: - """ - Normalize the video clip by mean subtraction and division by standard deviation - Args: - mean (3-tuple): pixel RGB mean - std (3-tuple): pixel RGB standard deviation - inplace (boolean): whether do in-place normalization - """ - - def __init__(self, mean, std, inplace=False): - self.mean = mean - self.std = std - self.inplace = inplace - - def __call__(self, clip): - """ - Args: - clip (torch.tensor): video clip must be normalized. Size is (C, T, H, W) - """ - return normalize(clip, self.mean, self.std, self.inplace) - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(mean={self.mean}, std={self.std}, inplace={self.inplace})" - - -class ToTensorVideo: - """ - Convert tensor data type from uint8 to float, divide value by 255.0 and - permute the dimensions of clip tensor - """ - - def __init__(self): - pass - - def __call__(self, clip): - """ - Args: - clip (torch.tensor, dtype=torch.uint8): Size is (T, C, H, W) - Return: - clip (torch.tensor, dtype=torch.float): Size is (T, C, H, W) - """ - return to_tensor(clip) - - def __repr__(self) -> str: - return self.__class__.__name__ - - -class RandomHorizontalFlipVideo: - """ - Flip the video clip along the horizontal direction with a given probability - Args: - p (float): probability of the clip being flipped. Default value is 0.5 - """ - - def __init__(self, p=0.5): - self.p = p - - def __call__(self, clip): - """ - Args: - clip (torch.tensor): Size is (T, C, H, W) - Return: - clip (torch.tensor): Size is (T, C, H, W) - """ - if random.random() < self.p: - clip = hflip(clip) - return clip - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(p={self.p})" - - -# ------------------------------------------------------------ -# --------------------- Sampling --------------------------- -# ------------------------------------------------------------ -class TemporalRandomCrop(object): - """Temporally crop the given frame indices at a random location. - - Args: - size (int): Desired length of frames will be seen in the model. - """ - - def __init__(self, size): - self.size = size - - def __call__(self, total_frames): - rand_end = max(0, total_frames - self.size - 1) - begin_index = random.randint(0, rand_end) - end_index = min(begin_index + self.size, total_frames) - return begin_index, end_index - - -if __name__ == "__main__": - import os - - import numpy as np - import torchvision.io as io - from torchvision import transforms - from torchvision.utils import save_image - - vframes, aframes, info = io.read_video(filename="./v_Archery_g01_c03.avi", - pts_unit="sec", - output_format="TCHW") - - trans = transforms.Compose([ - ToTensorVideo(), - RandomHorizontalFlipVideo(), - UCFCenterCropVideo(512), - # NormalizeVideo(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), - transforms.Normalize(mean=[0.5, 0.5, 0.5], - std=[0.5, 0.5, 0.5], - inplace=True), - ]) - - target_video_len = 32 - frame_interval = 1 - total_frames = len(vframes) - print(total_frames) - - temporal_sample = TemporalRandomCrop(target_video_len * frame_interval) - - # Sampling video frames - start_frame_ind, end_frame_ind = temporal_sample(total_frames) - # print(start_frame_ind) - # print(end_frame_ind) - assert end_frame_ind - start_frame_ind >= target_video_len - frame_indice = np.linspace(start_frame_ind, - end_frame_ind - 1, - target_video_len, - dtype=int) - print(frame_indice) - - select_vframes = vframes[frame_indice] - print(select_vframes.shape) - print(select_vframes.dtype) - - select_vframes_trans = trans(select_vframes) - print(select_vframes_trans.shape) - print(select_vframes_trans.dtype) - - select_vframes_trans_int = ((select_vframes_trans * 0.5 + 0.5) * - 255).to(dtype=torch.uint8) - print(select_vframes_trans_int.dtype) - print(select_vframes_trans_int.permute(0, 2, 3, 1).shape) - - io.write_video("./test.avi", - select_vframes_trans_int.permute(0, 2, 3, 1), - fps=8) - - for i in range(target_video_len): - save_image(select_vframes_trans[i], - os.path.join("./test000", "%04d.png" % i), - normalize=True, - value_range=(-1, 1)) diff --git a/examples/models/core/bert/.gitignore b/examples/models/core/bert/.gitignore deleted file mode 100644 index 4b9ff316e843..000000000000 --- a/examples/models/core/bert/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -bert* -*.log diff --git a/examples/models/core/bert/README.md b/examples/models/core/bert/README.md deleted file mode 100644 index 05c9e16ca4a3..000000000000 --- a/examples/models/core/bert/README.md +++ /dev/null @@ -1,117 +0,0 @@ -# BERT and BERT Variants - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document explains how to build the BERT family, specifically [BERT](https://huggingface.co/docs/transformers/model_doc/bert) and [RoBERTa](https://huggingface.co/docs/transformers/model_doc/roberta) model using TensorRT-LLM. It also describes how to run on a single GPU and two GPUs. - -## Overview - -The TensorRT LLM BERT family implementation can be found in [`tensorrt_llm/models/bert/model.py`](../../../../tensorrt_llm/models/bert/model.py). -The TensorRT LLM BERT family example code is located in [`examples/models/core/bert`](./). There are two main files in that folder: - - * [`convert_checkpoint.py`](./convert_checkpoint.py) to convert the BERT model into TensorRT LLM checkpoint format. - * [`run.py`](./run.py) to run the inference on an input text, - -## Convert Weights - -The `convert_checkpoint.py` script converts weights from HuggingFace format to TRT-LLM format. You need to prepare HuggingFace checkpoint files before you run the convert script. - -Use `--model_dir` to specify the HuggingFace checkpoint directory. - -Supported `model_name` options include: BertModel, BertForQuestionAnswering, BertForSequenceClassification. Use `--model` to specify the target BERT model. Please note that if you choose BertModel, [`convert_checkpoint.py`](./convert_checkpoint.py) will ignore the BERT model class specified in HuggingFace config file, and throw a warning to remind you. - -Use `--output_dir` to specify the converted checkpoint and configuration directory. The default value is `./tllm_checkpoint`. This directory will be used for next engine building phase. - -Take BertForQuestionAnswering for example, - -```bash -export hf_model_dir= -export model_name='bertqa' -export model='BertForQuestionAnswering' -export dtype='float16' - -# convert -python convert_checkpoint.py \ ---model $model \ ---model_dir $hf_model_dir \ ---output_dir ${model_name}_${dtype}_tllm_checkpoint \ ---dtype $dtype - -# convert tp=2 -python convert_checkpoint.py \ ---model $model \ ---model_dir $hf_model_dir \ ---output_dir ${model_name}_${dtype}_tllm_checkpoint \ ---dtype $dtype \ ---tp_size 2 - -``` - -## Build TensorRT engine(s) - -TensorRT LLM converts HuggingFace BERT family models into TensorRT engine(s). -To build the TensorRT engine, the basic command is: - -```bash -trtllm-build --checkpoint_dir ./${model_name}_${dtype}_tllm_checkpoint \ ---output_dir ${model_name}_engine_outputs \ -``` -Beside the basic engine build, TensorRT LLM provides these features by adding these flags with basic build command: - -- To use `bert_attention_plugin`, add `--bert_attention_plugin` to command. - -- To use remove input padding, add `--remove_input_padding=enable` and `--bert_attention_plugin` to command. Please note that remove input padding feature has to come with bert attention plugin. - -- To use FMHA kernels, add `--context_fmha=enable` or `--bert_context_fmha_fp32_acc=enable`(to enable FP32 accumulation). Note that these two flags should be used with `--bert_attention_plugin` - -Continue the BertForQuestionAnswering example: -```bash -# Build TensorRT engine for BertForQuestionAnswering model, with remove_input_padding enabled. -# TP=1 and TP=2 share the same build command -trtllm-build --checkpoint_dir ./${model_name}_${dtype}_tllm_checkpoint \ ---output_dir=${model_name}_engine_outputs \ ---remove_input_padding=enable \ ---bert_attention_plugin=${dtype} \ ---max_batch_size 8 \ ---max_input_len 512 -``` - -## Run TensorRT engine(s) -Run a TensorRT LLM BERT model using the engines generated by build command mentioned above. -Note that during model deployment, only the TensorRT engine files are needed. Previously downloaded model checkpoints and converted weights can be removed. - -[`run.py`](./run.py) provides an example for performing the inference and decoding the output. By default, it will use the task specific datasets as input text, for example, ['squad_v2'](https://huggingface.co/datasets/rajpurkar/squad_v2) for BertForQuestionAnswering. - -To run the TensorRT engine, the basic command is: - -```bash -run.py --engine_dir ./${model_name}_engine_outputs \ ---hf_model_dir $hf_model_dir \ # used for loading tokenizer -``` -Please note that: - -- To use remove input padding, add `--remove_input_padding` to command. This flag is used to tell the runtime how to process the input and decode the output. - -- To compare the result with HuggingFace model, add `--run_hf_test` to command. The runtime will load the HF model from `hf_model_dir` and compare the result. Refer to [`run.py`](./run.py) for more details. - -Continue the BertForQuestionAnswering example: -```bash -# Run TensorRT engine -python run.py \ ---engine_dir ./${model_name}_engine_outputs \ ---hf_model_dir=$hf_model_dir \ ---remove_input_padding \ ---run_hf_test - -# Run TP=2 inference -mpirun -n 2 \ -python run.py \ ---engine_dir ./${model_name}_engine_outputs \ ---hf_model_dir=$hf_model_dir \ ---remove_input_padding \ ---run_hf_test -``` diff --git a/examples/models/core/bert/__init__.py b/examples/models/core/bert/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/examples/models/core/bert/convert_checkpoint.py b/examples/models/core/bert/convert_checkpoint.py deleted file mode 100644 index f17624e8334e..000000000000 --- a/examples/models/core/bert/convert_checkpoint.py +++ /dev/null @@ -1,192 +0,0 @@ -import argparse -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Union - -from transformers import AutoConfig - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import (BertForQuestionAnswering, - BertForSequenceClassification, BertModel, - RobertaForQuestionAnswering, - RobertaForSequenceClassification, RobertaModel) -from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model', - required=True, - choices=[ - 'BertModel', - 'BertForQuestionAnswering', - 'BertForSequenceClassification', - 'RobertaModel', - 'RobertaForQuestionAnswering', - 'RobertaForSequenceClassification', - ]) - parser.add_argument('--model_dir', type=str, default=None) - 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('--dtype', - type=str, - default='float16', - choices=['float32', 'float16']) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - # Quantization args - parser.add_argument("--use_fp8", - action="store_true", - default=False, - help="Enable FP8 per-tensor quantization") - parser.add_argument( - '--quant_ckpt_path', - type=str, - default=None, - help='Path of a quantized model checkpoint in .safetensors format') - parser.add_argument( - '--calib_dataset', - type=str, - default='ccdv/cnn_dailymail', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - - parser.add_argument('--log_level', type=str, default='info') - - args = parser.parse_args() - - return args - - -def args_to_quant_config(args: argparse.Namespace) -> QuantConfig: - '''return config dict with quantization info based on the command line args - ''' - quant_config = QuantConfig() - - if args.use_fp8: - quant_config.quant_algo = QuantAlgo.FP8 - return quant_config - - -def convert_and_save_hf(args): - model_dir = args.model_dir - - world_size = args.tp_size * args.pp_size - #TODO: add override_fields if needed - # Need to convert the cli args to the kay-value pairs and override them in the generate config dict. - # Ideally these fields will be moved out of the config and pass them into build API, keep them here for compatibility purpose for now, - # before the refactor is done. - - #TODO: add fp8 support later - quant_config = args_to_quant_config(args) - - hf_config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True) - assert hf_config is not None, "Failed to load huggingface config, please check!" - - def convert_and_save_rank(args, rank, tllm_class: Union[ - BertModel, - RobertaModel, - BertForQuestionAnswering, - RobertaForQuestionAnswering, - BertForSequenceClassification, - RobertaForSequenceClassification, - ]): - mapping = Mapping( - world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size, - ) - tik = time.time() - tllm_bert = tllm_class.from_hugging_face( - model_dir, - args.dtype, - mapping=mapping, - quant_config=quant_config, - ) - print(f'Total time of reading and converting {time.time()-tik} s') - tik = time.time() - tllm_bert.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del tllm_bert - print(f'Total time of saving checkpoint {time.time()-tik} s') - - tllm_class = globals()[f'{args.model}'] - if not args.model == hf_config.architectures[0]: - logger.warning( - "The model doesn't match the architecture in huggingface config.") - - execute(args.workers, [convert_and_save_rank] * world_size, args, - tllm_class) - release_gc() - - -def execute(workers, func, args, - tllm_class: Union[BertModel, RobertaModel, BertForQuestionAnswering, - RobertaForQuestionAnswering, - BertForSequenceClassification, - RobertaForSequenceClassification]): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank, tllm_class) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [ - p.submit(f, args, rank, tllm_class) - for rank, f in enumerate(func) - ] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - logger.set_level(args.log_level) - - assert ((args.tp_size <= 2) - and (args.pp_size == 1)), "For now we only support TP = 2!" - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - assert args.model_dir is not None - convert_and_save_hf(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/core/bert/run.py b/examples/models/core/bert/run.py deleted file mode 100644 index 3796d2a28337..000000000000 --- a/examples/models/core/bert/run.py +++ /dev/null @@ -1,338 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import argparse -import json -import os - -# isort: off -import torch -import tensorrt as trt -# isort: on - -from transformers import AutoConfig, AutoTokenizer -from utils import (compare_bertcls_result, compare_bertqa_result, - decode_bertcls_output, decode_bertqa_output, get_engine_name, - intermediate_check, prepare_text_inputs, process_input, - temporary_datasets_config) - -import tensorrt_llm -from tensorrt_llm import logger -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import trt_dtype_to_torch -from tensorrt_llm.runtime import Session, TensorInfo - -from transformers import BertConfig, BertPreTrainedModel, BertForQuestionAnswering, BertForSequenceClassification, BertModel # isort:skip -from transformers import RobertaConfig, RobertaPreTrainedModel, RobertaForQuestionAnswering, RobertaForSequenceClassification, RobertaModel # isort:skip - -OUTPUT_NAME_MAPPING = { - 'BertModel': 'hidden_states', - 'BertForQuestionAnswering': 'logits', - 'BertForSequenceClassification': 'logits', - 'RobertaModel': 'hidden_states', - 'RobertaForQuestionAnswering': 'logits', - 'RobertaForSequenceClassification': 'logits' -} - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument('--engine_dir', type=str, default='bert_outputs') - parser.add_argument('--hf_model_dir', type=str, required=True) - parser.add_argument('--run_hf_test', action='store_true') - parser.add_argument('--remove_input_padding', action='store_true') - parser.add_argument('--debug', action='store_true') - - return parser.parse_args() - - -if __name__ == '__main__': - emit_engine_arch_deprecation("run.py") - args = parse_arguments() - - tensorrt_llm.logger.set_level(args.log_level) - - config_path = os.path.join(args.engine_dir, 'config.json') - with open(config_path, 'r') as f: - config = json.load(f) - - remove_padding = config['build_config']['plugin_config'][ - 'remove_input_padding'] - assert args.remove_input_padding == remove_padding, \ - f"The engine is build with remove_input_padding={remove_padding}, \ - but the inference runtime is performed with remove_input_padding={args.remove_input_padding}!" - - world_size = config['pretrained_config']['mapping']['world_size'] - assert world_size == tensorrt_llm.mpi_world_size(), \ - f'Engine world size ({world_size}) != Runtime world size ({tensorrt_llm.mpi_world_size()})' - - model_name = config['pretrained_config']['architecture'] - - # Roberta doesn't have token_type_ids, use all zeros to replace - is_roberta = "Roberta" in model_name - - runtime_rank = tensorrt_llm.mpi_rank() if world_size > 1 else 0 - - 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 = get_engine_name(runtime_rank) - serialize_path = os.path.join(args.engine_dir, serialize_path) - - stream = torch.cuda.current_stream().cuda_stream - logger.info(f'Loading engine from {serialize_path}') - with open(serialize_path, 'rb') as f: - engine_buffer = f.read() - logger.info(f'Creating session from engine') - session = Session.from_serialized_engine(engine_buffer) - if args.debug: session._print_engine_info() - - #NOTE: prepare input - with temporary_datasets_config(HF_DATASETS_OFFLINE=False): - test_inputs = prepare_text_inputs(model_name) - hf_tokenizer = AutoTokenizer.from_pretrained(args.hf_model_dir) - if args.remove_input_padding: - #NOTE:Remove padding - inputs_without_padding = hf_tokenizer(**test_inputs) - input_ids_list = [ - torch.tensor(ids).int().cuda() \ - for ids in inputs_without_padding['input_ids'] - ] - # attention_mask_list = inputs_without_padding['attention_mask'], - if is_roberta: - token_type_ids_list = [ - torch.zeros_like(torch.tensor(ids)).int().cuda() \ - for ids in inputs_without_padding['input_ids'] - ] - else: - token_type_ids_list = [ - torch.tensor(ids).int().cuda() \ - for ids in inputs_without_padding['token_type_ids'] - ] - - input_ids, input_lengths, token_type_ids, position_ids, max_input_length = \ - process_input(input_ids_list=input_ids_list, - token_type_ids_list=token_type_ids_list, - is_roberta=is_roberta, - padding_idx=config['pretrained_config']['pad_token_id']) - - else: - - #NOTE:Padding: pad to longest seq len - inputs_with_padding = hf_tokenizer( - **test_inputs, - padding=True, - ) - - inputs_without_padding = hf_tokenizer(**test_inputs) - input_ids = torch.tensor(inputs_with_padding['input_ids']).int().cuda() - input_lengths = [len(x) for x in inputs_without_padding['input_ids']] - input_lengths = torch.tensor(input_lengths, - device=input_ids.device, - dtype=torch.int32) - attention_mask = torch.tensor(inputs_with_padding['attention_mask'], - device=input_ids.device, - dtype=torch.int32) - if is_roberta: - token_type_ids = torch.zeros_like(torch.tensor( - inputs_with_padding['input_ids']), - device=input_ids.device, - dtype=torch.int32) - else: - token_type_ids = torch.tensor(inputs_with_padding['token_type_ids'], - device=input_ids.device, - dtype=torch.int32) - - # NOTE: TRT-LLM perform inference - output_name = OUTPUT_NAME_MAPPING[model_name] - if args.remove_input_padding: - # NOTE: Remove padding: - inputs = { - "input_ids": input_ids, - "input_lengths": input_lengths, - "token_type_ids": token_type_ids, - "position_ids": position_ids, - "max_input_length": max_input_length - } - output_info = session.infer_shapes([ - TensorInfo("input_ids", trt.DataType.INT32, input_ids.shape), - TensorInfo("input_lengths", trt.DataType.INT32, - input_lengths.shape), - TensorInfo("token_type_ids", trt.DataType.INT32, - token_type_ids.shape), - TensorInfo("position_ids", trt.DataType.INT32, position_ids.shape), - TensorInfo("max_input_length", trt.DataType.INT32, - max_input_length.shape) - ]) - else: - #NOTE: Padding: - inputs = { - 'input_ids': input_ids, - 'input_lengths': input_lengths, - 'token_type_ids': token_type_ids, - } - output_info = session.infer_shapes([ - TensorInfo('input_ids', trt.DataType.INT32, input_ids.shape), - TensorInfo('input_lengths', trt.DataType.INT32, - input_lengths.shape), - TensorInfo('token_type_ids', trt.DataType.INT32, - token_type_ids.shape) - ]) - - outputs = { - t.name: - torch.empty(tuple(t.shape), - dtype=trt_dtype_to_torch(t.dtype), - device='cuda') - for t in output_info - } - assert output_name in outputs, f'{output_name} not found in outputs, check if build.py set output name correctly' - - logger.info(f"Rank{runtime_rank} is running inference...") - ok = session.run(inputs=inputs, outputs=outputs, stream=stream) - assert ok, "Runtime execution failed" - torch.cuda.synchronize() - res = outputs[output_name] - if args.debug: logger.info(f"Outputs:{outputs.keys()}") - - # NOTE: load hf model and perform inference as reference (only on rank0) - if tensorrt_llm.mpi_rank() == 0: - logger.info(f"Rank{runtime_rank} is generating HF reference...") - if args.run_hf_test: - hf_bert = globals()[f'{model_name}'].from_pretrained( - args.hf_model_dir).cuda().to(torch.float16).eval() - hf_inputs = hf_tokenizer(**test_inputs, - padding=True, - return_tensors="pt") - hf_inputs = hf_inputs.to(hf_bert.device) - with torch.no_grad(): - hf_outputs = hf_bert.forward(output_hidden_states=args.debug, - **hf_inputs) - - torch.cuda.synchronize() - # NOTE: Decode output (only on rank0) - - if tensorrt_llm.mpi_rank() == 0: - logger.info(f"Rank{runtime_rank} is comparing with HF reference...") - if model_name == "BertModel" or model_name == "RobertaModel": - if args.remove_input_padding: - # reshape result back to [batch_size, ...], - # and then "padding" so we could compare the tensor - from torch.nn.utils.rnn import pad_sequence - res = torch.split(res, input_lengths.tolist(), dim=0) - res = pad_sequence(list(res), batch_first=True, padding_value=0) - else: - # applied attention mask on trtllm res - attention_mask_tmp = attention_mask.unsqueeze(-1) - res = res * attention_mask_tmp - if args.run_hf_test: - ref = hf_outputs.last_hidden_state - ref = ref * hf_inputs['attention_mask'].unsqueeze(-1) - - if args.debug: - intermediate_check(outputs, hf_outputs['hidden_states'], - attention_mask_tmp, logger) - - if world_size == 1: - torch.testing.assert_close(actual=res.half(), - expected=ref, - rtol=1.5e-2, - atol=1.5e-2) - else: - # the arithmetic order of TP>1 is different from HF ref, which is always TP=1 for convenience. - torch.testing.assert_close(actual=res.half(), - expected=ref, - rtol=4e-2, - atol=2e-2) - print(f"{model_name} result is all close to HF reference!") - - if model_name == 'BertForQuestionAnswering' or model_name == 'RobertaForQuestionAnswering': - if args.remove_input_padding: - - # [num_tokens, 2] -> [num_tokens, 1] - res_start_logits, res_end_logits = torch.split(res, 1, -1) - - # reshape result back to [batch_size, ...] - res_start_logits = torch.split(res_start_logits, - input_lengths.tolist(), - dim=0) - res_start_logits = tuple(t.squeeze() for t in res_start_logits) - res_end_logits = torch.split(res_end_logits, - input_lengths.tolist(), - dim=0) - res_end_logits = tuple(t.squeeze() for t in res_end_logits) - - else: - #NOTE: Padding - # [B, Padding_len, 2] -> [B, Padding_len, 1] - res_start_logits, res_end_logits = torch.split(res, 1, -1) - # [B, Padding_len, 1] -> [B, Padding_len] - res_start_logits = res_start_logits.squeeze() - res_end_logits = res_end_logits.squeeze() - res_start_logits = res_start_logits * attention_mask - res_end_logits = res_end_logits * attention_mask - - res_start_logits = torch.split(res_start_logits, 1, dim=0) - res_start_logits = tuple(t.squeeze(0) for t in res_start_logits) - res_end_logits = torch.split(res_end_logits, 1, dim=0) - res_end_logits = tuple(t.squeeze(0) for t in res_end_logits) - - - decode_res = decode_bertqa_output( inputs_text=test_inputs, \ - hf_tokenizer=hf_tokenizer, start_logits=res_start_logits, \ - end_logits=res_end_logits) - - if args.run_hf_test: - ref_start_logits = hf_outputs.start_logits - ref_end_logits = hf_outputs.end_logits - # when we use_plugin and have real-data model_dir and input - # We do not need to care about the output of padding positions: - ref_start_logits = ref_start_logits * hf_inputs['attention_mask'] - ref_end_logits = ref_end_logits * hf_inputs['attention_mask'] - - decode_ref = decode_bertqa_output( inputs_text=test_inputs, \ - hf_tokenizer=hf_tokenizer, start_logits=ref_start_logits, \ - end_logits=ref_end_logits) - compare_bertqa_result(inputs_text=test_inputs, - res_answers=decode_res, - ref_answers=decode_ref) - - elif model_name == 'BertForSequenceClassification' or model_name == 'RobertaForSequenceClassification': - hf_config = AutoConfig.from_pretrained(args.hf_model_dir) - decode_res = decode_bertcls_output(logits=res, - hf_model_config=hf_config, - inputs_text=test_inputs) - - if args.run_hf_test: - ref = hf_outputs.logits - if world_size == 1: - torch.testing.assert_close(actual=res.half(), - expected=ref, - rtol=1.5e-2, - atol=1.5e-2) - else: - # the arithmetic order of TP>1 is different from HF ref, which is always TP=1 for convenience. - torch.testing.assert_close(actual=res.half(), - expected=ref, - rtol=4e-2, - atol=2e-2) - decode_ref = decode_bertcls_output(logits=hf_outputs.logits, - hf_model_config=hf_config, - inputs_text=test_inputs) - compare_bertcls_result(inputs_text=test_inputs, - res_answers=decode_res, - ref_answers=decode_res) diff --git a/examples/models/core/bert/utils.py b/examples/models/core/bert/utils.py deleted file mode 100644 index b3714e523e87..000000000000 --- a/examples/models/core/bert/utils.py +++ /dev/null @@ -1,217 +0,0 @@ -from contextlib import contextmanager -from typing import Dict, List, Tuple - -# isort: off -import torch -# isort: on - -import os -from pathlib import Path -from typing import Optional - -import datasets -from datasets import load_dataset - -from transformers import BertConfig, BertPreTrainedModel, BertForQuestionAnswering, BertForSequenceClassification, BertModel # isort:skip -from transformers import RobertaConfig, RobertaPreTrainedModel, RobertaForQuestionAnswering, RobertaForSequenceClassification, RobertaModel # isort:skip - - -# NOTE: This routine is copied from from tests/unittests/utils/llm_data.py -def llm_models_root(check=False) -> Optional[Path]: - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ.get("LLM_MODELS_ROOT")) - - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - - if check: - assert root.exists(), \ - "You shall set LLM_MODELS_ROOT env or be able to access /home/scratch.trt_llm_data_ci to run this test" - - return root if root.exists() else None - - -def llm_datasets_root() -> Path: - return llm_models_root(check=True) / "datasets" - - -def prepare_text_inputs(model_name, batch_size=8): - print( - f"HF_DATASETS_OFFLINE inside function: {datasets.config.HF_DATASETS_OFFLINE}" - ) - if model_name == "BertForQuestionAnswering" or model_name == "RobertaForQuestionAnswering": - squad_dataset_root = str(llm_datasets_root( - )) + "/" if datasets.config.HF_DATASETS_OFFLINE else "" - squad_dataset_path = squad_dataset_root + "squad_v2" - squad_dataset = load_dataset(squad_dataset_path, trust_remote_code=True) - val_dataset = squad_dataset["validation"] - samples = val_dataset.select(range(batch_size)) - - qa_real_test_inputs = { - 'text': samples["question"], - 'text_pair': samples["context"] - } - return qa_real_test_inputs - elif model_name == "BertForSequenceClassification" or model_name == "RobertaForSequenceClassification": - yelp_dataset_root = str(llm_datasets_root( - )) + "/" if datasets.config.HF_DATASETS_OFFLINE else "fancyzhx/" - yelp_dataset_path = yelp_dataset_root + "yelp_polarity" - yelp_dataset = load_dataset(yelp_dataset_path, trust_remote_code=True) - val_dataset = yelp_dataset["test"] - samples = val_dataset.select(range(batch_size)) - - seqcls_real_test_inputs = {'text': samples['text']} - return seqcls_real_test_inputs - elif model_name == "BertModel" or model_name == "RobertaModel": - #NOTE: For BertModel, it is used as an encoder, so we use dummy input here, - # you can choose whatevert you like, but the numerical accuracy might vary. - test_input = 'To be or not to be: that is the question' - input_strings = [test_input for _ in range(batch_size)] - base_real_test_inputs = {'text': input_strings} - return base_real_test_inputs - - else: - raise NotImplementedError(f"Unknown model {model_name}") - - -def get_engine_name(rank): - return 'rank{}.engine'.format(rank) - - -def decode_bertqa_output(inputs_text, hf_tokenizer, - start_logits: Tuple[torch.Tensor], - end_logits: Tuple[torch.Tensor]): - question, context = inputs_text['text'], inputs_text['text_pair'] - assert len(context) == len(question) - batch_size = len(context) - - # regenerate inputs_ids because it is flatten for remove_input_padding=True - inputs = hf_tokenizer(**inputs_text, padding=True, return_tensors='pt') - inputs_ids = inputs['input_ids'] - answer_start_index = [logit.argmax(dim=0) for logit in start_logits] - answer_end_index = [logit.argmax(dim=0) for logit in end_logits] - decode_answer = [] - for i in range(batch_size): - predict_answer_tokens = inputs_ids[ - i, answer_start_index[i]:answer_end_index[i] + 1] - predict_text = hf_tokenizer.decode(predict_answer_tokens, - skip_special_tokens=True) - decode_answer.append(predict_text) - return decode_answer - - -def compare_bertqa_result(inputs_text, res_answers, ref_answers): - from difflib import SequenceMatcher - question, context = inputs_text['text'], inputs_text['text_pair'] - assert len(res_answers) == len(ref_answers) - batch_size = len(res_answers) - for i in range(batch_size): - print(f"Context: {context[i]}\nQuestion: {question[i]}") - print(f"Ref Answer: {ref_answers[i]}") - print(f"Res Answer: {res_answers[i]}") - match_rate = SequenceMatcher(None, "\n".join(res_answers[i]), - "\n".join(ref_answers[i])).ratio() - assert match_rate > 0.95 - print( - f"TRT-LLM results match HF results with literal match rate {match_rate}" - ) - - -def decode_bertcls_output(logits: torch.Tensor, hf_model_config, inputs_text): - text = inputs_text['text'] - id2label = hf_model_config.id2label - class_ids = logits.argmax(dim=1) - decode_answer = [] - batch_size = len(text) - for i in range(batch_size): - predicted_class_id = class_ids[i].item() - predicted_label = id2label[predicted_class_id] - decode_answer.append(predicted_label) - return decode_answer - - -def compare_bertcls_result(inputs_text, res_answers, ref_answers): - from difflib import SequenceMatcher - text = inputs_text['text'] - batch_size = len(text) - for i in range(batch_size): - print(f"Context: {text[i]}") - print(f"Ref Label: {ref_answers[i]}") - print(f"Res Label: {res_answers[i]}") - match_rate = SequenceMatcher(None, "\n".join(res_answers[i]), - "\n".join(ref_answers[i])).ratio() - assert match_rate > 0.95 - print( - f"TRT-LLM results match HF results with literal match rate {match_rate}" - ) - - -def process_input(input_ids_list: List[torch.Tensor], - token_type_ids_list: List[torch.Tensor], - is_roberta=False, - padding_idx=1): - input_lengths = [] - position_ids_list = [] - max_input_length = 0 - for i, input_ids in enumerate(input_ids_list): - input_len = len(input_ids) - assert input_len == len(token_type_ids_list[i]), f"sample {i}: len(input_ids)={len(input_ids)}, " \ - f"len(token_type_ids)={len(token_type_ids_list[i])}, not equal" - input_lengths.append(input_len) - position_ids = torch.arange(0, input_len, dtype=torch.int32) - if is_roberta: - position_ids = position_ids + 1 + padding_idx - - position_ids_list.append(position_ids) - max_input_length = max(max_input_length, input_len) - - # [num_tokens] - input_ids = torch.concat(input_ids_list).int().cuda() - token_type_ids = torch.concat(token_type_ids_list).int().cuda() - position_ids = torch.concat(position_ids_list).int().cuda() - - input_lengths = torch.tensor(input_lengths).int().cuda() # [batch_size] - max_input_length = torch.empty((max_input_length, )).int().cuda() - return input_ids, input_lengths, token_type_ids, position_ids, max_input_length - - -def intermediate_check(tllm_inter: Dict, hf_ref: Tuple[torch.Tensor], attn_mask, - logger): - - def apply_mask(x): - return x * attn_mask - - # minus one because there is an embedding output - num_layers = len(hf_ref) - 1 - - res = tllm_inter['embedding_output'] - res = apply_mask(res) - ref = hf_ref[0] - ref = apply_mask(ref) - torch.testing.assert_close(actual=res, expected=ref, rtol=1e-2, atol=1e-2) - logger.debug("Embedding are all close") - - for i in range(num_layers - 1): - res = tllm_inter[f'layer_{i}_output'] - res = apply_mask(res) - ref = hf_ref[i + 1] - ref = apply_mask(ref) - is_close = torch.allclose(res, ref, rtol=1e-2, atol=1e-2) - logger.debug(f'BertEncoderLayer_{i}_output is close: {is_close}') - - -@contextmanager -def temporary_datasets_config(**kwargs): - # Save original settings - original_settings = {} - for key, value in kwargs.items(): - original_settings[key] = getattr(datasets.config, key) - setattr(datasets.config, key, value) - try: - yield - finally: - # Restore original settings - for key, value in original_settings.items(): - setattr(datasets.config, key, value) diff --git a/examples/models/core/commandr/README.md b/examples/models/core/commandr/README.md deleted file mode 100644 index 12c978d15a0a..000000000000 --- a/examples/models/core/commandr/README.md +++ /dev/null @@ -1,213 +0,0 @@ -# Command R - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document explains how to build the [C4AI Command-R](https://huggingface.co/CohereForAI/c4ai-command-r-v01), [C4AI Command R+](https://huggingface.co/CohereForAI/c4ai-command-r-plus), [Aya-23-8B](https://huggingface.co/CohereForAI/aya-23-8B), [Aya-23-35B](https://huggingface.co/CohereForAI/aya-23-35B) models using TensorRT LLM and run on a single GPU or a single node with multiple GPUs. - -- [Command R](#Command-R) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [1. Download repo and weights from HuggingFace Transformers](#1-download-repo-and-weights-from-huggingface-transformers) - - [2. Convert weights from HF Transformers to TensorRT LLM format](#2-convert-weights-from-hf-transformers-to-tensorrt-llm-format) - - [3. Build TensorRT engine(s)](#3-build-tensorrt-engines) - - [4. Run inference](#4-run-inference) - - [Single node, single GPU](#single-node-single-gpu) - - [Single node, multi GPU](#single-node-multi-gpu) - - [5. Run summarization task](#5-run-summarization-task) - - [Weight Only quantization](#weight-only-quantization) - - -## Overview - -The TensorRT LLM Command-R implementation can be found in [`tensorrt_llm/models/commandr/model.py`](../../../../tensorrt_llm/models/commandr/model.py). -The TensorRT LLM Command-R example code is located in [`examples/models/core/commandr`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`run.py`](../../../run.py) to run the inference on an input text; -* [`summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - - * FP16 - * INT8 & INT4 Weight-Only - * Tensor Parallel - -## Usage - -The next section describe how to build the engine and run the inference demo. - -### 1. Download repo and weights from HuggingFace Transformers - -```bash -pip install -r requirements.txt -apt-get update -apt-get install git-lfs - -# clone one or more models we want to build -git clone https://huggingface.co/CohereForAI/c4ai-command-r-v01 command_r_v01 -git clone https://huggingface.co/CohereForAI/c4ai-command-r-plus command_r_plus -git clone https://huggingface.co/CohereForAI/aya-23-8B aya_23_8B -git clone https://huggingface.co/CohereForAI/aya-23-35B aya_23_35B -``` - -### 2. Convert weights from HF Transformers to TensorRT LLM format - -The [`convert_checkpoint.py`](./convert_checkpoint.py) script converts HF weights to TensorRT LLM checkpoints. The number of checkpoint files (in .safetensors format) is same to the number of GPUs used to run inference. - -```bash -# Command-R: single gpu, dtype float16 -python3 convert_checkpoint.py --model_dir command_r_v01 --output_dir trt_ckpt/command_r_v01/fp16/1-gpu - -# Command-R+: 4-way tensor parallelism -python3 convert_checkpoint.py --model_dir command_r_plus --tp_size 4 --output_dir trt_ckpt/command_r_plus/fp16/4-gpu - -# Aya-23-8B: single gpu, dtype float16 -python3 convert_checkpoint.py --model_dir aya_23_8B --output_dir trt_ckpt/aya_23_8B/fp16/1-gpu - -# Aya-23-35B: single gpu, dtype float16 -python3 convert_checkpoint.py --model_dir aya_23_35B --output_dir trt_ckpt/aya_23_35B/fp16/1-gpu -``` - -### 3. Build TensorRT engine(s) - -The `trtllm-build` command builds TensorRT LLM engines from TensorRT LLM checkpoints. The number of engine files is also same to the number of GPUs used to run inference. - -Normally, the `trtllm-build` command only requires a single GPU, but you can enable parallel building by passing the number of GPUs to the `--workers` argument. - -```bash -# Command-R: single-gpu engine with dtype float16, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir trt_ckpt/command_r_v01/fp16/1-gpu \ - --gemm_plugin float16 \ - --output_dir trt_engines/command_r_v01/fp16/1-gpu - -# Command-R+: 4-way tensor parallelism -trtllm-build --checkpoint_dir trt_ckpt/command_r_plus/fp16/4-gpu \ - --gemm_plugin float16 \ - --output_dir trt_engines/command_r_plus/fp16/4-gpu - -# Command-R: single-gpu engine with dtype float16, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir trt_ckpt/aya_23_8B/fp16/1-gpu \ - --gemm_plugin float16 \ - --output_dir trt_engines/aya_23_8B/fp16/1-gpu - -# Command-R: single-gpu engine with dtype float16, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir trt_ckpt/aya_23_35B/fp16/1-gpu \ - --gemm_plugin float16 \ - --output_dir trt_engines/aya_23_35B/fp16/1-gpu -``` - -If the engines are built successfully, you will see output like (Command-R as the example): - -```txt -...... -[09/19/2024-03:34:30] [TRT] [I] Engine generation completed in 26.9495 seconds. -[09/19/2024-03:34:30] [TRT] [I] [MemUsageStats] Peak memory usage of TRT CPU/GPU memory allocators: CPU 4 MiB, GPU 70725 MiB -[09/19/2024-03:34:55] [TRT] [I] [MemUsageStats] Peak memory usage during Engine building and serialization: CPU: 176260 MiB -[09/19/2024-03:34:55] [TRT-LLM] [I] Total time of building Unnamed Network 0: 00:00:52 -[09/19/2024-03:34:55] [TRT] [I] Serialized 26 bytes of code generator cache. -[09/19/2024-03:34:55] [TRT] [I] Serialized 315007 bytes of compilation cache. -[09/19/2024-03:34:55] [TRT] [I] Serialized 12 timing cache entries -[09/19/2024-03:34:55] [TRT-LLM] [I] Timing cache serialized to model.cache -[09/19/2024-03:34:55] [TRT-LLM] [I] Build phase peak memory: 176257.29 MB, children: 17.65 MB -[09/19/2024-03:34:55] [TRT-LLM] [I] Serializing engine to trt_engines/command_r_v01/fp16/1-gpu/rank0.engine... -[09/19/2024-03:35:20] [TRT-LLM] [I] Engine serialized. Total time: 00:00:25 -[09/19/2024-03:35:23] [TRT-LLM] [I] Total time of building all engines: 00:01:47 -``` - -### 4. Run inference - -#### Single node, single GPU - -```bash -# Run the default engine of Command-R on single GPU. -python3 ../../../run.py --max_output_len 50 \ - --tokenizer_dir command_r_v01 \ - --engine_dir trt_engines/command_r_v01/fp16/1-gpu - -# Run the default engine of Command-R on single GPU, using streaming output. -python3 ../../../run.py --max_output_len 50 \ - --tokenizer_dir command_r_v01 \ - --engine_dir trt_engines/command_r_v01/fp16/1-gpu \ - --streaming - -# Run the default engine of Aya-23-8B on single GPU. -python3 ../../../run.py --max_output_len 50 \ - --tokenizer_dir aya_23_8B \ - --engine_dir trt_engines/aya_23_8B/fp16/1-gpu - -# Run the default engine of Aya-23-35B on single GPU. -python3 ../../../run.py --max_output_len 50 \ - --tokenizer_dir aya_23_35B \ - --engine_dir trt_engines/aya_23_35B/fp16/1-gpu -``` - -#### Single node, multi GPU - -```bash -# Run the Tensor Parallel 4 engine of Command-R+ on 4 GPUs. -mpirun -n 4 \ - python ../../../run.py --max_output_len 50 \ - --tokenizer_dir command_r_plus \ - --engine_dir trt_engines/command_r_plus/fp16/4-gpu -``` - -If the engines are run successfully, you will see output like (Command-R as the example): - -```txt -...... -Input [Text 0]: "Born in north-east France, Soyer trained as a" -Output [Text 0 Beam 0]: " chef in Paris and worked in the kitchens of the French royal family. He came to England in 1814 and worked in a number of London hotels and restaurants, including the Reform Club and the London Tavern. He also opened his own restaurant" -``` - -### 5. Run summarization task - -```bash -# Run the summarization of Command-R task. -python3 ../../../summarize.py --test_trt_llm \ - --hf_model_dir command_r_v01 \ - --engine_dir trt_engines/command_r_v01/fp16/1-gpu -``` - -If the engines are run successfully, you will see output like (Command-R as the example): - -```txt -...... -[01/26/2024-02:51:56] [TRT-LLM] [I] TensorRT LLM (total latency: 81.05689692497253 sec) -[01/26/2024-02:51:56] [TRT-LLM] [I] TensorRT LLM (total output tokens: 2000) -[01/26/2024-02:51:56] [TRT-LLM] [I] TensorRT LLM (tokens per second: 24.67402621952367) -[01/26/2024-02:51:56] [TRT-LLM] [I] TensorRT LLM beam 0 result -[01/26/2024-02:51:56] [TRT-LLM] [I] rouge1 : 24.06804397902119 -[01/26/2024-02:51:56] [TRT-LLM] [I] rouge2 : 6.456513335555016 -[01/26/2024-02:51:56] [TRT-LLM] [I] rougeL : 16.77644999660741 -[01/26/2024-02:51:56] [TRT-LLM] [I] rougeLsum : 20.57359472317834 -``` - -### Weight Only quantization - -Use `--use_weight_only` to enable INT8-Weight-Only quantization, this will significantly lower the latency and memory footprint. Furthermore, use `--weight_only_precision int8` or `--weight_only_precision int4` to configure the data type of the weights. - -```bash -# Command-R: single gpu, int8 weight only quantization -python3 convert_checkpoint.py --model_dir command_r_v01 \ - --use_weight_only \ - --weight_only_precision int8 \ - --output_dir trt_ckpt/command_r_v01/int8_wo/1-gpu - -# Command-R: single-gpu engine with int8 weight only quantization, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir trt_ckpt/command_r_v01/int8_wo/1-gpu \ - --gemm_plugin float16 \ - --output_dir trt_engines/command_r_v01/int8_wo/1-gpu - -# Run inference. -python3 ../../../run.py --max_output_len 50 \ - --tokenizer_dir command_r_v01 \ - --engine_dir trt_engines/command_r_v01/int8_wo/1-gpu -``` diff --git a/examples/models/core/commandr/convert_checkpoint.py b/examples/models/core/commandr/convert_checkpoint.py deleted file mode 100644 index 666803143e02..000000000000 --- a/examples/models/core/commandr/convert_checkpoint.py +++ /dev/null @@ -1,187 +0,0 @@ -import argparse -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import CohereForCausalLM -from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - 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( - '--dtype', - type=str, - default='auto', - choices=['auto', 'float16', 'bfloat16', 'float32'], - help= - "The data type for the model weights and activations if not quantized. " - "If 'auto', the data type is automatically inferred from the source model; " - "however, if the source dtype is float32, it is converted to float16.") - 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( - '--disable_weight_only_quant_plugin', - default=False, - action="store_true", - help= - 'By default, using plugin implementation for weight quantization. Enabling disable_weight_only_quant_plugin flag will use ootb implementation instead of plugin.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - 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("--load_model_on_cpu", action="store_true") - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument('--log_level', type=str, default='info') - - args = parser.parse_args() - - return args - - -def args_to_quant_config(args: argparse.Namespace) -> QuantConfig: - '''return config dict with quantization info based on the command line args - ''' - quant_config = QuantConfig() - if args.use_weight_only: - if args.weight_only_precision == 'int8': - quant_config.quant_algo = QuantAlgo.W8A16 - elif args.weight_only_precision == 'int4': - quant_config.quant_algo = QuantAlgo.W4A16 - - return quant_config - - -def args_to_build_options(args): - return { - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'disable_weight_only_quant_plugin': - args.disable_weight_only_quant_plugin, - 'load_model_on_cpu': args.load_model_on_cpu, - } - - -def convert_and_save_hf(args): - model_dir = args.model_dir - world_size = args.tp_size * args.pp_size - # Need to convert the cli args to the kay-value pairs and override them in the generate config dict. - # Ideally these fields will be moved out of the config and pass them into build API, keep them here for compatibility purpose for now, - # before the refactor is done. - override_fields = {} - override_fields.update(args_to_build_options(args)) - - quant_config = args_to_quant_config(args) - - def convert_and_save_rank(args, rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - cohere = CohereForCausalLM.from_hugging_face( - model_dir, - args.dtype, - mapping=mapping, - quant_config=quant_config, - **override_fields, - ) - cohere.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del cohere - - execute(args.workers, [convert_and_save_rank] * world_size, args) - release_gc() - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - logger.set_level(args.log_level) - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - assert args.model_dir is not None - convert_and_save_hf(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/core/commandr/requirements.txt b/examples/models/core/commandr/requirements.txt deleted file mode 100644 index 2f8713d8658a..000000000000 --- a/examples/models/core/commandr/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score diff --git a/examples/models/core/enc_dec/README.md b/examples/models/core/enc_dec/README.md deleted file mode 100644 index 92ba03cd04c2..000000000000 --- a/examples/models/core/enc_dec/README.md +++ /dev/null @@ -1,444 +0,0 @@ -# Encoder-Decoder - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run an Encoder-Decoder (Enc-Dec) model in TensorRT LLM on NVIDIA GPUs. - -## Table of Contents - -- [Encoder-Decoder](#encoder-decoder) - - [Table of Contents](#table-of-contents) - - [Overview](#overview) - - [Usage](#usage) - - [Encoder-Decoder Model Support](#encoder-decoder-model-support) - - [Download weights from HuggingFace Transformers](#download-weights-from-huggingface-transformers) - - [Convert and Split Weights](#convert-and-split-weights) - - [Build TensorRT engine(s)](#build-tensorrt-engines) - - [Run](#run) - - [Run C++ runtime](#run-c-runtime) - - [Run with Triton Backend](#run-with-triton-backend) - - [Run Python runtime](#run-python-runtime) - - [Benchmark](#benchmark) - - [Benchmark C++ runtime](#benchmark-c-runtime) - - [Run BART with LoRA](#run-bart-with-lora) - - [Reminders](#reminders) - - [Attention Scaling Factors](#attention-scaling-factors) - - [Run FairSeq NMT (Neural Machine Translation) models](#run-fairseq-nmt-neural-machine-translation-models) - - [FP8 Post-Training Quantization](#fp8-post-training-quantization) - - [Get quantized checkpoint with ModelOpt](#get-quantized-checkpoint-with-modelopt) - -## 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/models/core/enc_dec`](./): - - * `trtllm-build` 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 example input text. - * Enc-Dec models can have specific implementations, such as the popular T5 family (T5, mT5, Flan-T5, ByT5), BART family (BART, mBART), and FairSeq family (WMTs). They are now merged into a single convert script: - * [`convert_checkpoint.py`](./convert_checkpoint.py) to convert weights from HuggingFace or FairSeq format to TRT-LLM format, and split weights for multi-GPU inference, -## Usage - -The TensorRT LLM Enc-Dec example code locates at [examples/models/core/enc_dec](./). It takes HuggingFace or FairSeq 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](https://huggingface.co/docs/transformers/model_doc/mt5) -- [BART](https://huggingface.co/docs/transformers/model_doc/bart) -- [mBART](https://huggingface.co/docs/transformers/model_doc/mbart) -- [FairSeq NMT](https://pytorch.org/hub/pytorch_fairseq_translation/) -- [ByT5](https://huggingface.co/docs/transformers/main/en/model_doc/byt5) -- [UL2 (coming)](https://huggingface.co/docs/transformers/model_doc/ul2) and [Flan-UL2 (coming)](https://huggingface.co/docs/transformers/model_doc/flan-ul2) - -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. BART models and FairSeq models can follow very similar steps by just replacing the model name. - -### 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 -git clone https://huggingface.co/facebook/bart-large-cnn tmp/hf_models/bart-large-cnn -git clone https://huggingface.co/facebook/mbart-large-50-many-to-one-mmt tmp/hf_models/mbart-large-50-many-to-one-mmt -git clone https://huggingface.co/google/byt5-small tmp/hf_models/byt5-small -``` - -### Convert and Split Weights -The `convert_checkpoint.py` script converts weights from HuggingFace or FairSeq format to TRT-LLM format, and splits weights for multi-GPU inference. `--tp_size` specifies the number of GPUs for tensor parallelism during inference. Pipeline Parallelism size can be set with `--pp_size` for distributed inference. - -The HuggingFace or Fairseq checkpoints of the enc-dec models mentioned in this Readme are all float32 precision. Use `--dtype` to set the target inference precision during the weight conversion. - -After weight conversion, TensorRT LLM converted weights and model configuration will be saved under `/` directory, which is the `--checkpoint_dir` input path you should give to the **next** engine building phase. - -Take T5 for example: - -```bash -# Example: build 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=1. -export MODEL_NAME="t5-small" # or "flan-t5-small" -export MODEL_TYPE="t5" -export INFERENCE_PRECISION="bfloat16" -export TP_SIZE=4 -export PP_SIZE=1 -export WORLD_SIZE=4 -export MAX_BEAM_WIDTH=1 -python convert_checkpoint.py --model_type ${MODEL_TYPE} \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --output_dir tmp/trt_models/${MODEL_NAME}/${INFERENCE_PRECISION} \ - --tp_size ${TP_SIZE} \ - --pp_size ${PP_SIZE} \ - --dtype ${INFERENCE_PRECISION} -``` - -### Build TensorRT engine(s) - -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. - -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`. - -The default value of `--max_input_len` is 1024. When building DecoderModel, specify decoder input length with `--max_input_len=1` for encoder-decoder model to start generation from decoder_start_token_id of length 1. If the start token is a single token (the default behavior of T5/BART/etc.), you should set `--max_input_len` as 1; if you want the decoder-only type of generation, set `--max_input_len` above 1 to get similar behavior as HF's `decoder_forced_input_ids`. - -EncoderModel does not generate prompt. `--max_seq_len` should be the same as `--max_input_len`. `--max_seq_len` would be set as `--max_input_len` if not specified. - -DecoderModel takes `--max_encoder_input_len` and `--max_input_len` as model inputs, `--max_encoder_input_len` is set to 1024 as default since `--max_input_len` is 1024 for EncoderModel. - -To be noted: -1. For T5, add `--context_fmha disable`. FMHA with T5's relative attention bias is not implemented. Add `--use_implicit_relative_attention` when `--max_seq_len` is extremely large, causing decoder engine size to be too large to fit in memory. Compute relative attention on-the-fly (implicitly, without pre-computation) instead. -2. `--bert_attention_plugin`, `--gpt_attention_plugin`, `--remove_input_padding`, `--gemm_plugin` require explicit disabling and setting, or else they'll be set to default value in `trtllm-build`. - -```bash -# --gpt_attention_plugin is necessary in Enc-Dec. -# Try --gemm_plugin to prevent accuracy issue. -# It is recommended to use --remove_input_padding along with --gpt_attention_plugin for better performance -trtllm-build --checkpoint_dir tmp/trt_models/${MODEL_NAME}/${INFERENCE_PRECISION}/encoder \ - --output_dir tmp/trt_engines/${MODEL_NAME}/${INFERENCE_PRECISION}/encoder \ - --paged_kv_cache disable \ - --moe_plugin disable \ - --max_beam_width ${MAX_BEAM_WIDTH} \ - --max_batch_size 8 \ - --max_input_len 1024 \ - --gemm_plugin ${INFERENCE_PRECISION} \ - --bert_attention_plugin ${INFERENCE_PRECISION} \ - --gpt_attention_plugin ${INFERENCE_PRECISION} \ - --remove_input_padding enable \ - --context_fmha disable - -# For decoder, refer to the above content and set --max_input_len correctly -trtllm-build --checkpoint_dir tmp/trt_models/${MODEL_NAME}/${INFERENCE_PRECISION}/decoder \ - --output_dir tmp/trt_engines/${MODEL_NAME}/${INFERENCE_PRECISION}/decoder \ - --moe_plugin disable \ - --max_beam_width ${MAX_BEAM_WIDTH} \ - --max_batch_size 8 \ - --max_input_len 1 \ - --max_seq_len 201 \ - --max_encoder_input_len 1024 \ - --gemm_plugin ${INFERENCE_PRECISION} \ - --bert_attention_plugin ${INFERENCE_PRECISION} \ - --gpt_attention_plugin ${INFERENCE_PRECISION} \ - --remove_input_padding enable \ - --context_fmha disable - -``` - -For BART, `--context_fmha` can be enabled. `trtllm-build` has the default setting to enable it. - -```bash -# Example: build bart-large-cnn using a single GPU, FP32, running greedy search -export MODEL_NAME="bart-large-cnn" # or "mbart-large-50-many-to-one-mmt" -export MODEL_TYPE="bart" -export INFERENCE_PRECISION="float32" -export TP_SIZE=1 -export PP_SIZE=1 -export WORLD_SIZE=1 -export MAX_BEAM_WIDTH=1 -python convert_checkpoint.py --model_type ${MODEL_TYPE} \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --output_dir tmp/trt_models/${MODEL_NAME}/${INFERENCE_PRECISION} \ - --tp_size ${TP_SIZE} \ - --pp_size ${PP_SIZE} \ - --dtype ${INFERENCE_PRECISION} - -# Note: non-T5 models can enable FMHA for the encoder part, for FP16/BF16, the default is enabled -trtllm-build --checkpoint_dir tmp/trt_models/${MODEL_NAME}/${INFERENCE_PRECISION}/encoder \ - --output_dir tmp/trt_engines/${MODEL_NAME}/${INFERENCE_PRECISION}/encoder \ - --paged_kv_cache disable \ - --moe_plugin disable \ - --max_beam_width ${MAX_BEAM_WIDTH} \ - --max_batch_size 8 \ - --max_input_len 1024 \ - --gemm_plugin ${INFERENCE_PRECISION} \ - --bert_attention_plugin ${INFERENCE_PRECISION} \ - --gpt_attention_plugin ${INFERENCE_PRECISION} \ - --remove_input_padding enable - # --context_fmha disable should be removed - -# Use the same command for decoder engine -trtllm-build --checkpoint_dir tmp/trt_models/${MODEL_NAME}/${INFERENCE_PRECISION}/decoder \ - --output_dir tmp/trt_engines/${MODEL_NAME}/${INFERENCE_PRECISION}/decoder \ - --moe_plugin disable \ - --max_beam_width ${MAX_BEAM_WIDTH} \ - --max_batch_size 8 \ - --max_input_len 1 \ - --max_seq_len 201 \ - --max_encoder_input_len 1024 \ - --gemm_plugin ${INFERENCE_PRECISION} \ - --bert_attention_plugin ${INFERENCE_PRECISION} \ - --gpt_attention_plugin ${INFERENCE_PRECISION} \ - --remove_input_padding enable - # --context_fmha disable should be removed - -``` - -### Run - -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. - -Different types of runtime are provided for encoder-decoder models. Following an order of serving performance and good usability, we recommend: -- (NEW) Python binding of C++ runtime w/ Paged KV Cache and Inflight Batching (IFB) -- Python runtime w/ Static Batching -- (NEW) C++ runtime w/ Paged KV Cache and Inflight Batching - -Please refer to the documentation for the details of [paged kv cache](../../../../docs/source/legacy/advanced/gpt-attention.md#paged-kv-cache) and [inflight batching](../../../../docs/source/legacy/advanced/gpt-attention.md#inflight-batching). - -#### Run C++ runtime -**Note: to use inflight batching and paged kv cache features in C++ runtime, please make sure you have set `--paged_kv_cache enable` (which is by default enabled) in the `trtllm-build` command of the decoder. Meanwhile, if using Python runtime, it is recommended to disable this flag by `--paged_kv_cache disable` to avoid any unnecessary overhead.** - -Note that for C++ runtime and Triton backend, Pipeline Parallelism (PP) is not supported yet, because PP usage is relatively rare for encoder-decoder models. If PP is really needed, it is recommended to use the Python runtime instead. - -For good usability, Python binding of the C++ runtime is provided. You can use the high-level C++ `ModelRunner` under the `examples/` root folder. - -```python -# Inferencing via python binding of C++ runtime with inflight batching (IFB) -python3 ../../../run.py --engine_dir tmp/trt_engines/${MODEL_NAME}/${INFERENCE_PRECISION} --tokenizer_dir tmp/hf_models/${MODEL_NAME} --max_output_len 64 --num_beams=1 --input_text "translate English to German: The house is wonderful." -``` - -You can specify `--kv_cache_free_gpu_memory_fraction` to control the percentage of free GPU memory to be used by KV cache (by default 0.9), and `--cross_kv_cache_fraction` to control the percentage of KV cache to be used by cross attention (by default 0.5, and rest of the KV cache will be used by self attention). - -For pure C++ runtime, there is no example given yet. Please check the [`Executor`](../../../../cpp/include/tensorrt_llm/executor/executor.h) API to implement your own end-to-end workflow. It is highly recommended to leverage more encapsulated solutions such as the above C++ Python binding or [Triton backend](https://github.com/triton-inference-server/tensorrtllm_backend). - -#### Run with Triton Backend -[Triton backend](https://github.com/triton-inference-server/tensorrtllm_backend/blob/main/docs/encoder_decoder.md) contains the tutorial on how to run encoder-decoder engines with Tritonserver. - -#### Run Python runtime - -For pure Python runtime, you can still use the encoder-decoder specific script under `examples/models/core/enc_dec/`. - -```bash -# Inferencing w/ single GPU greedy search, compare results with HuggingFace FP32 -python3 run.py --engine_dir tmp/trt_engines/${MODEL_NAME}/${INFERENCE_PRECISION} --engine_name ${MODEL_NAME} --model_name tmp/hf_models/${MODEL_NAME} --max_new_token=64 --num_beams=1 --compare_hf_fp32 - -# Inferencing w/ 4 GPUs (4-way TP, as configured during the engine building step), greedy search, compare results with HuggingFace FP32 -mpirun --allow-run-as-root -np ${WORLD_SIZE} python3 run.py --engine_dir tmp/trt_engines/${MODEL_NAME}/${INFERENCE_PRECISION} --engine_name ${MODEL_NAME} --model_name tmp/hf_models/${MODEL_NAME} --max_new_token=64 --num_beams=1 --compare_hf_fp32 -``` - -### Benchmark - -#### Benchmark C++ runtime - -The tutorial for encoder-decoder C++ runtime benchmark can be found in [`benchmarks/cpp`](../../benchmarks/cpp/README.md#2-launch-c-benchmarking-inflightv1-batching) - - -### Run BART with LoRA - -* Download the base model and lora model from HF: - -```bash -git clone https://huggingface.co/facebook/bart-large-cnn tmp/hf_models/bart-large-cnn -git clone https://huggingface.co/sooolee/bart-large-cnn-samsum-lora tmp/hf_models/bart-large-cnn-samsum-lora -``` - -If using customize models, just put both the base model and lora model dirs into `tmp/hf_models`. - -* Convert and Split Weights, setting `--hf_lora_dir`. - -```bash -export INFERENCE_PRECISION="float16" -python convert_checkpoint.py --model_type bart \ - --model_dir tmp/hf_models/bart-large-cnn \ - --output_dir tmp/trt_models/bart-large-cnn/${INFERENCE_PRECISION} \ - --tp_size 1 \ - --pp_size 1 \ - --dtype ${INFERENCE_PRECISION} -``` - -* Build engine, setting `--use_lora_plugin`. - -```bash - -trtllm-build --checkpoint_dir tmp/trt_models/bart-large-cnn/${INFERENCE_PRECISION}/encoder \ - --output_dir tmp/trt_engines/bart-large-cnn/${INFERENCE_PRECISION}/encoder \ - --paged_kv_cache disable \ - --moe_plugin disable \ - --max_beam_width 1 \ - --max_batch_size 8 \ - --max_input_len 1024 \ - --gemm_plugin ${INFERENCE_PRECISION} \ - --bert_attention_plugin ${INFERENCE_PRECISION} \ - --gpt_attention_plugin ${INFERENCE_PRECISION} \ - --remove_input_padding disable \ - --lora_plugin ${INFERENCE_PRECISION} \ - --lora_dir tmp/hf_models/bart-large-cnn-samsum-lora/ \ - --lora_target_modules attn_q attn_v - -trtllm-build --checkpoint_dir tmp/trt_models/bart-large-cnn/${INFERENCE_PRECISION}/decoder \ - --output_dir tmp/trt_engines/bart-large-cnn/${INFERENCE_PRECISION}/decoder \ - --moe_plugin disable \ - --max_beam_width 1 \ - --max_batch_size 8 \ - --max_input_len 1 \ - --max_seq_len 201 \ - --max_encoder_input_len 1024 \ - --gemm_plugin ${INFERENCE_PRECISION} \ - --bert_attention_plugin ${INFERENCE_PRECISION} \ - --gpt_attention_plugin ${INFERENCE_PRECISION} \ - --remove_input_padding disable \ - --lora_plugin ${INFERENCE_PRECISION} \ - --lora_dir tmp/hf_models/bart-large-cnn-samsum-lora/ \ - --lora_target_modules attn_q cross_attn_q attn_v cross_attn_v -``` - -* Run the engine, setting `--lora_dir` and `--lora_task_uids`. `--lora_task_uids` should be set as a list of uids which length equals to batch size. The following example is for batch size = 3: - -```bash -python run.py \ - --engine_dir tmp/trt_engines/bart-large-cnn/${INFERENCE_PRECISION}/ \ - --engine_name bart-large-cnn \ - --model_name tmp/hf_models/bart-large-cnn \ - --max_new_token=64 \ - --num_beams=1 \ - --lora_dir tmp/hf_models/bart-large-cnn-samsum-lora/ \ - --lora_task_uids 0 0 0 -``` - -* Run with multi-loRA, append `--lora_dir` with other lora directories and set `--lora_task_uids` according to the index of the lora directories. Set to "-1" to run with the base model: - -```bash -python run.py \ - --engine_dir tmp/trt_engines/bart-large-cnn/${INFERENCE_PRECISION}/ \ - --engine_name bart-large-cnn \ - --model_name tmp/hf_models/bart-large-cnn \ - --max_new_token=64 \ - --num_beams=1 \ - --lora_dir tmp/hf_models/bart-large-cnn-samsum-lora/ ... \ - --lora_task_uids 0 -1 -1 0 0 -1 -``` - -### Reminders - -- Flan-T5 models have known issues regarding FP16 precision and using BF16 precision is recommended, regardless of TRT-LLM. Please stay with FP32 or BF16 precision for Flan-T5 family. -- For T5 and Flan-T5 family that have relative attention bias design, the relative attention table is split along `num_heads` dimension in Tensor Parallelism mode. Therefore, `num_heads` must be divisible by `tp_size`. Please be aware of this when setting the TP parameter. -- For mBART, models that can control output languages (e.g. [`mbart-large-50-many-to-many-mmt`](https://huggingface.co/facebook/mbart-large-50-many-to-many-mmt)) are not currently supported, as the script does not support `ForcedBOSTokenLogitsProcessor` to control output languages. - -### Attention Scaling Factors - -The `q_scaling` convention in the TRT-LLM plugin is defined as follows: -``` -norm_factor = 1.f / (q_scaling * sqrt(head_size)) -``` -In the Multi-Head Attention (MHA) mechanism, the output of the `Q*K^T` product is scaled by this constant value `norm_factor` as `norm_factor * (Q*K^T)` for `softmax`. This scaling factor can be adjusted or neutralized based on the model's requirements. - -Handling in Different Models: -- BART/FairSeq NMT: For the BART models, `q_scaling` is set to `1.f`. Therefore, the `norm_factor` for BART becomes `1.f / sqrt(head_size)`. TRT-LLM uses the default value `q_scaling = 1.f`. Similar to FairSeq NMT models. -- T5: For the T5 models, `q_scaling` is `1.f/sqrt(head_size)`, leading to a `norm_factor` of `1.f`. This is handled in T5 by the TRT-LLM's `get_offset_q_scaling()` function, which reads `head_size` from the T5 model configuration and sets `q_scaling = 1.f/sqrt(head_size)` to effectively offset the `norm_factor` to `1.f`. - -### Run FairSeq NMT (Neural Machine Translation) models - -FairSeq model download and library dependency are different from HuggingFace ones. Especially if you are following the recommended docker container setup in [README](../../README.md), it has a custom PyTorch build but FairSeq installation will force upgrade the PyTorch version. As a workaround, we skip the `torch` and `torchaudio` dependencies in FairSeq to make everything work nicely inside the TRT-LLM container. - -```bash -# Download weights from HuggingFace Transformers -# Instructions from: https://github.com/facebookresearch/fairseq/blob/main/examples/translation/README.md#example-usage-cli-tools. Public model checkpoints are also listed there. Here we use WMT'14 Transformer model as an example. -mkdir -p tmp/fairseq_models && curl https://dl.fbaipublicfiles.com/fairseq/models/wmt14.en-fr.joined-dict.transformer.tar.bz2 | tar xvjf - -C tmp/fairseq_models --one-top-level=wmt14 --strip-components 1 --no-same-owner - -# Install FairSeq dependency -# avoid base torch to be upgraded by fairseq -pushd tmp && (git clone https://github.com/facebookresearch/fairseq.git || true) && pushd fairseq && sed -i '/torch>=/d;/torchaudio>=/d' setup.py && pip install -e . && pip install sacremoses subword_nmt && popd && popd - -# Convert and Split Weights, single GPU example -export TP_SIZE=1 -export PP_SIZE=1 -export WORLD_SIZE=1 -export INFERENCE_PRECISION="float32" -python convert_checkpoint.py --model_type nmt \ - --model_dir tmp/fairseq_models/wmt14 \ - --output_dir tmp/trt_models/wmt14/${INFERENCE_PRECISION} \ - --tp_size ${TP_SIZE} \ - --pp_size ${PP_SIZE} \ - --dtype ${INFERENCE_PRECISION} - -# Build TensorRT engine(s) -# Note: non-T5 models can enable FMHA for the encoder part, although only FP16/BF16 precisions are valid -trtllm-build --checkpoint_dir tmp/trt_models/wmt14/${INFERENCE_PRECISION}/encoder \ - --output_dir tmp/trt_engines/wmt14/${INFERENCE_PRECISION}/encoder \ - --paged_kv_cache disable \ - --moe_plugin disable \ - --max_beam_width 1 \ - --max_batch_size 8 \ - --max_input_len 1024 \ - --bert_attention_plugin ${INFERENCE_PRECISION} \ - --gpt_attention_plugin ${INFERENCE_PRECISION} \ - --remove_input_padding disable - -trtllm-build --checkpoint_dir tmp/trt_models/wmt14/${INFERENCE_PRECISION}/decoder \ - --output_dir tmp/trt_engines/wmt14/${INFERENCE_PRECISION}/decoder \ - --moe_plugin disable \ - --max_beam_width 1 \ - --max_batch_size 8 \ - --max_input_len 1 \ - --max_seq_len 201 \ - --max_encoder_input_len 1024 \ - --bert_attention_plugin ${INFERENCE_PRECISION} \ - --gpt_attention_plugin ${INFERENCE_PRECISION} \ - --remove_input_padding disable -# Run -mpirun --allow-run-as-root -np ${WORLD_SIZE} python3 run.py --engine_dir tmp/trt_engines/wmt14/${INFERENCE_PRECISION} --engine_name wmt14 --model_name tmp/fairseq_models/wmt14/${INFERENCE_PRECISION} --max_new_token=24 --num_beams=1 -``` - -### FP8 Post-Training Quantization - -The examples below uses the NVIDIA Modelopt (AlgorithMic Model Optimization) toolkit for the model quantization process. - -First make sure Modelopt toolkit `nvidia-modelopt>=0.22.1` is installed (see [examples/quantization/README.md](/examples/quantization/README.md#preparation)). - -> [!NOTE] -> Modelopt 0.22.1 is not yet released. - -#### Get quantized checkpoint with ModelOpt -Currently supported conversion are `bart-large-cnn` and `T5` family. For `bart`, please set `--dtype float16`; for `T5` family, please set `--dtype float32` due to known bug with apex+HF mentioned in [transformer:issue/34264](https://github.com/huggingface/transformers/issues/34264). - -```bash -# Example: quantize bart-large-cnn using 4-way tensor parallelism on a node with 8 GPUs (but only use 4 of them, for demonstration purpose) into FP8 weight and convert to TRTLLM checkpoint. -export MODEL_NAME="bart-large-cnn" -export MODEL_TYPE="bart" -export INFERENCE_PRECISION="float16" -export TP_SIZE=4 -export PP_SIZE=1 -export WORLD_SIZE=4 -export MAX_BEAM_WIDTH=1 -python ../quantization/quantize.py \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --dtype ${INFERENCE_PRECISION} \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir tmp/trt_models/${MODEL_NAME}/fp8 \ - --calib_size 512 \ - --batch_size 16 \ - --tp_size ${TP_SIZE} \ - --pp_size ${PP_SIZE} -``` - -The rest may follow the same command in [Build TensorRT engine(s)](#build-tensorrt-engines), with some notes: -* For `bart`, please add `--use_fp8_context_fmha enable` for fp8 context fmha support. For `t5`, context fmha is not supported due to relative attention bias. -* Please ensure `--paged_kv_cache enable` for decoder for fp8 paged kv cache. -* Please use `--gemm_plugin auto`, `--bert_attention_plugin auto`, `--gpt_attention_plugin auto` instead of setting precision to these plugins. -* Please use CPP runtime for better performance. diff --git a/examples/models/core/enc_dec/__init__.py b/examples/models/core/enc_dec/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/examples/models/core/enc_dec/convert_checkpoint.py b/examples/models/core/enc_dec/convert_checkpoint.py deleted file mode 100755 index 883a2fdebfff..000000000000 --- a/examples/models/core/enc_dec/convert_checkpoint.py +++ /dev/null @@ -1,1958 +0,0 @@ -import argparse -import configparser -import copy -import json -import logging -import os -import types -from ast import literal_eval -from collections import defaultdict -from datetime import datetime -from pathlib import Path - -import safetensors -from helper import (convert_weight_to_dtype, fairseq_sin_pos_embedding, - fuse_qkv_one_layer, reshape, split) -from transformers import (AutoModelForSeq2SeqLM, Blip2ForConditionalGeneration, - MBartForConditionalGeneration, NougatProcessor, - Pix2StructForConditionalGeneration, - T5ForConditionalGeneration, VisionEncoderDecoderModel) - -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import pad_vocab_size -from tensorrt_llm.functional import (LayerNormPositionType, LayerNormType, - MLPType) -from tensorrt_llm.layers import LanguageAdapterConfig -from tensorrt_llm.models import PretrainedConfig - -dir_path = os.path.dirname(os.path.realpath(__file__)) -LOGGER = logging.getLogger(__name__) - -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} - -# Constants for specific model configurations -ECLAIR_RADIO_MAX_POSITION_EMBEDDINGS = 20000 - - -def copy_args_to_component_config(component_config, args): - for arg in vars(args): - setattr(component_config, arg, getattr(args, arg)) - return component_config - - -def parse_t5_config(args, hf_model): - config = configparser.ConfigParser() - - config["encoder"] = {} - for key, val in hf_model.encoder.config.to_dict().items(): - config["encoder"][key] = f"{val}" - - # 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): - scaling = 1 / config.head_size**.5 - return scaling - - config["decoder"] = {} - for key, val in hf_model.decoder.config.to_dict().items(): - config["decoder"][key] = f"{val}" - - config["structure"] = dict() - config["structure"]["t5_with_bias"] = "false" - config["structure"]["use_gated_activation"] = str( - hf_model.encoder.config.is_gated_act) - config["structure"]["position_embedding_type"] = "relative" - config["structure"]["model_type"] = args.model_type - - def parse_t5_config_by_component(config, component, args): - component_config = types.SimpleNamespace() - component_config = copy_args_to_component_config(component_config, args) - component_config.n_head = config.getint(component, 'num_heads') - component_config.head_size = config.getint(component, 'd_kv') - component_config.hidden_size = config.getint(component, 'd_model') - component_config.ffn_hidden_size = config.getint(component, 'd_ff') - component_config.vocab_size = config.getint(component, 'vocab_size') - component_config.n_positions = config.getint(component, - 'n_positions', - fallback=512) - component_config.has_position_embedding = config.getboolean( - component, 'has_position_embedding', - fallback=False) # TODO: hardcoded here - - component_config.has_token_type_embedding = config.getboolean( - component, 'has_token_type_embedding', fallback=False) - component_config.has_embedding_layernorm = config.getboolean( - component, 'has_embedding_layernorm', fallback=False) - component_config.has_embedding_scale = config.getboolean( - component, 'has_embedding_scale', fallback=False) - component_config.q_scaling = get_offset_q_scaling(component_config) - component_config.has_attention_qkvo_bias = config.getboolean( - component, 'has_attention_qkvo_bias', - fallback=False) # TODO: hardcoded here - component_config.has_mlp_bias = config.getboolean(component, - 'has_mlp_bias', - fallback=False) - component_config.has_model_final_layernorm = config.getboolean( - component, 'has_model_final_layernorm', fallback=True) - component_config.layernorm_eps = config.getfloat( - component, 'layer_norm_epsilon') - component_config.layernorm_position = layernorm_position_map[config.get( - component, 'layernorm_position', - fallback='pre_layernorm')] # TODO: hardcoded here - component_config.layernorm_type = layernorm_type_map[config.get( - component, 'layernorm_type', fallback='RmsNorm')] - component_config.hidden_act = config.get(component, 'dense_act_fn') - component_config.gated_act = config.getboolean(component, - 'is_gated_act') - component_config.mlp_type = mlp_type_map['GatedMLP' if component_config. - gated_act else 'MLP'] - component_config.num_buckets = config.getint( - component, 'relative_attention_num_buckets') - component_config.max_distance = config.getint( - component, 'relative_attention_max_distance') - component_config.position_embedding_type = config.get( - 'structure', 'position_embedding_type') - component_config.logits_dtype = config.get(component, - 'logits_dtype', - fallback='float32') - - if component == 'encoder': - component_config.n_layer = config.getint(component, 'num_layers') - - component_config.relative_attention = config.get( - 'structure', 'position_embedding_type') == 'relative' - - elif component == 'decoder': - component_config.n_layer = config.getint(component, - 'num_decoder_layers') - component_config.has_lm_head_bias = config.getboolean( - component, # TODO: T5 with bias - 'has_lm_head_bias', - fallback=False) - component_config.relative_attention = config.getboolean( - component, 'relative_attention', fallback=True) - component_config.rescale_before_lm_head = config.getboolean( - component, 'tie_word_embeddings' - ) # default is True (for T5), but False for Flan-T5 - component_config.encoder_hidden_size = config.getint( - 'encoder', 'd_model') - component_config.encoder_num_heads = config.getint( - 'encoder', 'num_heads') - component_config.encoder_head_size = config.getint( - 'encoder', 'd_kv') - component_config.decoder_start_token_id = config.getint( - 'decoder', 'decoder_start_token_id', fallback=0) - component_config.eos_token_id = config.getint('decoder', - 'eos_token_id', - fallback=1) - bos_token_id = config.get('decoder', 'bos_token_id', fallback=None) - # T5 does not have bos_token_id - component_config.bos_token_id = int( - bos_token_id) if bos_token_id not in (None, "None") else None - component_config.pad_token_id = config.getint('decoder', - 'pad_token_id', - fallback=0) - - else: - assert False, 'Unsupported component!' - - return component_config - - encoder_config = parse_t5_config_by_component(config, "encoder", args) - decoder_config = parse_t5_config_by_component(config, "decoder", args) - - return encoder_config, decoder_config - - -def convert_t5_weights_to_tllm_safetensors(config, component, params): - weights = {} - - mapping = config.mapping - - convert_weight_to_dtype(params, config.dtype) - hidden_size = config.hidden_size - ffn_hidden_size = config.intermediate_size - num_layers = config.num_hidden_layers - n_head = config.num_attention_heads - head_size = config.head_size - attention_hidden_size = n_head * head_size # head size * num_heads not necessarily equals hidden_dim, such as Flan-T5 - - hf_param_prefix = f'{component}' - trtllm_layer_name = f'transformer.layers' - trtllm_attn_layer_name = 'attention' if component == 'encoder' else 'self_attention' - trtllm_attn_layernorm_name = 'self_attention_layernorm' if component == 'decoder' else 'attention_layernorm' - hf_component_idx = 1 if component == 'encoder' else 2 - - def get_attn_module_name(component, block, layer, attn_type): - return f'{component}.block.{int(block)}.layer.{int(layer)}.{attn_type}' - - weights['transformer.vocab_embedding.weight'] = reshape( - params['shared.weight'].clone(), - None) if not config.use_parallel_embedding else reshape( - split(params['shared.weight'].clone(), mapping.tp_size, - mapping.tp_rank, 0), None) - - layers_range = mapping.pp_layers(num_layers) - for layer_idx in layers_range: - local_layer_idx = layer_idx - layers_range[0] - trtllm_layer_name_prefix = f'{trtllm_layer_name}.{local_layer_idx}' - hf_layer_name_prefix = f'{hf_param_prefix}.block.{layer_idx}' - - hidden_layer_name_split = { - f'{hf_layer_name_prefix}.layer.0.SelfAttention.o.weight': { - "name": - f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}.dense.weight', - "shape": - (hidden_size, attention_hidden_size // mapping.tp_size), - "split_dim": -1 - }, - f'{hf_layer_name_prefix}.layer.{hf_component_idx}.DenseReluDense.wo.weight': - { - "name": f'{trtllm_layer_name_prefix}.mlp.proj.weight', - "shape": (hidden_size, ffn_hidden_size // mapping.tp_size), - "split_dim": -1 - }, - f'{hf_layer_name_prefix}.layer.{hf_component_idx}.DenseReluDense.wi.weight': - { - "name": f'{trtllm_layer_name_prefix}.mlp.fc.weight', - "shape": (ffn_hidden_size // mapping.tp_size, hidden_size), - "split_dim": 0 - }, - f'{hf_layer_name_prefix}.layer.{hf_component_idx}.DenseReluDense.wi_0.weight': - { - "name": f'{trtllm_layer_name_prefix}.mlp.fc.weight', - "shape": (ffn_hidden_size // mapping.tp_size, hidden_size), - "split_dim": 0 - }, - } - - hidden_layer_name_no_split = { - f'{hf_layer_name_prefix}.layer.0.layer_norm.weight': { - "name": - f'{trtllm_layer_name_prefix}.{trtllm_attn_layernorm_name}.weight', - "shape": None - }, - f'{hf_layer_name_prefix}.layer.{hf_component_idx}.layer_norm.weight': - { - "name": f'{trtllm_layer_name_prefix}.mlp_layernorm.weight', - "shape": None - }, - } - - if config.gated_act: - hidden_layer_name_split.update({ - f'{hf_layer_name_prefix}.layer.{hf_component_idx}.DenseReluDense.wi2.weight': - { - "name": f'{trtllm_layer_name_prefix}.mlp.gate.weight', - "shape": (ffn_hidden_size // mapping.tp_size, hidden_size), - "split_dim": 0 - }, - f'{hf_layer_name_prefix}.layer.{hf_component_idx}.DenseReluDense.wi_1.weight': - { - "name": f'{trtllm_layer_name_prefix}.mlp.gate.weight', - "shape": (ffn_hidden_size // mapping.tp_size, hidden_size), - "split_dim": 0 - }, - }) - - if component == 'decoder': - hidden_layer_name_split.update({ - f'{hf_layer_name_prefix}.layer.1.EncDecAttention.o.weight': { - "name": - f'{trtllm_layer_name_prefix}.cross_attention.dense.weight', - "shape": - (hidden_size, attention_hidden_size // mapping.tp_size), - "split_dim": -1 - }, - }) - hidden_layer_name_no_split.update({ - f'{hf_layer_name_prefix}.layer.1.layer_norm.weight': { - "name": - f'{trtllm_layer_name_prefix}.cross_attention_layernorm.weight', - "shape": None - }, - }) - self_attn_module_name = get_attn_module_name( - component, layer_idx, "1", 'EncDecAttention') - weights.update( - fuse_qkv_one_layer( - params, self_attn_module_name, - f'{trtllm_layer_name_prefix}.cross_attention', - mapping.tp_size, mapping.tp_rank, config.model_type, - (attention_hidden_size * 3 // mapping.tp_size, hidden_size), - None)) - - self_attn_module_name = get_attn_module_name(component, layer_idx, "0", - 'SelfAttention') - weights.update( - fuse_qkv_one_layer( - params, self_attn_module_name, - f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}', - mapping.tp_size, mapping.tp_rank, config.model_type, - (attention_hidden_size * 3 // mapping.tp_size, hidden_size), - None)) - - weights[ - f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}.rel_attn_table'] = reshape( - split( - params[ - f'{component}.block.0.layer.0.SelfAttention.relative_attention_bias.weight'] - .T, mapping.tp_size, mapping.tp_rank, 0), - (n_head // mapping.tp_size, config.num_buckets)) - - for hf_weight_name, weight_info in hidden_layer_name_split.items(): - if hf_weight_name in params.keys(): - weights[weight_info["name"]] = reshape( - split(params[hf_weight_name], - mapping.tp_size, - mapping.tp_rank, - dim=weight_info["split_dim"]), weight_info["shape"]) - for hf_weight_name, weight_info in hidden_layer_name_no_split.items(): - if hf_weight_name in params.keys(): - weights[weight_info["name"]] = reshape( - params[hf_weight_name].clone(), shape=weight_info["shape"]) - - weights['transformer.ln_f.weight'] = reshape( - params[f'{component}.final_layer_norm.weight'].clone(), None) - - if component == 'decoder': - weights['lm_head.weight'] = reshape( - split(params['lm_head.weight'], - mapping.tp_size, - mapping.tp_rank, - dim=0), (config.vocab_size // mapping.tp_size, hidden_size)) - if not config.use_implicit_relative_attention: - weights['rel_attn_table'] = reshape( - split( - params[ - f'{component}.block.0.layer.0.SelfAttention.relative_attention_bias.weight'] - .T, mapping.tp_size, mapping.tp_rank, 0), - (n_head // mapping.tp_size, config.num_buckets)) - - return weights - - -convert_blip2_weights_to_tllm_safetensors = convert_t5_weights_to_tllm_safetensors # func alias - - -def parse_nmt_config_by_component(config, component, args): - assert component in ('encoder', 'decoder'), 'Unsupported component!' - component_config = types.SimpleNamespace() - component_config = copy_args_to_component_config(component_config, args) - component_config.n_layer = config.getint(component, f'{component}_layers') - component_config.n_head = config.getint(component, - f'{component}_attention_heads') - component_config.hidden_size = config.getint( - component, f'{component}_embed_dim') # fairseq naming - component_config.head_size = config.getint( - component, - 'd_kv', - fallback=component_config.hidden_size // component_config.n_head) - component_config.ffn_hidden_size = config.getint( - component, f'{component}_ffn_embed_dim') # fairseq naming - component_config.vocab_size = config.getint(component, 'vocab_size') - component_config.n_positions = config.getint( - component, 'max_source_positions') # fairseq naming - component_config.has_position_embedding = not config.getboolean( - component, 'no_token_positional_embeddings', - fallback=False) # fairseq naming - component_config.has_token_type_embedding = config.getboolean( - component, 'has_token_type_embedding', fallback=False) - component_config.has_embedding_layernorm = config.getboolean( - component, 'layernorm_embedding', fallback=True) # fairseq naming - component_config.has_embedding_scale = not config.getboolean( - component, 'no_scale_embedding') # fairseq naming - component_config.q_scaling = config.getfloat(component, - 'q_scaling', - fallback=1.0) - component_config.has_attention_qkvo_bias = config.getboolean('structure', - 't5_with_bias', - fallback=True) - component_config.has_mlp_bias = config.getboolean('structure', - 't5_with_bias', - fallback=True) - component_config.has_model_final_layernorm = config.getboolean( - component, 'has_model_final_layernorm') - component_config.layernorm_eps = config.getfloat( - component, 'layer_norm_epsilon', fallback=1e-5) # fairseq naming - - normalize_before = config.getboolean( - component, f'{component}_normalize_before') # fairseq naming - component_config.layernorm_position = layernorm_position_map[ - 'pre_layernorm' if normalize_before else 'post_layernorm'] - - component_config.layernorm_type = layernorm_type_map[config.get( - component, 'layernorm_type', fallback='LayerNorm')] - component_config.hidden_act = config.get(component, - 'activation_fn') # fairseq naming - component_config.gated_act = config.getboolean(component, - 'is_gated_act', - fallback=False) - component_config.mlp_type = mlp_type_map['GatedMLP' if component_config. - gated_act else 'MLP'] - component_config.relative_attention = config.get( - 'structure', 'position_embedding_type') == 'relative' - - component_config.num_buckets = config.getint( - component, 'relative_attention_num_buckets', fallback=0) - component_config.max_distance = config.getint( - component, 'relative_attention_max_distance', fallback=0) - component_config.position_embedding_type = config.get( - 'structure', 'position_embedding_type') - component_config.logits_dtype = config.get(component, - 'logits_dtype', - fallback='float32') - if component == 'decoder': - component_config.rescale_before_lm_head = config.getboolean( - component, 'rescale_before_lm_head') - - component_config.encoder_hidden_size = config.getint( - 'encoder', 'encoder_embed_dim') # fairseq naming - component_config.encoder_num_heads = config.getint( - 'encoder', 'encoder_attention_heads') - component_config.encoder_head_size = config.getint( - 'encoder', - 'd_kv', - fallback=component_config.encoder_hidden_size // - component_config.encoder_num_heads) - component_config.decoder_start_token_id = None - component_config.eos_token_id = None - component_config.bos_token_id = None - component_config.pad_token_id = None - - return component_config - - -def parse_nmt_config(args, model): - config = configparser.ConfigParser() - fairseq_config = vars(model.cfg.model) # Namespace --> dict - - config['encoder'] = dict() - for key, val in fairseq_config.items(): - config["encoder"][key] = f"{val}" - config["encoder"]["q_scaling"] = '1' - # NMT has final layernorm for pre-norm model architecture. - config['encoder']['has_model_final_layernorm'] = config['encoder'][ - 'encoder_normalize_before'] - config['encoder']['vocab_size'] = str(len(model.src_dict)) # fairseq naming - - config['decoder'] = dict() - for key, val in fairseq_config.items(): - config["decoder"][key] = f"{val}" - config["decoder"]["q_scaling"] = '1' - config["decoder"]["rescale_before_lm_head"] = 'false' - config['decoder']['has_model_final_layernorm'] = str( - config['decoder'].getboolean('decoder_normalize_before', False) - and not config['decoder'].getboolean('no_decoder_final_norm', False)) - config['decoder']['vocab_size'] = str(len(model.tgt_dict)) # fairseq naming - - config["structure"] = dict() - config["structure"]["t5_with_bias"] = "true" - config["structure"]["use_gated_activation"] = "false" - config["structure"][ - "position_embedding_type"] = "learned_absolute" # "sinusoid" - config["structure"]["model_type"] = args.model_type - - encoder_config = parse_nmt_config_by_component(config, "encoder", args) - decoder_config = parse_nmt_config_by_component(config, "decoder", args) - - return encoder_config, decoder_config - - -def convert_nmt_weights_to_tllm_safetensors(config, component, params, - sin_pos_embedding): - weights = {} - - mapping = config.mapping - - hidden_size = config.hidden_size - - convert_weight_to_dtype(params, config.dtype) - ffn_hidden_size = config.intermediate_size - vocab_size = config.vocab_size - - hf_param_prefix = f'models.0.{component}' - trtllm_layer_name = f'transformer.layers' - trtllm_attn_layer_name = 'attention' if component == 'encoder' else 'self_attention' - trtllm_attn_layernorm_name = 'self_attention_layernorm' if component == 'decoder' else 'attention_layernorm' - - hidden_layer_name_split = { - 'self_attn.out_proj.weight': { - "name": f'{trtllm_attn_layer_name}.dense.weight', - "shape": (hidden_size, hidden_size // mapping.tp_size), - "split_dim": -1 - }, - 'fc1.weight': { - "name": 'mlp.fc.weight', - "shape": (ffn_hidden_size // mapping.tp_size, hidden_size), - "split_dim": 0 - }, - 'fc1.bias': { - "name": 'mlp.fc.bias', - "shape": (ffn_hidden_size // mapping.tp_size), - "split_dim": 0 - }, - 'fc2.weight': { - "name": 'mlp.proj.weight', - "shape": (hidden_size, ffn_hidden_size // mapping.tp_size), - "split_dim": -1 - }, - } - - hidden_layer_name_no_split = { - 'self_attn.out_proj.bias': { - "name": f'{trtllm_attn_layer_name}.dense.bias', - "shape": (hidden_size) - }, - 'self_attn_layer_norm.weight': { - "name": f'{trtllm_attn_layernorm_name}.weight', - "shape": None - }, - 'self_attn_layer_norm.bias': { - "name": f'{trtllm_attn_layernorm_name}.bias', - "shape": None - }, - 'fc2.bias': { - "name": 'mlp.proj.bias', - "shape": (hidden_size) - }, - 'final_layer_norm.weight': { - "name": 'mlp_layernorm.weight', - "shape": None - }, - 'final_layer_norm.bias': { - "name": 'mlp_layernorm.bias', - "shape": None - }, - } - - if component == "decoder": - hidden_layer_name_split.update({ - 'encoder_attn.out_proj.weight': { - "name": 'cross_attention.dense.weight', - "shape": (hidden_size, hidden_size // mapping.tp_size), - "split_dim": -1 - }, - }) - hidden_layer_name_no_split.update({ - 'encoder_attn.out_proj.bias': { - "name": 'cross_attention.dense.bias', - "shape": (hidden_size) - }, - 'encoder_attn_layer_norm.weight': { - "name": 'cross_attention_layernorm.weight', - "shape": None, - }, - 'encoder_attn_layer_norm.bias': { - "name": 'cross_attention_layernorm.bias', - "shape": None - }, - }) - - def get_attn_module_name(component, layer, attn_type): - return f'models.0.{component}.layers.{int(layer)}.{attn_type}' - - weights["transformer.vocab_embedding.weight"] = reshape( - params[f'{hf_param_prefix}.embed_tokens.weight'].clone(), - (vocab_size, -1)) if not config.use_parallel_embedding else reshape( - split(params[f'{hf_param_prefix}.embed_tokens.weight'].clone(), - mapping.tp_size, mapping.tp_rank, 0), - (vocab_size // mapping.tp_size, -1)) - weights["transformer.position_embedding.weight"] = reshape( - sin_pos_embedding, (config.max_position_embeddings, hidden_size)) - - num_layers = config.num_hidden_layers - - layers_range = mapping.pp_layers(num_layers) - for layer_idx in layers_range: - local_layer_idx = layer_idx - layers_range[0] - hf_layer_name_prefix = f'{hf_param_prefix}.layers.{layer_idx}' - trtllm_layer_name_prefix = f'{trtllm_layer_name}.{local_layer_idx}' - - for hf_weight_name, weight_info in hidden_layer_name_split.items(): - weights[ - f'{trtllm_layer_name_prefix}.{weight_info["name"]}'] = reshape( - split(params[f'{hf_layer_name_prefix}.{hf_weight_name}'], - mapping.tp_size, - mapping.tp_rank, - dim=weight_info["split_dim"]), weight_info["shape"]) - - for hf_weight_name, weight_info in hidden_layer_name_no_split.items(): - trtllm_layer_fullname = f'{trtllm_layer_name_prefix}.{weight_info["name"]}' - hf_layer_fullname = f'{hf_layer_name_prefix}.{hf_weight_name}' - weights[trtllm_layer_fullname] = reshape( - params[hf_layer_fullname].clone(), shape=weight_info["shape"]) - - self_attn_module_name = get_attn_module_name(component, layer_idx, - 'self_attn') - weights.update( - fuse_qkv_one_layer( - params, self_attn_module_name, - f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}', - mapping.tp_size, mapping.tp_rank, config.model_type, - (hidden_size * 3 // mapping.tp_size, hidden_size), - (hidden_size * 3 // mapping.tp_size))) - if component == 'decoder': - cross_attn_module_name = get_attn_module_name( - component, layer_idx, 'encoder_attn') - weights.update( - fuse_qkv_one_layer( - params, cross_attn_module_name, - f'{trtllm_layer_name_prefix}.cross_attention', - mapping.tp_size, mapping.tp_rank, config.model_type, - (hidden_size * 3 // mapping.tp_size, hidden_size), - (hidden_size * 3 // mapping.tp_size))) - - if component == 'decoder': - weights['lm_head.weight'] = reshape( - split(params[f'{hf_param_prefix}.output_projection.weight'], - mapping.tp_size, - mapping.tp_rank, - dim=0), (config.vocab_size // mapping.tp_size, hidden_size)) - - if config.has_model_final_layernorm: - weights['transformer.ln_f.weight'] = params[ - f'{hf_param_prefix}.layer_norm.weight'].clone() - weights['transformer.ln_f.bias'] = params[ - f'{hf_param_prefix}.layer_norm.bias'].clone() - - return weights - - -def parse_bart_config(args, hf_model): - - config = configparser.ConfigParser() - - config['decoder'] = dict() - if args.eclair_radio: - for key, val in hf_model.config.to_dict().items(): - config["decoder"][key] = f"{val}" - else: - for key, val in hf_model.model.decoder.config.to_dict().items(): - config["decoder"][key] = f"{val}" - config["decoder"]["q_scaling"] = '1' - config["decoder"]["rescale_before_lm_head"] = str(False) - config['decoder']['has_model_final_layernorm'] = str( - args.nougat or args.eclair_radio - or isinstance(hf_model, MBartForConditionalGeneration)) - - if args.nougat or args.eclair_radio: - # These flags are true for mbart decoders, but missing in HF config - config['decoder']['normalize_before'] = str(True) - config['decoder']['normalize_embeddings'] = str(True) - - config['encoder'] = dict() - # Init few encoder configs, needed by build, from decoder config - encoder_config_keys = [ - "encoder_ffn_dim", "encoder_layers", "encoder_attention_heads", - "encoder_layerdrop", "d_model" - ] - for key in encoder_config_keys: - config['encoder'][key] = config['decoder'][key] - else: - config['encoder'] = dict() - for key, val in hf_model.model.encoder.config.to_dict().items(): - config["encoder"][key] = f"{val}" - config["encoder"]["q_scaling"] = '1' - - # mBART has final layernorm, BART does not - config['encoder']['has_model_final_layernorm'] = str( - isinstance(hf_model, MBartForConditionalGeneration)) - - config["structure"] = dict() - config["structure"]["t5_with_bias"] = "true" - config["structure"]["use_gated_activation"] = "false" - config["structure"]["position_embedding_type"] = "learned_absolute" - config["structure"]["model_type"] = args.model_type - - def parse_bart_config_by_component(config, component, args): - assert component in ('encoder', 'decoder'), 'Unsupported component!' - component_config = types.SimpleNamespace() - component_config = copy_args_to_component_config(component_config, args) - component_config.n_layer = config.getint(component, - f'{component}_layers') - component_config.n_head = config.getint(component, - f'{component}_attention_heads') - component_config.hidden_size = config.getint(component, 'd_model') - component_config.head_size = config.getint( - component, - 'd_kv', - fallback=component_config.hidden_size // component_config.n_head) - component_config.ffn_hidden_size = config.getint( - component, f'{component}_ffn_dim') - component_config.vocab_size = config.getint(component, 'vocab_size') - component_config.n_positions = config.getint(component, - 'max_position_embeddings') - component_config.has_position_embedding = config.getboolean( - component, 'has_position_embedding', - fallback=True) # TODO: hardcoded here - component_config.has_token_type_embedding = config.getboolean( - component, 'has_token_type_embedding', fallback=False) - component_config.has_embedding_layernorm = config.getboolean( - component, 'has_embedding_layernorm', fallback=True) - component_config.has_embedding_scale = config.getboolean( - component, 'scale_embedding') - component_config.q_scaling = config.getfloat(component, - 'q_scaling', - fallback=1.0) - component_config.has_attention_qkvo_bias = config.getboolean( - 'structure', 't5_with_bias', fallback=True) - component_config.has_mlp_bias = config.getboolean('structure', - 't5_with_bias', - fallback=True) - component_config.has_model_final_layernorm = config.getboolean( - component, 'has_model_final_layernorm') - component_config.layernorm_eps = config.getfloat(component, - 'layer_norm_epsilon', - fallback=False) - - normalize_before = config.getboolean(component, 'normalize_before') - component_config.layernorm_position = layernorm_position_map[ - 'pre_layernorm' if normalize_before else 'post_layernorm'] - - component_config.layernorm_type = layernorm_type_map[config.get( - component, 'layernorm_type', fallback='LayerNorm')] - component_config.hidden_act = config.get(component, - 'activation_function') - component_config.gated_act = config.getboolean(component, - 'is_gated_act', - fallback=False) - component_config.mlp_type = mlp_type_map['GatedMLP' if component_config. - gated_act else 'MLP'] - component_config.relative_attention = config.get( - 'structure', 'position_embedding_type') == 'relative' - - component_config.num_buckets = config.getint( - component, 'relative_attention_num_buckets', fallback=0) - component_config.max_distance = config.getint( - component, 'relative_attention_max_distance', fallback=0) - component_config.max_lora_rank = config.getint(component, - 'max_lora_rank', - fallback=0) - component_config.lora_target_modules = literal_eval( - config.get(component, 'lora_target_modules', fallback="[]")) - component_config.hf_modules_to_trtllm_modules = literal_eval( - config.get(component, 'hf_modules_to_trtllm_modules', - fallback="{}")) - component_config.trtllm_modules_to_hf_modules = literal_eval( - config.get(component, 'trtllm_modules_to_hf_modules', - fallback="{}")) - component_config.logits_dtype = config.get(component, - 'logits_dtype', - fallback='float32') - component_config.position_embedding_type = config.get( - 'structure', 'position_embedding_type') - - if component == 'decoder': - component_config.rescale_before_lm_head = config.getboolean( - component, 'rescale_before_lm_head') - - component_config.encoder_hidden_size = config.getint( - 'encoder', 'd_model') - component_config.encoder_num_heads = config.getint( - 'encoder', 'encoder_attention_heads') - component_config.encoder_head_size = config.getint( - 'encoder', - 'd_kv', - fallback=component_config.encoder_hidden_size // - component_config.encoder_num_heads) - - # nougat has decoder_start_token_id = None, special handling - decoder_start_token_id = config.get('decoder', - 'decoder_start_token_id') - component_config.decoder_start_token_id = int( - decoder_start_token_id - ) if decoder_start_token_id != "None" else None - component_config.eos_token_id = config.getint( - 'decoder', 'eos_token_id') - component_config.bos_token_id = config.getint( - 'decoder', 'bos_token_id') - component_config.pad_token_id = config.getint( - 'decoder', 'pad_token_id') - - return component_config - - encoder_config = None - if not (args.nougat or args.eclair_radio): - encoder_config = parse_bart_config_by_component(config, "encoder", args) - decoder_config = parse_bart_config_by_component(config, "decoder", args) - - # Override n_positions for eclair_radio model - if args.eclair_radio: - decoder_config.n_positions = ECLAIR_RADIO_MAX_POSITION_EMBEDDINGS - - return encoder_config, decoder_config - - -def convert_bart_weights_to_tllm_safetensors(config, component, params): - weights = {} - - mapping = config.mapping - - hidden_size = config.hidden_size - - convert_weight_to_dtype(params, config.dtype) - ffn_hidden_size = config.intermediate_size - vocab_size = config.vocab_size - - hf_param_prefix = f'model.{component}' - trtllm_layer_name = f'transformer.layers' - trtllm_attn_layer_name = 'attention' if component == 'encoder' else 'self_attention' - trtllm_attn_layernorm_name = 'self_attention_layernorm' if component == 'decoder' else 'attention_layernorm' - embedding_layer_names = { - 'embed_tokens.weight': { - "name": 'transformer.vocab_embedding.weight', - "shape": (vocab_size, -1) - }, - 'embed_positions.weight': { - "name": 'transformer.position_embedding.weight', - "shape": (config.max_position_embeddings, hidden_size) - }, - 'layernorm_embedding.weight': { - "name": 'transformer.ln_embed.weight', - "shape": None - }, - 'layernorm_embedding.bias': { - "name": 'transformer.ln_embed.bias', - "shape": None - }, - } - - hidden_layer_name_split = { - 'self_attn.out_proj.weight': { - "name": f'{trtllm_attn_layer_name}.dense.weight', - "shape": (hidden_size, hidden_size // mapping.tp_size), - "split_dim": -1 - }, - 'fc1.weight': { - "name": 'mlp.fc.weight', - "shape": (ffn_hidden_size // mapping.tp_size, hidden_size), - "split_dim": 0 - }, - 'fc1.bias': { - "name": 'mlp.fc.bias', - "shape": (ffn_hidden_size // mapping.tp_size), - "split_dim": 0 - }, - 'fc2.weight': { - "name": 'mlp.proj.weight', - "shape": (hidden_size, ffn_hidden_size // mapping.tp_size), - "split_dim": -1 - }, - } - - hidden_layer_name_no_split = { - 'self_attn.out_proj.bias': { - "name": f'{trtllm_attn_layer_name}.dense.bias', - "shape": (hidden_size) - }, - 'self_attn_layer_norm.weight': { - "name": f'{trtllm_attn_layernorm_name}.weight', - "shape": None - }, - 'self_attn_layer_norm.bias': { - "name": f'{trtllm_attn_layernorm_name}.bias', - "shape": None - }, - 'fc2.bias': { - "name": 'mlp.proj.bias', - "shape": (hidden_size) - }, - 'final_layer_norm.weight': { - "name": 'mlp_layernorm.weight', - "shape": None - }, - 'final_layer_norm.bias': { - "name": 'mlp_layernorm.bias', - "shape": None - }, - } - - if config.model_type == 'mbart': - hidden_layer_name_split['layer_norm.weight'] = { - "name": 'transformer.ln_f.weight', - "shape": None, - "split_dim": 0 - } - hidden_layer_name_no_split['layer_norm.bias'] = { - "name": 'transformer.ln_f.bias', - "shape": None, - "split_dim": 0 - } - - if component == "decoder": - hidden_layer_name_split.update({ - 'encoder_attn.out_proj.weight': { - "name": 'cross_attention.dense.weight', - "shape": (hidden_size, hidden_size // mapping.tp_size), - "split_dim": -1 - } - }) - hidden_layer_name_no_split.update({ - 'encoder_attn.out_proj.bias': { - "name": 'cross_attention.dense.bias', - "shape": (hidden_size) - }, - 'encoder_attn_layer_norm.weight': { - "name": 'cross_attention_layernorm.weight', - "shape": None - }, - 'encoder_attn_layer_norm.bias': { - "name": 'cross_attention_layernorm.bias', - "shape": None - }, - }) - - def get_attn_module_name(component, layer, attn_type): - return f'model.{component}.layers.{int(layer)}.{attn_type}' - - for hf_weight_name, weight_info in embedding_layer_names.items(): - if 'position' in hf_weight_name: - weights[weight_info["name"]] = params[ - f'{hf_param_prefix}.{hf_weight_name}'][2:].clone() - else: - weights[weight_info["name"]] = params[ - f'{hf_param_prefix}.{hf_weight_name}'].clone() - weights[weight_info["name"]] = reshape(weights[weight_info["name"]], - weight_info["shape"]) - - weights["embedding.vocab_embedding.weight"] = reshape( - params[f'{hf_param_prefix}.embed_tokens.weight'].clone(), - (vocab_size, -1)) if not config.use_parallel_embedding else reshape( - split(params[f'{hf_param_prefix}.embed_tokens.weight'].clone(), - mapping.tp_size, mapping.tp_rank, 0), - (vocab_size // mapping.tp_size, -1)) - - num_layers = config.num_hidden_layers - - layers_range = mapping.pp_layers(num_layers) - for layer_idx in layers_range: - local_layer_idx = layer_idx - layers_range[0] - hf_layer_name_prefix = f'{hf_param_prefix}.layers.{layer_idx}' - trtllm_layer_name_prefix = f'{trtllm_layer_name}.{local_layer_idx}' - - for hf_weight_name, weight_info in hidden_layer_name_split.items(): - weights[ - f'{trtllm_layer_name_prefix}.{weight_info["name"]}'] = reshape( - split(params[f'{hf_layer_name_prefix}.{hf_weight_name}'], - mapping.tp_size, - mapping.tp_rank, - dim=weight_info["split_dim"]), weight_info["shape"]) - - for hf_weight_name, weight_info in hidden_layer_name_no_split.items(): - trtllm_layer_fullname = f'{trtllm_layer_name_prefix}.{weight_info["name"]}' - hf_layer_fullname = f'{hf_layer_name_prefix}.{hf_weight_name}' - weights[trtllm_layer_fullname] = reshape( - params[hf_layer_fullname].clone(), shape=weight_info["shape"]) - - self_attn_module_name = get_attn_module_name(component, layer_idx, - 'self_attn') - weights.update( - fuse_qkv_one_layer( - params, self_attn_module_name, - f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}', - mapping.tp_size, mapping.tp_rank, config.model_type, - (hidden_size * 3 // mapping.tp_size, hidden_size), - (hidden_size * 3 // mapping.tp_size))) - if component == 'decoder': - cross_attn_module_name = get_attn_module_name( - component, layer_idx, 'encoder_attn') - weights.update( - fuse_qkv_one_layer( - params, cross_attn_module_name, - f'{trtllm_layer_name_prefix}.cross_attention', - mapping.tp_size, mapping.tp_rank, config.model_type, - (hidden_size * 3 // mapping.tp_size, hidden_size), - (hidden_size * 3 // mapping.tp_size))) - - if component == 'decoder': - import torch - lm_head_weights = params['lm_head.weight'].clone().detach() - vocab_size = config.vocab_size - if params['lm_head.weight'].shape[0] % mapping.tp_size != 0: - vocab_size_padded = pad_vocab_size(config.vocab_size, - mapping.tp_size) - pad_width = vocab_size_padded - config.vocab_size - - lm_head_weights = torch.nn.functional.pad(lm_head_weights, - (0, 0, 0, pad_width), - 'constant', - value=0) - vocab_size = vocab_size_padded - weights['lm_head.weight'] = reshape( - split(lm_head_weights, mapping.tp_size, mapping.tp_rank, dim=0), - (vocab_size // mapping.tp_size, hidden_size)) - - if config.has_model_final_layernorm: - weights['transformer.ln_f.weight'] = params[ - f'{hf_param_prefix}.layer_norm.weight'].clone() - weights['transformer.ln_f.bias'] = params[ - f'{hf_param_prefix}.layer_norm.bias'].clone() - - return weights - - -def parse_pix2struct_config(args, hf_model): - # manually set q_scaling to offset attention scaling's effect. - # TODO: modify kernels to control whether to disable attention scaling - config = configparser.ConfigParser() - - def get_offset_q_scaling(config) -> str: - d_model = config.hidden_size - num_heads = config.num_heads - head_size = d_model / num_heads - scaling = 1 / head_size**.5 - return str(scaling) - - config["decoder"] = {} - for key, val in hf_model.decoder.config.to_dict().items(): - config["decoder"][key] = f"{val}" - - config["decoder"]["q_scaling"] = get_offset_q_scaling( - hf_model.decoder.config) - - config["structure"] = dict() - config["structure"]["pix2struct_with_bias"] = "false" - config["structure"]["use_gated_activation"] = "false" - config["structure"]["position_embedding_type"] = "relative" - config["structure"]["model_type"] = args.model_type - - def parse_pix2struct_config_by_component(config, component, args): - if component == 'decoder': - 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, 'hidden_size') - 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', - fallback=512) - args.has_position_embedding = config.getboolean( - 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( - component, 'has_embedding_layernorm', fallback=False) - args.has_embedding_scale = config.getboolean(component, - 'has_embedding_scale', - 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) - 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=True) - args.layernorm_eps = config.getfloat(component, - 'layer_norm_epsilon') - args.layernorm_position = layernorm_position_map[config.get( - 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 = True - 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=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.rescale_before_lm_head = config.getboolean( - component, 'tie_word_embeddings' - ) # default is True (for T5), but False for Flan-T5 - args.encoder_hidden_size = config.getint('decoder', 'hidden_size') - args.encoder_num_heads = config.getint('decoder', 'num_heads') - args.encoder_head_size = config.getint('decoder', 'd_kv') - args.position_embedding_type = config.get( - 'structure', 'position_embedding_type') - args.decoder_start_token_id = config.getint( - 'decoder', 'decoder_start_token_id') - args.eos_token_id = config.getint('decoder', 'eos_token_id') - bos_token_id = config.get('decoder', 'bos_token_id') - # pix2struct does not have bos_token_id - args.bos_token_id = int( - bos_token_id) if bos_token_id != "None" else None - args.pad_token_id = config.getint('decoder', 'pad_token_id') - - else: - assert False, 'Unsupported component!' - return args - - decoder_args = parse_pix2struct_config_by_component(config, "decoder", args) - return None, decoder_args - - -def convert_pix2struct_weights_to_tllm_safetensors(config, component, params): - weights = {} - - mapping = config.mapping - - convert_weight_to_dtype(params, config.dtype) - hidden_size = config.hidden_size - ffn_hidden_size = config.intermediate_size - num_layers = config.num_hidden_layers - n_head = config.num_attention_heads - head_size = config.head_size - attention_hidden_size = n_head * head_size # head size * num_heads not necessarily equals hidden_dim, such as Flan-T5 - - hf_param_prefix = f'{component}' - trtllm_layer_name = f'transformer.layers' - trtllm_attn_layer_name = 'self_attention' - trtllm_attn_layernorm_name = 'self_attention_layernorm' - - def get_attn_module_name(component, layer, attn_type): - return f'{component}.layer.{int(layer)}.{attn_type}.attention' - - weights['transformer.vocab_embedding.weight'] = reshape( - params[f'{hf_param_prefix}.embed_tokens.weight'].clone(), - None) if not config.use_parallel_embedding else reshape( - split(params[f'{hf_param_prefix}.embed_tokens.weight'].clone(), - mapping.tp_size, mapping.tp_rank, 0), None) - - layers_range = mapping.pp_layers(num_layers) - for layer_idx in layers_range: - local_layer_idx = layer_idx - layers_range[0] - trtllm_layer_name_prefix = f'{trtllm_layer_name}.{local_layer_idx}' - hf_layer_name_prefix = f'{hf_param_prefix}.layer.{layer_idx}' - - hidden_layer_name_split = { - f'{hf_layer_name_prefix}.self_attention.attention.output.weight': { - "name": - f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}.dense.weight', - "shape": - (hidden_size, attention_hidden_size // mapping.tp_size), - "split_dim": -1 - }, - f'{hf_layer_name_prefix}.mlp.DenseReluDense.wo.weight': { - "name": f'{trtllm_layer_name_prefix}.mlp.proj.weight', - "shape": (hidden_size, ffn_hidden_size // mapping.tp_size), - "split_dim": -1 - }, - f'{hf_layer_name_prefix}.mlp.DenseReluDense.wi_0.weight': { - "name": f'{trtllm_layer_name_prefix}.mlp.fc.weight', - "shape": (ffn_hidden_size // mapping.tp_size, hidden_size), - "split_dim": 0 - }, - } - - hidden_layer_name_no_split = { - f'{hf_layer_name_prefix}.self_attention.layer_norm.weight': { - "name": - f'{trtllm_layer_name_prefix}.{trtllm_attn_layernorm_name}.weight', - "shape": None - }, - f'{hf_layer_name_prefix}.mlp.layer_norm.weight': { - "name": f'{trtllm_layer_name_prefix}.mlp_layernorm.weight', - "shape": None - }, - } - - if config.gated_act: - hidden_layer_name_split.update({ - f'{hf_layer_name_prefix}.mlp.DenseReluDense.wi_1.weight': { - "name": f'{trtllm_layer_name_prefix}.mlp.gate.weight', - "shape": (ffn_hidden_size // mapping.tp_size, hidden_size), - "split_dim": 0 - }, - }) - - hidden_layer_name_split.update({ - f'{hf_layer_name_prefix}.encoder_decoder_attention.attention.output.weight': - { - "name": - f'{trtllm_layer_name_prefix}.cross_attention.dense.weight', - "shape": - (hidden_size, attention_hidden_size // mapping.tp_size), - "split_dim": -1 - }, - }) - hidden_layer_name_no_split.update({ - f'{hf_layer_name_prefix}.encoder_decoder_attention.layer_norm.weight': - { - "name": - f'{trtllm_layer_name_prefix}.cross_attention_layernorm.weight', - "shape": None - }, - }) - self_attn_module_name = get_attn_module_name( - component, layer_idx, 'encoder_decoder_attention') - weights.update( - fuse_qkv_one_layer( - params, self_attn_module_name, - f'{trtllm_layer_name_prefix}.cross_attention', mapping.tp_size, - mapping.tp_rank, config.model_type, - (attention_hidden_size * 3 // mapping.tp_size, hidden_size), - None)) - - self_attn_module_name = get_attn_module_name(component, layer_idx, - 'self_attention') - weights.update( - fuse_qkv_one_layer( - params, self_attn_module_name, - f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}', - mapping.tp_size, mapping.tp_rank, config.model_type, - (attention_hidden_size * 3 // mapping.tp_size, hidden_size), - None)) - - weights[ - f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}.rel_attn_table'] = reshape( - split( - params[ - f'{component}.layer.0.self_attention.attention.relative_attention_bias.weight'] - .T, mapping.tp_size, mapping.tp_rank, 0), - (n_head // mapping.tp_size, config.num_buckets)) - - for hf_weight_name, weight_info in hidden_layer_name_split.items(): - if hf_weight_name in params.keys(): - weights[weight_info["name"]] = reshape( - split(params[hf_weight_name], - mapping.tp_size, - mapping.tp_rank, - dim=weight_info["split_dim"]), weight_info["shape"]) - for hf_weight_name, weight_info in hidden_layer_name_no_split.items(): - if hf_weight_name in params.keys(): - weights[weight_info["name"]] = reshape( - params[hf_weight_name].clone(), shape=weight_info["shape"]) - - weights[f'transformer.ln_f.weight'] = reshape( - params[f'{component}.final_layer_norm.weight'].clone(), None) - - weights['lm_head.weight'] = reshape( - split(params[f'{component}.lm_head.weight'], - mapping.tp_size, - mapping.tp_rank, - dim=0), (config.vocab_size // mapping.tp_size, hidden_size)) - if not config.use_implicit_relative_attention: - weights[f'rel_attn_table'] = reshape( - split( - params[ - f'{component}.layer.0.self_attention.attention.relative_attention_bias.weight'] - .T, mapping.tp_size, mapping.tp_rank, 0), - (n_head // mapping.tp_size, config.num_buckets)) - - return weights - - -def parse_language_adapter_config(args, model): - config = configparser.ConfigParser() - config.read(args.model_dir + "/config.ini") - # rename from "sinusoid" to "learned_absolute" - config["structure"]["position_embedding_type"] = "learned_absolute" - - locale_list = ['de_DE', 'en_US', 'es_ES', 'fr_FR'] # to be changed by user - - encoder_config = parse_nmt_config_by_component(config, "encoder", args) - encoder_config.residual_scaling = config.getfloat("encoder", - 'residual_scaling', - fallback=1.0) - encoder_config.language_adapter_config = LanguageAdapterConfig( - num_languages=config.getint("encoder", 'adapter_langs', fallback=None), - ffn_hidden_size=config.getint("encoder", - 'encoder_adapter_embed_dim', - fallback=None), - language_list=locale_list).to_dict() - - decoder_config = parse_nmt_config_by_component(config, "decoder", args) - decoder_config.residual_scaling = config.getfloat("decoder", - 'residual_scaling', - fallback=1.0) - decoder_config.language_adapter_config = LanguageAdapterConfig( - num_languages=config.getint("decoder", 'adapter_langs', fallback=None), - ffn_hidden_size=config.getint("decoder", - 'decoder_adapter_embed_dim', - fallback=None), - language_list=locale_list).to_dict() - - decoder_config.decoder_start_token_id = 2 - decoder_config.eos_token_id = 2 - decoder_config.bos_token_id = 2 - decoder_config.pad_token_id = 0 - - return encoder_config, decoder_config - - -def convert_language_adapter_weights_to_tllm_safetensors( - config, component, params): - weights = {} - - mapping = config.mapping - - convert_weight_to_dtype(params, config.dtype) - - param_prefix = f'{component}' - trtllm_layer_name = f'transformer.layers' - trtllm_attn_layer_name = 'attention' if component == 'encoder' else 'self_attention' - trtllm_attn_layernorm_name = 'self_attention_layernorm' if component == 'decoder' else 'attention_layernorm' - mlp_param_prefix = '' if f'{param_prefix}.0.fc1.weight' in params else 'mlp.' - - hidden_size = params[ - f'{param_prefix}.layers.0.self_attn.out_proj.weight'].shape[0] - ffn_hidden_size = params[ - f'{param_prefix}.layers.0.{mlp_param_prefix}fc1.weight'].shape[0] - - hidden_layer_name_split = { - 'self_attn.out_proj.weight': { - "name": f'{trtllm_attn_layer_name}.dense.weight', - "shape": (hidden_size, hidden_size // mapping.tp_size), - "split_dim": -1 - }, - f'{mlp_param_prefix}fc1.weight': { - "name": 'mlp.fc.weight', - "shape": (ffn_hidden_size // mapping.tp_size, hidden_size), - "split_dim": 0 - }, - f'{mlp_param_prefix}fc1.bias': { - "name": 'mlp.fc.bias', - "shape": (ffn_hidden_size // mapping.tp_size), - "split_dim": 0 - }, - f'{mlp_param_prefix}fc2.weight': { - "name": 'mlp.proj.weight', - "shape": (hidden_size, ffn_hidden_size // mapping.tp_size), - "split_dim": -1 - }, - } - - hidden_layer_name_no_split = { - 'self_attn.out_proj.bias': { - "name": f'{trtllm_attn_layer_name}.dense.bias', - "shape": (hidden_size) - }, - 'self_attn_layer_norm.weight': { - "name": f'{trtllm_attn_layernorm_name}.weight', - "shape": None - }, - 'self_attn_layer_norm.bias': { - "name": f'{trtllm_attn_layernorm_name}.bias', - "shape": None - }, - f'{mlp_param_prefix}fc2.bias': { - "name": 'mlp.proj.bias', - "shape": (hidden_size) - }, - 'final_layer_norm.weight': { - "name": 'mlp_layernorm.weight', - "shape": None - }, - 'final_layer_norm.bias': { - "name": 'mlp_layernorm.bias', - "shape": None - }, - } - - if component == "decoder": - hidden_layer_name_split.update({ - 'encoder_attn.out_proj.weight': { - "name": 'cross_attention.dense.weight', - "shape": (hidden_size, hidden_size // mapping.tp_size), - "split_dim": -1 - }, - }) - hidden_layer_name_no_split.update({ - 'encoder_attn.out_proj.bias': { - "name": 'cross_attention.dense.bias', - "shape": (hidden_size) - }, - 'encoder_attn_layer_norm.weight': { - "name": 'cross_attention_layernorm.weight', - "shape": None, - }, - 'encoder_attn_layer_norm.bias': { - "name": 'cross_attention_layernorm.bias', - "shape": None - }, - }) - - def get_attn_module_name(layer, attn_type): - return f'{param_prefix}.layers.{int(layer)}.{attn_type}' - - # support MostlyFreezedEmbedding in 5.5B model - embed_tokens_weight_name = f'{param_prefix}.embed_tokens.weight' - if embed_tokens_weight_name not in params: - embed_tokens_weight_name = f'{param_prefix}.embed_tokens.weight_' - - weights['transformer.vocab_embedding.weight'] = reshape( - params[embed_tokens_weight_name].clone(), - None) if not config.use_parallel_embedding else reshape( - split(params[embed_tokens_weight_name].clone(), mapping.tp_size, - mapping.tp_rank, 0), None) - - weights[ - 'transformer.position_embedding.weight'] = fairseq_sin_pos_embedding( - config.max_position_embeddings, - params[embed_tokens_weight_name].shape[1]) - - num_layers = config.num_hidden_layers - layers_range = mapping.pp_layers(num_layers) - - for layer_idx in layers_range: - local_layer_idx = layer_idx - layers_range[0] - hf_layer_name_prefix = f'{param_prefix}.layers.{layer_idx}' - trtllm_layer_name_prefix = f'{trtllm_layer_name}.{local_layer_idx}' - - for hf_weight_name, weight_info in hidden_layer_name_split.items(): - weights[ - f'{trtllm_layer_name_prefix}.{weight_info["name"]}'] = reshape( - split(params[f'{hf_layer_name_prefix}.{hf_weight_name}'], - mapping.tp_size, - mapping.tp_rank, - dim=weight_info["split_dim"]), weight_info["shape"]) - - for hf_weight_name, weight_info in hidden_layer_name_no_split.items(): - trtllm_layer_fullname = f'{trtllm_layer_name_prefix}.{weight_info["name"]}' - hf_layer_fullname = f'{hf_layer_name_prefix}.{hf_weight_name}' - weights[trtllm_layer_fullname] = reshape( - params[hf_layer_fullname].clone(), shape=weight_info["shape"]) - - self_attn_module_name = get_attn_module_name(layer_idx, 'self_attn') - weights.update( - fuse_qkv_one_layer( - params, self_attn_module_name, - f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}', - mapping.tp_size, mapping.tp_rank, config.model_type, - (hidden_size * 3 // mapping.tp_size, hidden_size), - (hidden_size * 3 // mapping.tp_size))) - - if component == 'decoder': - cross_attn_module_name = get_attn_module_name( - layer_idx, 'encoder_attn') - weights.update( - fuse_qkv_one_layer( - params, cross_attn_module_name, - f'{trtllm_layer_name_prefix}.cross_attention', - mapping.tp_size, mapping.tp_rank, config.model_type, - (hidden_size * 3 // mapping.tp_size, hidden_size), - (hidden_size * 3 // mapping.tp_size))) - assert len(config.language_adapter_config['language_list']) > 0 - - language_adapter_weights = defaultdict(list) - language_adapter_weight_info = { - 'fc1.weight': { - "name": f'{trtllm_layer_name_prefix}.adapter.layers.fc.weight', - "shape": None - }, - 'fc1.bias': { - "name": f'{trtllm_layer_name_prefix}.adapter.layers.fc.bias', - "shape": None - }, - 'fc2.weight': { - "name": - f'{trtllm_layer_name_prefix}.adapter.layers.proj.weight', - "shape": None - }, - 'fc2.bias': { - "name": f'{trtllm_layer_name_prefix}.adapter.layers.proj.bias', - "shape": None - }, - } - - for language in config.language_adapter_config['language_list']: - for key in language_adapter_weight_info.keys(): - language_adapter_weights[key].append(params[ - f'{param_prefix}.layers.{layer_idx}.adapter.{language}.{key}'] - .unsqueeze(0)) - - import torch - for key, weight_info in language_adapter_weight_info.items(): - weights[weight_info["name"]] = torch.cat( - language_adapter_weights[key], dim=0) - - weights[ - f'{trtllm_layer_name_prefix}.adapter_layer_norm.weight'] = params[ - f'{param_prefix}.layers.{layer_idx}.adapter_layer_norm.weight'] - weights[f'{trtllm_layer_name_prefix}.adapter_layer_norm.bias'] = params[ - f'{param_prefix}.layers.{layer_idx}.adapter_layer_norm.bias'] - - if component == 'decoder': - - # share_decoder_input_output_embed=True, output_proj = embed_tokens.transpose() - lm_head_weight_name = f'{param_prefix}.output_projection.weight' - if lm_head_weight_name not in params: - lm_head_weight_name = embed_tokens_weight_name - weights['lm_head.weight'] = reshape( - split(params[lm_head_weight_name], - mapping.tp_size, - mapping.tp_rank, - dim=0), (config.vocab_size // mapping.tp_size, hidden_size)) - - return weights - - -def get_model(args): - if args.model_type == "t5": - model = T5ForConditionalGeneration.from_pretrained(args.model_dir) - elif args.model_type == "nmt": - from fairseq.models.transformer import TransformerModel - model = TransformerModel.from_pretrained(args.model_dir) - elif args.model_type == "bart": - if args.nougat: - model = VisionEncoderDecoderModel.from_pretrained(args.model_dir) - model = model.get_decoder() - elif args.eclair_radio: - import torch - - class RadioWithNeck(torch.nn.Module): - - def __init__(self): - super().__init__() - - self.model_encoder = torch.hub.load("NVlabs/RADIO", - "radio_model", - version="radio_v2.5-h") - self.model_encoder.summary_idxs = torch.tensor(4) - - self.conv1 = torch.nn.Conv1d(1280, 1024, 1) - self.layer_norm1 = torch.nn.LayerNorm( - 1024, eps=1e-6, elementwise_affine=True) - self.conv2 = torch.nn.Conv2d(1024, - 1024, - kernel_size=(1, 4), - stride=(1, 4), - padding=0, - bias=False) - self.layer_norm2 = torch.nn.LayerNorm( - 1024, eps=1e-6, elementwise_affine=True) - - def forward(self, pixel_values): - _, feature = self.model_encoder(pixel_values) - output = self.conv1(feature.permute(0, 2, - 1)).permute(0, 2, 1) - output = self.layer_norm1(output).permute(0, 2, 1) - - b, d, _ = output.shape - h = pixel_values.shape[-2] // 16 - w = pixel_values.shape[-1] // 16 - output = self.conv2(output.reshape(b, d, h, w)) - output = output.flatten(-2, -1).permute(0, 2, 1) - output = self.layer_norm2(output) - return output - - def get_processor(): - processor = NougatProcessor.from_pretrained( - "facebook/nougat-base") - - special_tokens = { - "output_plain_index": "", - "output_markdown_index": "", - "output_no_text_index": "", - "output_ocr_index": "", - "predict_bbox_index": "", - "no_bbox_index": "", - "bbox_start_index": "", # not used but can keep - # "bbox_end_index": "", # not used but can keep - "no_class_index": "", - "predict_classes_index": "", - } - for key, special_t in special_tokens.items(): - processor.tokenizer.add_special_tokens( - {"additional_special_tokens": [special_t]}) - setattr(processor.tokenizer, key, - processor.tokenizer.encode(special_t)[1]) - - # Add regular tokens for boxes - processor.tokenizer.add_tokens( - [f"" for x_i in range(1024)]) - processor.tokenizer.add_tokens( - [f"" for y_i in range(1280)]) - # Add regular tokens for classes - #"" - possible_classes = [ - "Text", "Title", "Section-header", "List-item", "TOC", - "Bibliography", "Footnote", "Page-header", "Page-footer", - "Picture", "Formula", "Page-number", "Table", "Caption" - ] - processor.tokenizer.add_tokens( - [f"" for cls in possible_classes]) - return processor - - processor = get_processor() - model = VisionEncoderDecoderModel.from_pretrained( - "facebook/nougat-base") - model.encoder = RadioWithNeck() - model.decoder.resize_token_embeddings(len(processor.tokenizer), - pad_to_multiple_of=64) - model.config.decoder_start_token_id = processor.tokenizer.eos_token_id # 2 - model.config.pad_token_id = processor.tokenizer.pad_token_id # 1 - from transformers.models.mbart.modeling_mbart import \ - MBartLearnedPositionalEmbedding - _, d_model = model.device, model.config.decoder.d_model - - with torch.inference_mode(): - # Inspect checkpoint shapes - safetensors.torch.load_model(model, - os.path.join( - args.model_dir, - "model.safetensors"), - strict=False) - model.decoder.model.decoder.embed_positions = MBartLearnedPositionalEmbedding( - ECLAIR_RADIO_MAX_POSITION_EMBEDDINGS, d_model) - model.decoder.model.decoder.embed_positions.weight.data.zero_() - model.decoder.model.decoder.embed_positions.weight.requires_grad_( - True) - model.decoder.lm_head.weight = model.decoder.get_input_embeddings( - ).weight - - model.eval() - model = model.get_decoder() - - else: - model = AutoModelForSeq2SeqLM.from_pretrained(args.model_dir) - elif args.model_type == "pix2struct": - model = Pix2StructForConditionalGeneration.from_pretrained( - args.model_dir) - elif args.model_type == "blip2": - model = Blip2ForConditionalGeneration.from_pretrained( - args.model_dir).language_model - elif args.model_type == "language_adapter": - import torch - - class DummyTorchModel: - - def __init__(self, model) -> None: - self.model = model - - def state_dict(self): - return self.model['model'] - - model = torch.load(args.model_dir + "/model.pt", weights_only=False) - return DummyTorchModel(model) - - return model - - -def convert_checkpoint(args): - - model = get_model(args) - - saved_dir = Path(args.output_dir) - saved_dir.mkdir(parents=True, exist_ok=True) - - encoder_saved_dir = saved_dir / "encoder" - encoder_saved_dir.mkdir(parents=True, exist_ok=True) - decoder_saved_dir = saved_dir / "decoder" - decoder_saved_dir.mkdir(parents=True, exist_ok=True) - - world_size = args.tp_size * args.pp_size - - kv_cache_quant_algo = None - quant_algo = None - - model_type = args.model_type if args.model_type != "blip2" else "t5" - parse_config_mapper = { - 't5': parse_t5_config, - 'pix2struct': parse_pix2struct_config, - 'blip2': parse_t5_config, # blip2 uses t5 config parser - 'language_adapter': parse_language_adapter_config, - 'nmt': parse_nmt_config, - 'bart': parse_bart_config, - } - encoder_config, decoder_config = parse_config_mapper[model_type](args, - model) - - additional_settings = ["gated_act"] - if model_type == 'language_adapter': - additional_settings += ["residual_scaling", "language_adapter_config"] - - if not (args.nougat - or args.eclair_radio) and args.model_type != "pix2struct": - tllm_encoder_config = { - 'architecture': "EncoderModel", - 'dtype': args.dtype, - 'logits_dtype': encoder_config.logits_dtype, - 'num_hidden_layers': encoder_config.n_layer, - 'num_attention_heads': encoder_config.n_head, - 'hidden_size': encoder_config.hidden_size, - 'norm_epsilon': encoder_config.layernorm_eps, - 'vocab_size': encoder_config.vocab_size, - 'position_embedding_type': encoder_config.position_embedding_type, - 'hidden_act': encoder_config.hidden_act, - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'max_position_embeddings': encoder_config.n_positions, - 'num_key_value_heads': encoder_config.n_head, - 'head_size': encoder_config.head_size, - 'has_position_embedding': encoder_config.has_position_embedding, - 'layernorm_type': encoder_config.layernorm_type, - 'has_attention_qkvo_bias': encoder_config.has_attention_qkvo_bias, - 'has_mlp_bias': encoder_config.has_mlp_bias, - 'has_model_final_layernorm': - encoder_config.has_model_final_layernorm, - 'has_embedding_layernorm': encoder_config.has_embedding_layernorm, - 'has_embedding_scale': encoder_config.has_embedding_scale, - 'intermediate_size': encoder_config.ffn_hidden_size, - 'q_scaling': encoder_config.q_scaling, - 'layernorm_position': encoder_config.layernorm_position, - 'mlp_type': encoder_config.mlp_type, - 'relative_attention': encoder_config.relative_attention, - 'max_distance': encoder_config.max_distance, - 'num_buckets': encoder_config.num_buckets, - 'model_type': encoder_config.model_type, - } - - for additional_setting in additional_settings: - if hasattr(encoder_config, additional_setting): - tllm_encoder_config.update({ - additional_setting: - getattr(encoder_config, additional_setting) - }) - - with (encoder_saved_dir / "config.json").open('w') as f: - json.dump(tllm_encoder_config, f, indent=4) - - encoder_convert_args = dict(params=model.state_dict(), - component="encoder") - tllm_decoder_config = { - 'architecture': "DecoderModel", - 'dtype': args.dtype, - 'logits_dtype': decoder_config.logits_dtype, - 'num_hidden_layers': decoder_config.n_layer, - 'num_attention_heads': decoder_config.n_head, - 'hidden_size': decoder_config.hidden_size, - 'norm_epsilon': decoder_config.layernorm_eps, - 'vocab_size': decoder_config.vocab_size, - 'position_embedding_type': decoder_config.position_embedding_type, - 'hidden_act': decoder_config.hidden_act, - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'max_position_embeddings': decoder_config.n_positions, - 'head_size': decoder_config.head_size, - 'has_position_embedding': decoder_config.has_position_embedding, - 'layernorm_type': decoder_config.layernorm_type, - 'has_attention_qkvo_bias': decoder_config.has_attention_qkvo_bias, - 'has_mlp_bias': decoder_config.has_mlp_bias, - 'has_model_final_layernorm': decoder_config.has_model_final_layernorm, - 'has_embedding_layernorm': decoder_config.has_embedding_layernorm, - 'has_embedding_scale': decoder_config.has_embedding_scale, - 'intermediate_size': decoder_config.ffn_hidden_size, - 'q_scaling': decoder_config.q_scaling, - 'layernorm_position': decoder_config.layernorm_position, - 'mlp_type': decoder_config.mlp_type, - 'relative_attention': decoder_config.relative_attention, - 'max_distance': decoder_config.max_distance, - 'num_buckets': decoder_config.num_buckets, - 'model_type': decoder_config.model_type, - 'rescale_before_lm_head': decoder_config.rescale_before_lm_head, - 'encoder_hidden_size': decoder_config.encoder_hidden_size, - 'encoder_num_heads': decoder_config.encoder_num_heads, - 'encoder_head_size': decoder_config.encoder_head_size, - 'skip_cross_kv': args.skip_cross_kv, - 'use_implicit_relative_attention': args.use_implicit_relative_attention, - 'decoder_start_token_id': decoder_config.decoder_start_token_id, - 'eos_token_id': decoder_config.eos_token_id, - 'bos_token_id': decoder_config.bos_token_id, - 'pad_token_id': decoder_config.pad_token_id, - } - for additional_setting in additional_settings: - if hasattr(decoder_config, additional_setting): - tllm_decoder_config.update({ - additional_setting: - getattr(decoder_config, additional_setting) - }) - - with (decoder_saved_dir / "config.json").open('w') as f: - json.dump(tllm_decoder_config, f, indent=4) - - decoder_convert_args = dict(params=model.state_dict(), component="decoder") - - if args.model_type == "nmt": - fairseq_config = vars(model.cfg.model) # Namespace --> dict - num_embeddings = fairseq_config['max_source_positions'] - embedding_dim = fairseq_config['encoder_embed_dim'] - padding_idx = model.models[0].encoder.embed_tokens.padding_idx # 1 - - sin_pos_embedding = model.models[ - 0].encoder.embed_positions.get_embedding( - padding_idx + 1 + num_embeddings, - embedding_dim, - padding_idx=padding_idx) # [2 + num_embeddings, embed_dim] - sin_pos_embedding = sin_pos_embedding[2:, :] # remove offset embeddings - - encoder_convert_args["sin_pos_embedding"] = sin_pos_embedding - decoder_convert_args["sin_pos_embedding"] = sin_pos_embedding - - if args.workers == 1: - if not (args.nougat - or args.eclair_radio) and args.model_type != "pix2struct": - convert(0, world_size, args, tllm_encoder_config, - encoder_convert_args, encoder_saved_dir) - convert(0, world_size, args, tllm_decoder_config, decoder_convert_args, - decoder_saved_dir) - else: - if args.workers > world_size: - args.workers = world_size - LOGGER.info(f'Convert checkpoint using {args.workers} workers.') - import torch.multiprocessing as mp - if not (args.nougat - or args.eclair_radio) and args.model_type != "pix2struct": - mp.spawn(convert, - nprocs=args.workers, - args=(world_size, args, tllm_encoder_config, - encoder_convert_args, encoder_saved_dir)) - mp.spawn(convert, - nprocs=args.workers, - args=(world_size, args, tllm_decoder_config, - decoder_convert_args, decoder_saved_dir)) - - -def convert(worker_rank, world_size, args, model_config, convert_args, - saved_dir): - for rank in range(worker_rank, world_size, args.workers): - rank_config = copy.deepcopy(PretrainedConfig.from_dict(model_config)) - rank_config.set_rank(rank) - weights = globals( - )[f'convert_{rank_config.model_type}_weights_to_tllm_safetensors']( - config=rank_config, **convert_args) - safetensors.torch.save_file(weights, - f'{saved_dir}/rank{rank}.safetensors') - - -if __name__ == "__main__": - emit_engine_arch_deprecation("convert_checkpoint.py") - parser = argparse.ArgumentParser( - formatter_class=argparse.RawTextHelpFormatter) - parser.add_argument( - '--model_type', - type=str, - default='t5', - choices=[ - 't5', 'nmt', 'bart', 'pix2struct', 'blip2', 'language_adapter' - ], - help= - 'Multimodal type when this script is used for multimodal conversion.') - - 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("--model_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( - "--workers", - type=int, - help="How many workers to spawn for conversion (default: 4)", - default=4) - parser.add_argument("--nougat", - action="store_true", - help="Model which uses vision encoder + mbart decoder") - parser.add_argument("--eclair_radio", - action="store_true", - help="Model which uses vision encoder + mbart decoder") - parser.add_argument("--verbose", - action="store_true", - help="Provide verbose messages") - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharding is only enabled when embedding_sharding_dim = 0' - ) - 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( - '--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( - '--skip_cross_kv', - action='store_true', - help= - 'Skip redundant cross qkv computation by using TensorRT IfConditional switch (experimental).' - ) - parser.add_argument( - '--use_implicit_relative_attention', - action='store_true', - help= - 'Compute relative attention bias on the fly instead of pre-compute a relative attention bias table.' - ) - 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/models/core/enc_dec/helper.py b/examples/models/core/enc_dec/helper.py deleted file mode 100755 index ed3628c1bbbd..000000000000 --- a/examples/models/core/enc_dec/helper.py +++ /dev/null @@ -1,121 +0,0 @@ -import math -import typing -from typing import Union - -import numpy as np -import torch # pytype: disable=import-error - -from tensorrt_llm._utils import str_dtype_to_torch - - -def split(v: Union[np.ndarray, torch.Tensor], - tp_size: int, - tp_rank: int, - dim=0): - if tp_size == 1: - if isinstance(v, np.ndarray): - return np.ascontiguousarray(v.copy()) - else: - return v.clone().detach() - 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: - 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 reshape(v: torch.Tensor, shape=None): - if shape is None: - return v.contiguous() - else: - return v.reshape(shape).contiguous() - - -def fuse_qkv_one_layer(params, attn_module_name, trtllm_layer_name, tp_size, - tp_rank, model_type, weight_shape, bias_shape): - - qkv_module_names = get_qkv_module_name(model_type) - - weight = {} - - # fuse weights of q, k, v - q_w = params[f'{attn_module_name}.{qkv_module_names["q"]}.weight'] - k_w = params[f'{attn_module_name}.{qkv_module_names["k"]}.weight'] - v_w = params[f'{attn_module_name}.{qkv_module_names["v"]}.weight'] - - # fuse qkv weight - shape = q_w.shape # (do, din) - qkv_w = torch.cat([q_w, k_w, v_w], - dim=0).reshape([3, shape[0], shape[1]]) # (3, do, din) - qkv_w = split(qkv_w, tp_size, tp_rank, dim=1) - weight[f'{trtllm_layer_name}.qkv.weight'] = reshape(qkv_w, - shape=weight_shape) - - # fuse qkv biases if present - if f'{attn_module_name}.{qkv_module_names["q"]}.bias' in params.keys( - ) and params[f'{attn_module_name}.{qkv_module_names["q"]}.bias'] is not None: - q_b = params[f'{attn_module_name}.{qkv_module_names["q"]}.bias'] - k_b = params[f'{attn_module_name}.{qkv_module_names["k"]}.bias'] - v_b = params[f'{attn_module_name}.{qkv_module_names["v"]}.bias'] - shape = q_b.shape[0] # (do,) - qkv_b = torch.cat([q_b, k_b, v_b], dim=0).reshape([3, shape]) # (3, do) - qkv_b = split(qkv_b, tp_size, tp_rank, dim=1) - weight[f'{trtllm_layer_name}.qkv.bias'] = reshape(qkv_b, - shape=bias_shape) - return weight - - -def get_qkv_module_name(model_type): - if model_type in ["t5", "blip2"]: - q = "q" - k = "k" - v = "v" - elif model_type == "bart" or model_type == "nmt" or model_type == "language_adapter": - q = "q_proj" - k = "k_proj" - v = "v_proj" - elif model_type == "pix2struct": - q = "query" - k = "key" - v = "value" - return {"q": q, "k": k, "v": v} - - -def convert_weight_to_dtype(params: typing.Dict[str, torch.Tensor], - dtype: typing.Optional[np.dtype] = None): - if dtype is not None: - assert isinstance(dtype, - str), f"dtype must be str, but get type {type(dtype)}" - for name in params.keys(): - params[name] = params[name].to(str_dtype_to_torch(dtype)) - - -def fairseq_sin_pos_embedding(num_embeddings: int, embedding_dim: int): - ''' - generate fairseq specific sinusoidal position embedding [sin, sin, ... cos, cos...] - https://github.com/facebookresearch/fairseq/blob/main/fairseq/modules/sinusoidal_positional_embedding.py - ''' - padding_offset = 2 - half_dim = embedding_dim // 2.0 - emb = math.log(10000) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=torch.float16) * -emb) - emb = torch.arange(num_embeddings + padding_offset, - dtype=torch.float16).unsqueeze(1) * emb.unsqueeze(0) - emb = torch.cat([torch.sin(emb), torch.cos(emb)], - dim=1).view(num_embeddings + padding_offset, -1) - if embedding_dim % 2 == 1: - # zero pad - emb = torch.cat( - [emb, torch.zeros(num_embeddings + padding_offset, 1)], - dim=1, - dtype=torch.float16) - ''' - remove first 2 column to match position_id setup difference between fairseq & trt - fairseq position_id starts with 2, ex: [2, 3, 4 ..] - trt position_id starts with 0 [0, 1, 2] - ''' - return emb[padding_offset:, :] diff --git a/examples/models/core/enc_dec/run.py b/examples/models/core/enc_dec/run.py deleted file mode 100644 index 8ae81960e6c5..000000000000 --- a/examples/models/core/enc_dec/run.py +++ /dev/null @@ -1,512 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import argparse -import json -import time - -import numpy as np -import torch -from transformers import (AutoConfig, AutoModelForSeq2SeqLM, AutoTokenizer, - BartForConditionalGeneration, - MBartForConditionalGeneration, - T5ForConditionalGeneration) - -import tensorrt_llm -from tensorrt_llm import logger -from tensorrt_llm.runtime import EncDecModelRunner - - -def print_tensor(tensor_name, tensor, num_elements=10): - if tensor.dtype in (torch.int32, torch.int64): - tensor = tensor.to(dtype=float) - print( - f'{tensor_name}: mean={tensor.abs().mean().item():.3f}, sum={tensor.abs().sum().item():.3f}, max={tensor.abs().max().item():.3f}' - ) - # Pass num_elements=-1 will print the whole tensor - if num_elements < 0: - num_elements = torch.numel(tensor) - print(f'{tensor.flatten()[:num_elements]}') - print("Tensor Shape: ", tensor.size()) - print("") - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument("--max_new_tokens", type=int, default=64) - parser.add_argument("--log_level", type=str, default="error") - parser.add_argument("--engine_dir", "-i", type=str, default="trt_engines") - parser.add_argument("--engine_name", type=str, default="enc_dec") - parser.add_argument("--model_name", - type=str, - help="HuggingFace model name or FairSeq model path", - 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", - help="Whether or not to turn on the debug mode", - action='store_true') - parser.add_argument("--compare_hf_fp32", - help="Compare results with HuggingFace FP32", - action='store_true') - parser.add_argument('--lora_dir', type=str, default=None, nargs="+") - parser.add_argument('--lora_task_uids', type=str, default=None, nargs="+") - parser.add_argument("--output_npy", - type=str, - default=None, - help="Store input/output tensors C++ runtime testing") - return parser.parse_args() - - -def test_fairseq_models(args): - ## Note: NMT is the only FairSeq model. Adding FairSeq dependency is too heavy for the CI workflow, hence we used fixed input/output ids for correctness check and leave FairSeq code in comments. Users can follow Encoder-Decoder's README to install FairSeq and test locally. - ''' - from fairseq.models.transformer import TransformerModel - - fairseq_model = TransformerModel.from_pretrained(model_name_or_path=args.model_name, data_name_or_path=args.model_name, bpe='subword_nmt', tokenizer='moses').cuda() - - input_text = "Good Morning! How are you doing today?" - input_ids = fairseq_model.encode(input_text) - - tik = time.time() - # Note: FairSeq sampling=True results are not deterministic, disable during accuracy check - fairseq_output_ids = fairseq_model.generate(input_ids, beam=1, sampling=False) # - tik = time.time() - - fairseq_output_ids = fairseq_output_ids[0]['tokens'] - fairseq_output_text = fairseq_model.decode(fairseq_output_ids) - - print("--------------------------------------") - print("input text: ", input_text) - print("input ids: ", input_ids) # [9938, 5384, 9328, 812, 3619, 53, 181, 3829, 1735, 171, 2] - print("fairseq_output ids: ", fairseq_output_ids) # [9804, 391, 4, 4625, 167, 25, 1003, 5123, 17, 167, 1466, 1234, 171, 2] - print("fairseq_output text: ", fairseq_output_text) # "Bonjour, Comment vous en tirez-vous aujourd'hui ?" - print(f"FairSeq E2E time {(tok-tik)*1000}ms") - print("--------------------------------------") - ''' - - max_new_tokens = args.max_new_tokens - bos_token_id = 2 - pad_token_id = 0 - eos_token_id = 2 - decoder_start_token_id = bos_token_id - - input_ids = torch.tensor( - [9938, 5384, 9328, 812, 3619, 53, 181, 3829, 1735, 171, 2]) - fairseq_output_ids = torch.tensor( - [9804, 391, 4, 4625, 167, 25, 1003, 5123, 17, 167, 1466, 1234, 171, 2]) - input_ids = torch.tensor([input_ids.tolist()]).type(torch.IntTensor).cuda() - decoder_input_ids = torch.IntTensor([[decoder_start_token_id]]).cuda() - decoder_input_ids = decoder_input_ids.repeat((input_ids.shape[0], 1)) - - tllm_model = EncDecModelRunner.from_engine(args.engine_name, - args.engine_dir, - debug_mode=args.debug_mode) - - inference_dtype = tllm_model.encoder_model_config.dtype - - return_dict = False # when set return_dict=True, get outputs by key - tik = time.time() - tllm_output = tllm_model.generate( - encoder_input_ids=input_ids, - decoder_input_ids=decoder_input_ids, - max_new_tokens=max_new_tokens, - num_beams=args.num_beams, - bos_token_id=bos_token_id, - pad_token_id=pad_token_id, - eos_token_id=eos_token_id, - debug_mode=args.debug_mode, - ) - torch.cuda.synchronize() - tok = time.time() - - if return_dict: - tllm_output_ids = tllm_output['output_ids'] - else: - tllm_output_ids = tllm_output - - if tensorrt_llm.mpi_rank() == 0: - output_ids = tllm_output_ids[:, 0, :] - output_ids = output_ids[output_ids != eos_token_id] - fairseq_output_ids = fairseq_output_ids[fairseq_output_ids != - eos_token_id] - - print("--------------------------------------") - print("TRT-LLM output_ids: ", output_ids) - print(f"TRT-LLM E2E time {(tok-tik)*1000}ms") - print("Precision:", inference_dtype) - print("--------------------------------------") - - assert output_ids.tolist() == fairseq_output_ids.tolist( - ), f"TRT-LLM output ids {output_ids} does not match Fairseq ids {fairseq_output_ids}" - - -def test_language_adapter_models(args): - # TRT-LLM runtime - tllm_model = EncDecModelRunner.from_engine(args.engine_name, - args.engine_dir, - debug_mode=args.debug_mode) - inference_dtype = tllm_model.encoder_model_config.dtype - - tokenized_inputs = [[ - 34901, 3048, 3011, 123250, 9517, 3018, 45732, 3048, 3003, 6553, 3781, - 383416, 33356, 3032, 97339, 3382, 3003, 19143, 3022, 169460, 3001, - 87966, 35848, 2996, 3002, 6358, 7387, 25864, 3032, 3011, 4570, 3022, - 7235, 182168, 2992, 3003, 2991, 39861, 2997, 26629, 98419, 5339, 2993, - 423511, 2544, 2 - ], - [ - 34901, 3048, 3011, 123250, 9517, 3018, 45732, 3048, - 3003, 6553, 3781, 383416, 33356, 3032, 97339, 3382, - 3003, 19143, 3022, 169460, 3001, 87966, 35848, 2996, - 3002, 6358, 7387, 25864, 3032, 3011, 4570, 3022, - 7235, 182168, 2992, 3003, 2991, 39861, 2997, 26629, - 98419, 5339, 2993, 423512, 2712, 2 - ]] - language_task_uids = [2, 3] - - target_outputs = [[ - 4094, 82383, 3501, 3073, 12672, 3535, 45217, 3018, 45732, 3158, 3116, - 400231, 3010, 7212, 12398, 52837, 3046, 391725, 3164, 3116, 40625, 2994, - 204507, 3001, 402406, 35848, 2996, 3002, 3003, 8317, 2994, 3007, 80104, - 55333, 3046, 3073, 4755, 2994, 7235, 182168, 2992, 3030, 4005, 2994, - 63261, 60932, 3010, 2991, 39861, 2993 - ], - [ - 62366, 3099, 14803, 3056, 9517, 3056, 3495, 36942, - 3975, 292422, 3262, 3315, 3010, 53857, 41472, 9823, - 3010, 6493, 26179, 151498, 3062, 286084, 3453, 3315, - 45059, 2994, 286488, 3001, 53771, 16240, 35848, 2996, - 3002, 22161, 3072, 3315, 25864, 51019, 3062, 3072, - 3063, 2999, 10657, 2994, 7235, 182168, 2992, 3030, - 7109, 3077, 2999, 109181, 51563, 3366, 2991, 39861, - 2993 - ]] - - max_new_tokens = args.max_new_tokens - input_ids = torch.IntTensor(tokenized_inputs) - - with open(f"{args.engine_dir}/decoder/config.json", "r") as f: - model_config = json.load(f) - decoder_start_token_id = model_config['pretrained_config'][ - 'decoder_start_token_id'] - decoder_input_ids = torch.IntTensor([[decoder_start_token_id]]).to('cuda') - decoder_input_ids = decoder_input_ids.repeat((input_ids.shape[0], 1)) - - def get_language_adapter_routings(language_uids, input_ids): - language_adapter_routing_masks = torch.tensor(language_uids, - dtype=torch.int32) - language_adapter_routings = [] - - for i, input_id in enumerate(input_ids): - mask = language_adapter_routing_masks[i].repeat(len(input_id), 1) - language_adapter_routings.append(mask) - - return torch.cat(language_adapter_routings) - - encoder_language_adapter_routings = get_language_adapter_routings( - language_task_uids, input_ids) - decoder_language_adapter_routings = get_language_adapter_routings( - language_task_uids, decoder_input_ids) - - bos_token_id = 2 - pad_token_id = 0 - eos_token_id = 2 - - return_dict = True # when set return_dict=True, get outputs by key - tik = time.time() - tllm_output = tllm_model.generate( - encoder_input_ids=input_ids, - decoder_input_ids=decoder_input_ids, - max_new_tokens=max_new_tokens, - num_beams=args.num_beams, - bos_token_id=bos_token_id, - pad_token_id=pad_token_id, - eos_token_id=eos_token_id, - debug_mode=args.debug_mode, - return_dict=return_dict, - encoder_language_adapter_routings=encoder_language_adapter_routings, - decoder_language_adapter_routings=decoder_language_adapter_routings, - return_encoder_output=args.output_npy and tensorrt_llm.mpi_rank() == 0) - torch.cuda.synchronize() - tok = time.time() - - if tensorrt_llm.mpi_rank() == 0: - tllm_output_ids = tllm_output['output_ids'] - - output_ids = tllm_output_ids[:, 0, :] - output_ids_list = [ - output_id[output_id != eos_token_id].tolist() - for output_id in output_ids - ] - - decoder_input_lengths = (decoder_input_ids != pad_token_id).sum(dim=1) - output_gen_lengths = (output_ids != eos_token_id).sum( - dim=1) - decoder_input_lengths - - print( - f"------ TRT-LLM beam = {args.num_beams} --------------------------------" - ) - if 'encoder_output' in tllm_output: - encoder_output = tllm_output['encoder_output'] - print_tensor('TRT-LLM encoder_output:', encoder_output) - print("TRT-LLM output_ids: ", output_ids) - print("TRT-LLM output generated lengths: ", output_gen_lengths) - print(f"TRT-LLM E2E time {(tok-tik)*1000}ms") - print("Precision:", inference_dtype) - print("--------------------------------------") - - assert output_ids_list == target_outputs, f"TRT-LLM output ids {output_ids_list} does not match Fairseq ids {target_outputs}" - - if args.output_npy: - output_npy(args, tokenized_inputs, tllm_output, output_ids) - - -def output_npy(args, tokenized_inputs, tllm_output, output_ids): - os.makedirs(args.output_npy, exist_ok=True) - - if hasattr(tokenized_inputs, "attention_mask"): - input_lengths = tokenized_inputs.attention_mask.sum(dim=1).type( - torch.IntTensor) - input_ids = tokenized_inputs.input_ids.type(torch.IntTensor) - else: - input_lengths = torch.IntTensor( - [len(input_ids) for input_ids in tokenized_inputs]) - input_ids = torch.IntTensor(tokenized_inputs) - - input_ids_flatten = torch.cat( - [input_ids[i][:input_lengths[i]] for i in range(len(input_lengths))]) - encoder_output = tllm_output['encoder_output'].type(torch.float16) - - def save_npy(tensor, name): - np.save(os.path.join(args.output_npy, f'{name}.npy'), - tensor.cpu().numpy()) - - print( - f"Saving input/output tensors to {args.output_npy} for C++ runtime testing" - ) - save_npy(input_ids_flatten, 'input_ids') # [num_tokens] - save_npy(input_lengths, 'input_lengths') # [batch_size] - save_npy(encoder_output, 'encoder_output') # [num_tokens, hidden_size] - save_npy( - output_ids, f'output_ids_beam{args.num_beams}' - ) # [batch_size, max_output_tokens], max_output_tokens = decoder_input_tokens + max_new_tokens - - -if __name__ == "__main__": - import os - - os.environ["TOKENIZERS_PARALLELISM"] = "false" - args = parse_arguments() - logger.set_level(args.log_level) - - # FairSeq NMT test logic is different from HuggingFace models - if 'wmt' in args.model_name: - test_fairseq_models(args) - exit() - - # language adapter test logic is different from HuggingFace models - if 'language_adapter' in args.engine_name: - test_language_adapter_models(args) - exit() - - test_remove_padding = True - if not test_remove_padding: - if 't5' in args.model_name: - 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." - elif 'bart' in args.model_name: - input_text = "The tower is 324 metres (1,063 ft) tall, about the same height as an 81-storey building, and the tallest structure in Paris. Its base is square, measuring 125 metres (410 ft) on each side. During its construction, the Eiffel Tower surpassed the Washington Monument to become the tallest man-made structure in the world, a title it held for 41 years until the Chrysler Building in New York City was finished in 1930. It was the first structure to reach a height of 300 metres. Due to the addition of a broadcasting aerial at the top of the tower in 1957, it is now taller than the Chrysler Building by 5.2 metres (17 ft). Excluding transmitters, the Eiffel Tower is the second tallest free-standing structure in France after the Millau Viaduct." - else: - raise RuntimeError('Unsupported model type!') - - else: - input_text = [ - "translate English to German: The house is wonderful.", - "summarize: I am a high-performance inference optimizer and runtime.", - "During its construction, the Eiffel Tower surpassed the Washington Monument to become the tallest man-made structure in the world", - ] - - # TRT-LLM runtime - tllm_model = EncDecModelRunner.from_engine(args.engine_name, - args.engine_dir, - args.lora_dir, - args.lora_task_uids, - debug_mode=args.debug_mode) - - inference_dtype = tllm_model.encoder_model_config.dtype - if inference_dtype == 'float32': - if "byt5" in args.model_name: - print( - "ByT5 models tokenize input by bytes instead of words, causing the input text in this example to be longer than the default value during build stage. Please adjust --max_input_len during trtllm-build to select the right length limit for ByT5 models." - ) - else: - input_text.append( - "Summarize this article in one sentence.\n\nKristine Watts (Molie Weeks) is broken apart, missing her lover; she is not able to overcome her love for him that is lost in the past. She hires a stranger (Douglas Davis) and gives a list of her mistakes to him with things to fix. But time is irreversible and sometimes the cure for the pain is a tragic end.\n\nThe first point that impresses in \"The Cure\" is the stylish cinematography that alternates black and white with color. The concise and sharp screenplay is capable to develop a tragic and bleak tale of love with an unexpected plot point in the very end in less than eight minutes. The soundtrack is beautiful but the volume is a little loud and associated to the fact that English is not my native language, in some moments I needed to repeat some words whispered by the narrator. The unknown lead actress has magnificent performance and is extremely gorgeous. I hope to have a chance to see her again on the screen. Last but not the least, the debut of the director and writer Ryan Jafri could not be better. My vote is nine.\n\nTitle (Brazil): Not Available", - ) - - tokenizer = AutoTokenizer.from_pretrained( - args.model_name) # TODO: use model path instead - 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("--------------------------------------") - - model_config = AutoConfig.from_pretrained(args.model_name) - - # start_id for decoder (could add more input_ids as forced_decoder_ids) - 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: - hf_model = AutoModelForSeq2SeqLM.from_pretrained( - args.model_name, # TODO: use model path instead - # dtype=torch.float16 if '16' in dtype else torch.float32, # TODO: use matched torch dtype - ).to('cuda').eval() # TODO: create config model path instead - assert type(hf_model) in ( - T5ForConditionalGeneration, BartForConditionalGeneration, - MBartForConditionalGeneration), 'Unsupported model!' - - if args.lora_dir is not None: - assert len(args.lora_dir - ) >= 1, "At least one lora model dir is required" - # we can only test single lora with HF - from peft import PeftModel - hf_model = PeftModel.from_pretrained( - hf_model, args.lora_dir[0]).to('cuda').eval() - - tik = time.time() - hf_gen_output = 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, - # control logits processors - no_repeat_ngram_size=0, # disable no repeat post-processor - forced_bos_token_id=None, # disable forced first/last token - forced_eos_token_id=None, - min_length=0, - # for debug - output_scores=True, - output_hidden_states=True, - return_dict_in_generate=True) - # get hf output scores - hf_output_ids = hf_gen_output.sequences - # convert to logits - 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( - f"------ HF beam = {args.num_beams} --------------------------------" - ) - 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("--------------------------------------") - - return_dict = True # when set return_dict=True, get outputs by key - tik = time.time() - tllm_output = tllm_model.generate( - encoder_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, - debug_mode=args.debug_mode, - return_dict=return_dict, - attention_mask=tokenized_inputs.attention_mask, - time_encoder=True, - return_encoder_output=args.output_npy and tensorrt_llm.mpi_rank() == 0) - torch.cuda.synchronize() - tok = time.time() - - if tensorrt_llm.mpi_rank() == 0: - if return_dict: - tllm_output_ids = tllm_output['output_ids'] - else: - tllm_output_ids = tllm_output - - 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( - f"------ TRT-LLM beam = {args.num_beams} --------------------------------" - ) - if 'encoder_output' in tllm_output: - encoder_output = tllm_output['encoder_output'] - print_tensor('TRT-LLM encoder_output:', encoder_output) - 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("--------------------------------------") - - # save input/output tensors for C++ runtime testing - if args.output_npy: - output_npy(args, tokenized_inputs, tllm_output, output_ids) - - # 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) - if inference_dtype != "float32": - print("") - print( - f"[CAVEAT] Comparing TRT-LLM {inference_dtype} results with HF float32 results. Close match are not expected!" - ) - assert match_rate > 0.8, f"Incorrect results! Match rate {match_rate}" - else: - 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}" - ) diff --git a/examples/models/core/gemma/README.md b/examples/models/core/gemma/README.md index 7c94f7ca52ff..1315755e85b4 100644 --- a/examples/models/core/gemma/README.md +++ b/examples/models/core/gemma/README.md @@ -1,786 +1,10 @@ -# Run Gemma on TensorRT-LLM +# Gemma (PyTorch Backend) -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. +Gemma 4 runs on the TensorRT LLM **PyTorch backend** — HuggingFace checkpoints are +loaded directly. The legacy TensorRT engine flow (`convert_checkpoint.py` / +`trtllm-build`) is no longer required. -## Table Of Contents - -- [Run Gemma on TensorRT-LLM](#run-gemma-on-tensorrt-llm) - - [Table Of Contents](#table-of-contents) - - [Support Matrix](#support-matrix) - - [Common scripts](#common-scripts) - - [Convert checkpoint](#convert-checkpoint) - - [Build engine](#build-engine) - - [Run inference](#run-inference) - - [Specific commands](#specific-commands) - - [Run Gemma 2B](#run-gemma-2b) - - [Run inference under bfloat16 for HF checkpoint](#run-inference-under-bfloat16-for-hf-checkpoint) - - [Run inference under FP8 for keras checkpoint](#run-inference-under-fp8-for-keras-checkpoint) - - [Run 2B inference under SmoothQuant for jax checkpoint](#run-2b-inference-under-smoothquant-for-jax-checkpoint) - - [Run inference under weight only for jax checkpoint](#run-inference-under-weight-only-for-jax-checkpoint) - - [Run inference under INT8 KV caches for jax checkpoint](#run-inference-under-int8-kv-caches-for-jax-checkpoint) - - [Run Gemma 7B](#run-gemma-7b) - - [Run inference under bfloat16 for torch checkpoint](#run-inference-under-bfloat16-for-torch-checkpoint) - - [Run inference under FP8 for jax checkpoint](#run-inference-under-fp8-for-jax-checkpoint) - - [Run 7B inference under SmoothQuant for jax checkpoint](#run-7b-inference-under-smoothquant-for-jax-checkpoint) - - [Run inference under weight only for keras checkpoint](#run-inference-under-weight-only-for-keras-checkpoint) - - [Run inference under INT8 KV caches for keras checkpoint](#run-inference-under-int8-kv-caches-for-keras-checkpoint) - - [Run Gemma 2](#run-gemma-2) - - [Run inference under bfloat16 for torch checkpoint](#run-inference-under-bfloat16-for-torch-checkpoint-1) - - [Run Gemma 3](#run-gemma-3) - - [Run inference under bfloat16 for HF checkpoint](#run-inference-under-bfloat16-for-hf-checkpoint-1) - - [Disaggregated Serving](#disaggregated-serving) - - [Dynamo](#dynamo) - - [Run Gemma 4](#run-gemma-4) - - [Serve with `trtllm-serve` (OpenAI-compatible API)](#serve-with-trtllm-serve-openai-compatible-api) - - [Accuracy evaluation with `trtllm-eval`](#accuracy-evaluation-with-trtllm-eval) - - [Run Modelopt Quantization](#run-modelopt-quantization) - - [Requirements](#requirements) - - [Quantize Checkpoints](#quantize-checkpoints) - - [Build Engines](#build-engines) - - [Accuracy Results on MMLU](#accuracy-results-on-mmlu) - -## Support Matrix - * FP32/FP16/BF16/INT8 Weight-Only/INT4 AWQ/SmoothQuant/FP8 - * For SmoothQuant, TRT-LLM only supports FP16 higher precision now. - * checkpoint type: Jax, Torch, Keras, Huggingface (HF) - * STRONGLY TYPED - * python runtime and triton backend - -## Common scripts - -### Convert checkpoint - -Please install required packages first: - -```bash -pip install -r requirements.txt -``` - -Users can use `convert_checkpoint.py` to convert the different source checkpoint to unified TensorRT LLM checkpoint format. Users could set `--dtype` to determine the inference data type, and set the quantization options like `--enable_fp8`, `--fp8_kv_cache` `--use_smooth_quant`, `--calibrate_kv_cache` (for INT8 kv cache) and `--use-weight-only-with-precision` (weight only). Users could also control the source checkpoint type by `--ckpt-type`. Currently, supported checkpoint types are `jax`, `torch` and `keras`. - -```bash -CKPT_PATH=/tmp/models/gemma_nv/checkpoints/tmp_2b_it -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_2b_it_tensorrt_llm/bf16/tp1/ - -python3 ./convert_checkpoint.py \ - --ckpt-type jax \ - --model-dir ${CKPT_PATH} \ - --dtype bfloat16 \ - --world-size 1 \ - --output-model-dir ${UNIFIED_CKPT_PATH} -``` - -### Build engine - -After getting checkpoint, we can use `trtllm-build` command to build TensorRT LLM engines from TensorRT LLM checkpoints. - -```bash -ENGINE_PATH=/tmp/gemma/2B/bf16/1-gpu/ -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 8 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} -``` - -### Run inference - -We provide three examples to run inference `run.py`, `summarize.py` and `mmlu.py`. `run.py` only run inference with `input_text` and show the output. - -`summarize.py` runs summarization on [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset and evaluate the model by [ROUGE](https://en.wikipedia.org/wiki/ROUGE_(metric)) scores and use the `ROUGE-1` score to validate the implementation. - -`mmlu.py` runs MMLU to evaluate the model by accuracy. - -Note that we need to download the dataset of MMLU first and the evaluation of MMLU requires more time. - -* run.py - -```bash -VOCAB_FILE_PATH=/tmp/models/gemma_nv/checkpoints/tmp_vocab.model -python3 ../../../run.py --engine_dir ${ENGINE_PATH} \ - --max_output_len 30 \ - --vocab_file ${VOCAB_FILE_PATH} - -[TensorRT-LLM] TensorRT LLM version: 0.9.0.dev2024020600Input [Text 0]: " Born in north-east France, Soyer trained as a" -Output [Text 0 Beam 0]: "chef in the renowned kitchens of Lyon. After honing his skills in various Michelin-starred establishments, he embarked on a solo venture, establishing his own restaurant" -``` - -* summarize.py - -```bash -python3 ../../../summarize.py --test_trt_llm \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 \ - --vocab_file ${VOCAB_FILE_PATH} - -[02/06/2024-10:08:54] [TRT-LLM] [I] TensorRT LLM (total latency: 3.2821836471557617 sec) -[02/06/2024-10:08:54] [TRT-LLM] [I] TensorRT LLM (total output tokens: 1989) -[02/06/2024-10:08:54] [TRT-LLM] [I] TensorRT LLM (tokens per second: 605.9989975648089) -[02/06/2024-10:08:54] [TRT-LLM] [I] TensorRT LLM beam 0 result -[02/06/2024-10:08:55] [TRT-LLM] [I] rouge1 : 26.376388677070615 -[02/06/2024-10:08:55] [TRT-LLM] [I] rouge2 : 7.468157586877296 -[02/06/2024-10:08:55] [TRT-LLM] [I] rougeL : 17.953060795106556 -[02/06/2024-10:08:55] [TRT-LLM] [I] rougeLsum : 22.410938121151652 -``` - -* mmlu.py - -Download the dataset first - -```bash -mkdir data -wget https://people.eecs.berkeley.edu/~hendrycks/data.tar -O data/mmlu.tar -tar -xf data/mmlu.tar -C data -mv data/data data/mmlu -``` - -Evaluate on MMLU dataset. - -```bash -python3 ../../../mmlu.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} - -Average accuracy 0.358 - social sciences -Average accuracy 0.359 - other (business, health, misc.) -Average accuracy: 0.329 -``` - -## Specific commands - -In this section, we demonstrate the scripts to convert checkpoint, building engine and run inference on different settings. We will not demonstrate all combinations here because there are too many cases. We choose some important cases to demonstrate. - -### Run Gemma 2B - -#### Run inference under bfloat16 for HF checkpoint - -```bash -git clone git@hf.co:google/gemma-2b -CKPT_PATH=gemma-2b/ -UNIFIED_CKPT_PATH=/tmp/ckpt/hf/gemma/2b/1-gpu/ -ENGINE_PATH=/tmp/engines/gemma/2B/bf16/1-gpu/ -VOCAB_FILE_PATH=gemma-2b/ - -python3 ./convert_checkpoint.py \ - --ckpt-type hf \ - --model-dir ${CKPT_PATH} \ - --dtype bfloat16 \ - --world-size 1 \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 8 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --tokenizer_dir ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -[03/05/2024-02:24:39] [TRT-LLM] [I] TensorRT LLM (total latency: 3.0897433757781982 sec) -[03/05/2024-02:24:39] [TRT-LLM] [I] TensorRT LLM (total output tokens: 2141) -[03/05/2024-02:24:39] [TRT-LLM] [I] TensorRT LLM (tokens per second: 692.9378073221881) -[03/05/2024-02:24:39] [TRT-LLM] [I] TensorRT LLM beam 0 result -[03/05/2024-02:24:39] [TRT-LLM] [I] rouge1 : 21.042873132085678 -[03/05/2024-02:24:39] [TRT-LLM] [I] rouge2 : 6.322669223228836 -[03/05/2024-02:24:39] [TRT-LLM] [I] rougeL : 16.450116567540338 -[03/05/2024-02:24:39] [TRT-LLM] [I] rougeLsum : 18.836567173262736 -``` - -#### Run inference under FP8 for keras checkpoint - -WARNING: This way of running FP8 will introduce noticeable accuracy drop. To avoid that, use Modelopt quantization mentioned in this readme. - -In this example, we demonstrate how to run FP8 inference on Gemma. Note that `convert_checkpoint.py` only uses identity activation scales, so the accuracy might be little worse than higher precision in some cases, but it is still very good because we don't do any calibration. This also shows the stability of FP8 compared to INT8. - -```bash -git clone git@hf.co:google/gemma-2b-it-keras -GIT_LFS_SKIP_SMUDGE=1 git clone git@hf.co:google/gemma-2b-it-flax # clone tokenizer model -cd gemma-2b-it-flax -git lfs pull -I tokenizer.model - -CKPT_PATH=gemma-2b-it-keras -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_2b_en_tensorrt_llm/fp8/tp1/ -ENGINE_PATH=/tmp/gemma/2B/fp8/1-gpu/ -VOCAB_FILE_PATH=gemma-2b-it-flax/tokenizer.model - -python3 ./convert_checkpoint.py \ - --ckpt-type keras \ - --model-dir ${CKPT_PATH} \ - --dtype bfloat16 \ - --world-size 1 \ - --enable_fp8 \ - --fp8_kv_cache \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --max_batch_size 8 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -[02/08/2024-10:37:15] [TRT-LLM] [I] TensorRT LLM (total latency: 3.116227149963379 sec) -[02/08/2024-10:37:15] [TRT-LLM] [I] TensorRT LLM (total output tokens: 2419) -[02/08/2024-10:37:15] [TRT-LLM] [I] TensorRT LLM (tokens per second: 776.259201781368) -[02/08/2024-10:37:15] [TRT-LLM] [I] TensorRT LLM beam 0 result -[02/08/2024-10:37:15] [TRT-LLM] [I] rouge1 : 20.206082692133098 -[02/08/2024-10:37:15] [TRT-LLM] [I] rouge2 : 5.902141189518428 -[02/08/2024-10:37:15] [TRT-LLM] [I] rougeL : 15.403458457907643 -[02/08/2024-10:37:15] [TRT-LLM] [I] rougeLsum : 17.44535527417846 - -python3 ../mmlu.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} - -Average accuracy 0.390 - social sciences -Average accuracy 0.405 - other (business, health, misc.) -Average accuracy: 0.356 -``` - -#### Run 2B inference under SmoothQuant for jax checkpoint - -```bash -git clone git@hf.co:google/gemma-2b-it-flax -CKPT_PATH=gemma-2b-it-flax/2b-it/ -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_2b_it_tensorrt_llm/sq/tp1 -ENGINE_PATH=/tmp/gemma/2B/int8_sq/1-gpu/ -VOCAB_FILE_PATH=gemma-2b-it-flax/tokenizer.model - -python3 ./convert_checkpoint.py \ - --ckpt-type jax \ - --model-dir ${CKPT_PATH} \ - --dtype float16 \ - --use_smooth_quant_plugin 0.5 \ - --tokenizer_dir ${VOCAB_FILE_PATH} \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 8 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -[02/08/2024-04:42:06] [TRT-LLM] [I] TensorRT LLM (total latency: 3.460859775543213 sec) -[02/08/2024-04:42:06] [TRT-LLM] [I] TensorRT LLM (total output tokens: 1786) -[02/08/2024-04:42:06] [TRT-LLM] [I] TensorRT LLM (tokens per second: 516.0567361385428) -[02/08/2024-04:42:06] [TRT-LLM] [I] TensorRT LLM beam 0 result -[02/08/2024-04:42:06] [TRT-LLM] [I] rouge1 : 22.534044843245525 -[02/08/2024-04:42:06] [TRT-LLM] [I] rouge2 : 5.940093176022924 -[02/08/2024-04:42:06] [TRT-LLM] [I] rougeL : 16.258991712579736 -[02/08/2024-04:42:06] [TRT-LLM] [I] rougeLsum : 19.60977626046262 -``` - -#### Run inference under weight only for jax checkpoint - -Available precisions: `int8` and `int4` - -Note that `int4-weight-only` might not be able to keep the accuracies on all models. If users want to use int4 to run inference, we recommend using `int4_awq`. - -* `int8` - -```bash -git clone git@hf.co:google/gemma-2b-it-flax -CKPT_PATH=gemma-2b-it-flax/2b-it/ -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_2b_it_tensorrt_llm/w8_a16/tp1/ -ENGINE_PATH=/tmp/gemma/2B/w8_a16/1-gpu/ -VOCAB_FILE_PATH=gemma-2b-it-flax/tokenizer.model - -python3 ./convert_checkpoint.py \ - --ckpt-type jax \ - --model-dir ${CKPT_PATH} \ - --use-weight-only-with-precision int8 \ - --dtype bfloat16 \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 32 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -[02/08/2024-04:44:54] [TRT-LLM] [I] TensorRT LLM (total latency: 3.5987987518310547 sec) -[02/08/2024-04:44:54] [TRT-LLM] [I] TensorRT LLM (total output tokens: 1797) -[02/08/2024-04:44:54] [TRT-LLM] [I] TensorRT LLM (tokens per second: 499.3332842203787) -[02/08/2024-04:44:54] [TRT-LLM] [I] TensorRT LLM beam 0 result -[02/08/2024-04:44:54] [TRT-LLM] [I] rouge1 : 24.48521318679745 -[02/08/2024-04:44:54] [TRT-LLM] [I] rouge2 : 7.240543314565931 -[02/08/2024-04:44:54] [TRT-LLM] [I] rougeL : 17.857921729984078 -[02/08/2024-04:44:54] [TRT-LLM] [I] rougeLsum : 21.214162155642896 -``` - -#### Run inference under INT8 KV caches for jax checkpoint - -```bash -git clone git@hf.co:google/gemma-2b-it-flax -CKPT_PATH=gemma-2b-it-flax/2b-it/ -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_2b_it_tensorrt_llm/int8kv/tp1 -ENGINE_PATH=/tmp/gemma/2B/int8kv/1-gpu/ -VOCAB_FILE_PATH=gemma-2b-it-flax/tokenizer.model - -python3 ./convert_checkpoint.py \ - --ckpt-type jax \ - --model-dir ${CKPT_PATH} \ - --world-size 1 \ - --dtype bfloat16 \ - --calibrate_kv_cache \ - --tokenizer_dir ${VOCAB_FILE_PATH} \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 32 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -[02/08/2024-04:52:22] [TRT-LLM] [I] TensorRT LLM (total latency: 3.5348474979400635 sec) -[02/08/2024-04:52:22] [TRT-LLM] [I] TensorRT LLM (total output tokens: 1819) -[02/08/2024-04:52:22] [TRT-LLM] [I] TensorRT LLM (tokens per second: 514.5907994786265) -[02/08/2024-04:52:22] [TRT-LLM] [I] TensorRT LLM beam 0 result -[02/08/2024-04:52:22] [TRT-LLM] [I] rouge1 : 24.0397941580232 -[02/08/2024-04:52:22] [TRT-LLM] [I] rouge2 : 7.325311340360227 -[02/08/2024-04:52:22] [TRT-LLM] [I] rougeL : 17.54210044633271 -[02/08/2024-04:52:22] [TRT-LLM] [I] rougeLsum : 20.627861723682177 -``` - -### Run Gemma 7B - -#### Run inference under bfloat16 for torch checkpoint - -Since torch model does not have model config, we need to add it manually in `CKPT_PATH` with file name `config.json`. - -```bash -git clone git@hf.co:google/gemma-7b-pytorch - -CKPT_PATH=gemma-7b-pytorch/ -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_7b_it_tensorrt_llm/bf16/tp1/ -ENGINE_PATH=/tmp/gemma/7B/bf16/1-gpu/ -VOCAB_FILE_PATH=gemma-7b-pytorch/tokenizer.model - -python3 ./convert_checkpoint.py \ - --ckpt-type torch \ - --model-dir ${CKPT_PATH} \ - --dtype bfloat16 \ - --world-size 1 \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 8 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -python3 ../../../mmlu.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} - -Average accuracy 0.739 - social sciences -Average accuracy 0.697 - other (business, health, misc.) -Average accuracy: 0.630 -``` - -#### Run inference under FP8 for jax checkpoint - -WARNING: This way of running FP8 will introduce noticeable accuracy drop. To avoid that, use Modelopt quantization mentioned in this readme. - -In this example, we demonstrate how to run FP8 inference on Gemma. Note that `convert_checkpoint.py` only uses identity activation scales, so the accuracy might be little worse than higher precision in some cases, but it is still very good because we don't do any calibration. This also shows the stability of FP8 compared to INT8. - -```bash -CKPT_PATH=/tmp/models/gemma_nv/checkpoints/tmp_7b_it -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_7b_it_tensorrt_llm/fp8/tp1/ -ENGINE_PATH=/tmp/gemma/7B/fp8/1-gpu/ -VOCAB_FILE_PATH=/tmp/models/gemma_nv/checkpoints/tmp_vocab.model - -python3 ./convert_checkpoint.py \ - --ckpt-type jax \ - --model-dir ${CKPT_PATH} \ - --dtype bfloat16 \ - --world-size 1 \ - --enable_fp8 \ - --fp8_kv_cache \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 8 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -[02/08/2024-06:42:13] [TRT-LLM] [I] TensorRT LLM (total latency: 5.884302377700806 sec) -[02/08/2024-06:42:13] [TRT-LLM] [I] TensorRT LLM (total output tokens: 2694) -[02/08/2024-06:42:13] [TRT-LLM] [I] TensorRT LLM (tokens per second: 457.8282737830064) -[02/08/2024-06:42:13] [TRT-LLM] [I] TensorRT LLM beam 0 result -[02/08/2024-06:42:13] [TRT-LLM] [I] rouge1 : 27.18633861010837 -[02/08/2024-06:42:13] [TRT-LLM] [I] rouge2 : 7.734928823230158 -[02/08/2024-06:42:13] [TRT-LLM] [I] rougeL : 19.32537431798716 -[02/08/2024-06:42:13] [TRT-LLM] [I] rougeLsum : 22.82522575944535 -``` - -#### Run 7B inference under SmoothQuant for jax checkpoint - -```bash -git clone git@hf.co:google/gemma-7b-it-flax -CKPT_PATH=gemma-7b-it-flax/7b-it/ -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_7b_it_tensorrt_llm/sq/tp1 -ENGINE_PATH=/tmp/gemma/7B/int8_sq/1-gpu/ -VOCAB_FILE_PATH=gemma-7b-it-flax/tokenizer.model - -python3 ./convert_checkpoint.py \ - --ckpt-type jax \ - --model-dir ${CKPT_PATH} \ - --dtype float16 \ - --use_smooth_quant_plugin 0.5 \ - --tokenizer_dir ${VOCAB_FILE_PATH} \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 8 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -[02/19/2024-10:02:53] [TRT-LLM] [I] --------------------------------------------------------- -[02/19/2024-10:03:09] [TRT-LLM] [I] TensorRT LLM (total latency: 13.65670919418335 sec) -[02/19/2024-10:03:09] [TRT-LLM] [I] TensorRT LLM (total output tokens: 8351) -[02/19/2024-10:03:09] [TRT-LLM] [I] TensorRT LLM (tokens per second: 611.494312521266) -[02/19/2024-10:03:09] [TRT-LLM] [I] TensorRT LLM beam 0 result -[02/19/2024-10:03:09] [TRT-LLM] [I] rouge1 : 28.8107815115074 -[02/19/2024-10:03:09] [TRT-LLM] [I] rouge2 : 8.623835512061866 -[02/19/2024-10:03:09] [TRT-LLM] [I] rougeL : 19.7277195532959 -[02/19/2024-10:03:09] [TRT-LLM] [I] rougeLsum : 23.434950511855114 -``` - -#### Run inference under weight only for keras checkpoint - -Available precisions: `int8` and `int4` - -Note that `int4-weight-only` might not be able to keep the accuracies on all models. If users want to use int4 to run inference, we recommend using `int4_awq`. - -* `int8` - -```bash -git clone git@hf.co:google/gemma-7b-it-keras -GIT_LFS_SKIP_SMUDGE=1 git clone git@hf.co:google/gemma-7b-it-flax # clone tokenizer model -cd gemma-7b-it-flax -git lfs pull -I tokenizer.model - -CKPT_PATH=gemma-7b-it-keras -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_7b_it_tensorrt_llm/w8_a16/tp1/ -ENGINE_PATH=/tmp/gemma/7B/w8_a16/1-gpu/ -VOCAB_FILE_PATH=gemma-7b-it-flax/tokenizer.model - -python3 ./convert_checkpoint.py \ - --ckpt-type keras \ - --model-dir ${CKPT_PATH} \ - --use-weight-only-with-precision int8 \ - --dtype bfloat16 \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 32 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -[02/08/2024-07:38:15] [TRT-LLM] [I] TensorRT LLM (total latency: 8.49835753440857 sec) -[02/08/2024-07:38:15] [TRT-LLM] [I] TensorRT LLM (total output tokens: 2654) -[02/08/2024-07:38:15] [TRT-LLM] [I] TensorRT LLM (tokens per second: 312.2956393931832) -[02/08/2024-07:38:15] [TRT-LLM] [I] TensorRT LLM beam 0 result -[02/08/2024-07:38:16] [TRT-LLM] [I] rouge1 : 20.396209981234687 -[02/08/2024-07:38:16] [TRT-LLM] [I] rouge2 : 5.73302850102211 -[02/08/2024-07:38:16] [TRT-LLM] [I] rougeL : 16.001683776127507 -[02/08/2024-07:38:16] [TRT-LLM] [I] rougeLsum : 18.36957526315223 -``` - -#### Run inference under INT8 KV caches for keras checkpoint - -```bash -CKPT_PATH=/tmp/models/gemma_keras/keras/gemma_7b_en/ -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_7b_it_tensorrt_llm/int8kv/tp1 -ENGINE_PATH=/tmp/gemma/7B/int8kv/1-gpu/ -VOCAB_FILE_PATH=/tmp/models/gemma_nv/checkpoints/tmp_vocab.model - -python3 ./convert_checkpoint.py \ - --ckpt-type keras \ - --model-dir ${CKPT_PATH} \ - --world-size 1 \ - --dtype bfloat16 \ - --calibrate_kv_cache \ - --tokenizer_dir ${VOCAB_FILE_PATH} \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 32 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -[02/08/2024-07:51:11] [TRT-LLM] [I] TensorRT LLM (total latency: 8.73880124092102 sec) -[02/08/2024-07:51:11] [TRT-LLM] [I] TensorRT LLM (total output tokens: 2771) -[02/08/2024-07:51:11] [TRT-LLM] [I] TensorRT LLM (tokens per second: 317.09154649544956) -[02/08/2024-07:51:11] [TRT-LLM] [I] TensorRT LLM beam 0 result -[02/08/2024-07:51:11] [TRT-LLM] [I] rouge1 : 20.934864626327627 -[02/08/2024-07:51:11] [TRT-LLM] [I] rouge2 : 4.954721611692932 -[02/08/2024-07:51:11] [TRT-LLM] [I] rougeL : 15.307592049634444 -[02/08/2024-07:51:11] [TRT-LLM] [I] rougeLsum : 17.94213019528988 -``` - - -### Run Gemma 2 - -Gemma 2 currently has following limitations: - - Only HF style checkpoints are supported. - - The maximum sequence length allowed is 4096. -#### Run inference under bfloat16 for torch checkpoint -```bash -variant=9b # 27b -git clone git@hf.co:google/gemma-2-$variant-it - -CKPT_PATH=gemma-2-$variant-it/ -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_$variant_it_tensorrt_llm/bf16/tp1/ -ENGINE_PATH=/tmp/gemma2/$variant/bf16/1-gpu/ -VOCAB_FILE_PATH=gemma-2-$variant-it/tokenizer.model - -python3 ./examples/models/core/gemma/convert_checkpoint.py \ - --ckpt-type hf \ - --model-dir ${CKPT_PATH} \ - --dtype bfloat16 \ - --world-size 1 \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 8 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -python3 ../../../mmlu.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} - -Average accuracy 0.739 - social sciences -Average accuracy 0.697 - other (business, health, misc.) -Average accuracy: 0.630 -``` - -### Run Gemma 3 - -Gemma 3's text generation model from HuggingFace is supported. Gemma3 1B model interleaves 5 local layers between each global layer. While local layers use sliding window attention with a short span of 512 tokens, global layers attend to the long context. TRTLLM support layerwise sliding-window attention and the sliding window size for each layer could be passed in using the `--max_attention_window_size` parameter at runtime. If a subpattern is provided, TRTLLM can extrapolate the complete pattern and the extrapolation logic is printed to terminal. - -#### Run inference under bfloat16 for HF checkpoint -```bash -git clone https://huggingface.co/google/gemma-3-1b-it - -CKPT_PATH=gemma-3-1b-it/ -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_1b_it_tensorrt_llm/bf16/tp1/ -ENGINE_PATH=/tmp/gemma3/1b/bf16/1-gpu/ -VOCAB_FILE_PATH=gemma-3-1b-it/tokenizer.model - -python3 ./convert_checkpoint.py \ - --ckpt-type hf \ - --model-dir ${CKPT_PATH} \ - --dtype bfloat16 \ - --world-size 1 \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 8 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 1 \ - --max_ite 5 \ - --max_attention_window_size 512 512 512 512 512 3100 - -... -[TensorRT-LLM][INFO] TRTGptModel mMaxAttentionWindowSize: (512, 512, 512, 512, 512, 3100) * 4 + (512, 512) -... -[04/09/2025-18:28:26] [TRT-LLM] [I] TensorRT LLM (total latency: 1.6197962760925293 sec) -[04/09/2025-18:28:26] [TRT-LLM] [I] TensorRT LLM (total output tokens: 475) -[04/09/2025-18:28:26] [TRT-LLM] [I] TensorRT LLM (tokens per second: 293.2467539349165) -[04/09/2025-18:28:26] [TRT-LLM] [I] TensorRT LLM beam 0 result -[04/09/2025-18:28:26] [TRT-LLM] [I] rouge1: 22.780314381954003 -[04/09/2025-18:28:26] [TRT-LLM] [I] rouge2: 4.331099231480823 -[04/09/2025-18:28:26] [TRT-LLM] [I] rougeL: 15.26751867562475 -[04/09/2025-18:28:26] [TRT-LLM] [I] rougeLsum: 20.14696930976001 -``` - -#### Disaggregated Serving - -To serve the model in disaggregated mode, you should launch context and generation servers using `trtllm-serve`. - -For example, you can launch a single context server on port 8001 with: - -```bash -export TRTLLM_USE_UCX_KVCACHE=1 - -cat >./ctx_config.yml < output_ctx_8001 & -``` - -Then launch a single generation server on port 8002 with: - -```bash -cat >./gen_config.yml < output_gen_8002 & -``` - -Finally, you can launch the disaggregated server which will accept requests from the client and do -the orchestration between the context and generation servers with: - -```bash -cat >./disagg-config.yaml < argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument("--ckpt-type", - type=CheckpointType, - choices=list(CheckpointType)) - parser.add_argument("--model-dir", "--model_dir", type=Path, required=True) - parser.add_argument("--output-model-dir", - "--output_dir", - type=Path, - required=True) - parser.add_argument("--world-size", - type=int, - default=1, - help="world size, only support tensor parallelism now") - 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', - "--use-weight-only-with-precision", - choices=["int8", "int4", "w4a8_awq", "w4a16_awq"], - help= - "help='Quantize weights for the various GEMMs to INT4/INT8. Define the precision for the weights.", - ) - parser.add_argument( - "--use-int8-weight-only-embedding", - action="store_true", - help= - "Use weight only on embedding table and lm_head. (Only supported on Hopper GPU)", - ) - parser.add_argument( - '--dtype', - type=str, - default='auto', - choices=['auto', 'float16', 'bfloat16', 'float32'], - help= - "The data type for the model weights and activations if not quantized. " - "If 'auto', the data type is automatically inferred from the source model; " - "however, if the source dtype is float32, it is converted to float16.") - parser.add_argument( - "--enable_fp8", - action="store_true", - help="Use FP8 Linear layer for Attention QKV/Dense and MLP.") - parser.add_argument( - "--fp8_kv_cache", - action="store_true", - help= - "By default, we use dtype for KV cache. fp8_kv_cache chooses fp8 quantization for KV", - ) - parser.add_argument( - "--quant_ckpt_path", - default=None, - help= - "Path of a directory to quantized model checkpoints in .safetensors format or \ - path of a quantized model checkpoint in .npz format") - parser.add_argument( - '--calib_dataset', - type=str, - default='ccdv/cnn_dailymail', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - parser.add_argument('--use_smooth_quant', - action="store_true", - help="Use smooth quant.") - parser.add_argument( - "--int8_kv_cache", - "--calibrate_kv_cache", - "-kv", - action="store_true", - help= - "Generate scaling factors for KV cache. Used for storing KV cache in int8." - ) - parser.add_argument( - '--per_channel', - 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', - 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( - "--smoothquant", - "--use_smooth_quant_plugin", - "-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( - '--tokenizer_dir', - default=None, - help='tokenizer path; defaults to jax_model_dir if left unspecified') - parser.add_argument("--load_model_on_cpu", action="store_true") - args = parser.parse_args() - - if args.use_weight_only: - assert args.weight_only_precision is not None - return args - - -CKPT_PARSER: Dict[CheckpointType, Type[Parsers]] = { - CheckpointType.jax: JAXParser, - CheckpointType.keras: KerasParser, - CheckpointType.torch: TorchParser, - CheckpointType.hf: HfParser -} - - -def compute_quant_algo(args: argparse.Namespace) -> Optional[QuantAlgo]: - if args.weight_only_precision: - return { - "int8": QuantAlgo.W8A16, - "int4": QuantAlgo.W4A16, - "w4a8_awq": QuantAlgo.W4A8_AWQ, - "w4a16_awq": QuantAlgo.W4A16_AWQ, - }[args.weight_only_precision] - elif args.enable_fp8: - return QuantAlgo.FP8 - if args.use_smooth_quant: - return QuantAlgo.W8A8_SQ_PER_CHANNEL - elif args.smoothquant is not None: - if args.per_token and args.per_channel: - return QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN - elif not args.per_token and not args.per_channel: - return QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN - elif not args.per_token and args.per_channel: - return QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN - elif args.per_token and not args.per_channel: - return QuantAlgo.W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN - return None - - -def create_quant_config(args: argparse.Namespace) -> QuantConfig: - quant_algo = compute_quant_algo(args) - GemmaForCausalLM.assert_valid_quant_algo(quant_algo) - quant_config = QuantConfig(quant_algo=quant_algo) - if args.smoothquant is not None: - quant_config.smoothquant_val = args.smoothquant - - if args.fp8_kv_cache: - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - if args.int8_kv_cache: - quant_config.kv_cache_quant_algo = QuantAlgo.INT8 - - if args.weight_only_precision: - use_awq = args.weight_only_precision.endswith("awq") - use_int4 = args.weight_only_precision.endswith("int4") - if use_awq: - quant_config.group_size = 128 - - if use_awq or use_int4 or not args.use_int8_weight_only_embedding: - quant_config.has_zero_point = False - quant_config.pre_quant_scale = True - else: - quant_config.exclude_modules = ['router'] - - return quant_config - - -def main() -> None: - emit_engine_arch_deprecation("convert_checkpoint.py") - args = parse_arguments() - tik = time.time() - quant_config = create_quant_config(args) - - ckpt_parser = CKPT_PARSER[args.ckpt_type]() - - mapping = Mapping( - world_size=args.world_size, - tp_size=args.world_size, - pp_size=1, - ) - """We don't support pipeline parallelism yet for Gemma.""" - - if isinstance(ckpt_parser, HfParser): - trt_llm_config = GemmaConfig.from_hugging_face( - args.model_dir, - args.dtype, - mapping=mapping, - quant_config=quant_config, - ) - else: - print(f"Loading source parameters from {args.model_dir.absolute()}") - ckpt_params = ckpt_parser.load_parameters(args.model_dir) - input_embedding_weights = ckpt_parser.embedding_weights(ckpt_params) - num_embed, _ = input_embedding_weights.shape - ckpt_params_dtype = str(input_embedding_weights.dtype).split(".")[ - -1] # np.bfloat16 -> bfloat16 - ckpt_config = ckpt_parser.get_config(args.model_dir, ckpt_params, - num_embed) - # 2B TransformerConfig(num_layers=18, num_embed=256128, embed_dim=2048, hidden_dim=16384, num_heads=8, head_dim=256, num_kv_heads=1) - # 7B TransformerConfig(...) - - del ckpt_params - - print(f"Source configuration determined from parameters: {ckpt_config}") - - trt_llm_config = tensorrt_llm.models.GemmaConfig( - architecture=GEMMA_ARCHITECTURE, - dtype=infer_dtype(args.dtype, ckpt_params_dtype), - logits_dtype="float32", - vocab_size=ckpt_config.num_embed, - max_position_embeddings=8192, - hidden_size=ckpt_config.embed_dim, - num_hidden_layers=ckpt_config.num_layers, - num_attention_heads=ckpt_config.num_heads, - num_key_value_heads=ckpt_config.num_kv_heads, - head_size=ckpt_config.head_dim, - hidden_act="gelu", - intermediate_size=ckpt_config.hidden_dim, - norm_epsilon=1e-6, # hard-coded in RMSNorm from gemma/layers.py - position_embedding_type="rope_gpt_neox", - mapping=mapping, - gpus_per_node=8, - quantization=quant_config, - use_parallel_embedding=mapping.tp_size > 1, - ) - if hasattr(ckpt_config, - "model_type") and ckpt_config.model_type == "gemma2": - trt_llm_config.inter_layernorms = True - trt_llm_config.final_logit_softcapping = ckpt_config.final_logit_softcapping - trt_llm_config.attn_logit_softcapping = ckpt_config.attn_logit_softcapping - trt_llm_config.query_pre_attn_scalar = ckpt_config.query_pre_attn_scalar - - trt_llm_config_dict = trt_llm_config.to_dict() - print(f"Determined TensorRT LLM configuration {trt_llm_config_dict}") - - save_config(trt_llm_config, output_dir=args.output_model_dir, log=True) - - for config in trt_llm_config.for_each_rank(): - hf_weights = load_gemma_weights( - parameters_or_model_dir=args.model_dir, - trt_llm_config=config, - ckpt_parser=ckpt_parser, - load_model_on_cpu=args.load_model_on_cpu) - ranked_weights = non_modelopt_quantize_if_needed( - hf_weights, - model_dir=args.model_dir, - quantize_modifiers=QuantizeModifiers.from_args(args), - trt_llm_config=config) - save_checkpoint(output_dir=args.output_model_dir, - weights=ranked_weights, - rank=config.mapping.rank) - - elapsed = time.strftime("%H:%M:%S", time.gmtime(time.time() - tik)) - print(f"Total time of converting checkpoints: {elapsed}") - - -if __name__ == "__main__": - main() diff --git a/examples/models/core/gemma/requirements.txt b/examples/models/core/gemma/requirements.txt deleted file mode 100644 index a1bbed25b68a..000000000000 --- a/examples/models/core/gemma/requirements.txt +++ /dev/null @@ -1,14 +0,0 @@ --f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html --c ../../../constraints.txt -# WAR the new posting of "nvidia-cudnn-cu12~=9.0". -# "jax[cuda12_pip]~=0.4.19" specifies "nvidia-cudnn-cu12>=8.9" but actually requires "nvidia-cudnn-cu12~=8.9". -nvidia-cudnn-cu12~=8.9; platform_machine == "x86_64" -tensorrt_llm>=0.0.0.dev0 -flax~=0.8.0 -# jax[cuda12_pip]~=0.4.19 -safetensors~=0.4.1 -sentencepiece>=0.1.99 -h5py~=3.12.1 -rouge_score -nltk -datasets==3.1.0 diff --git a/examples/models/core/glm-4-9b/README.md b/examples/models/core/glm-4-9b/README.md deleted file mode 100644 index b2789d59f92f..000000000000 --- a/examples/models/core/glm-4-9b/README.md +++ /dev/null @@ -1,292 +0,0 @@ -# ChatGLM - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document explains how to build the [glm-4-9b](https://huggingface.co/THUDM/glm-4-9b) models using TensorRT LLM and run on a single GPU, a single node with multiple GPUs or multiple nodes with multiple GPUs. - -- [glm-4-9b](#glm-4-9b) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Model comparison](#model-comparison) - - [Tokenizer and special tokens comparison](#tokenizer-and-special-tokens-comparison) - - [Usage](#usage) - - [1. Download repo and weights from HuggingFace Transformers](#1-download-repo-and-weights-from-huggingface-transformers) - - [2. Convert weights from HF Transformers to TensorRT LLM format](#2-convert-weights-from-hf-transformers-to-tensorrt-llm-format) - - [3. Build TensorRT engine(s)](#3-build-tensorrt-engines) - - [Enable plugins](#enable-plugins) - - [In-flight batching](#in-flight-batching) - - [4. Run inference](#4-run-inference) - - [Single node, single GPU](#single-node-single-gpu) - - [Single node, multi GPU](#single-node-multi-gpu) - - [5. Run summarization task](#5-run-summarization-task) - - [Weight Only quantization](#weight-only-quantization) - - [Smooth Quantization (SQ)](#smooth-quantization-sq) - - [Activation-aware Weight Quantization (AWQ)](#activation-aware-weight-quantization-awq) - - [FP8 Quantization](#fp8-quantization) - - [Benchmark](#benchmark) - - -## Overview - -The TensorRT LLM ChatGLM implementation can be found in [`tensorrt_llm/models/chatglm/model.py`](../../../../tensorrt_llm/models/chatglm/model.py). -The TensorRT LLM ChatGLM example code is located in [`examples/models/core/glm-4-9b`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`run.py`](../../../run.py) to run the inference on an input text; -* [`summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - -| Model Name | FP16 | FMHA | WO | SQ | AWQ | FP8 | TP | PP | ST | C++ | benchmark | IFB | -| :--------------: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :-------: | :---: | -| glm_4_9b | Y | Y | Y | | | | Y | | | Y | | | -| glm_4_9b_chat | Y | Y | Y | | | | Y | | | Y | | | - -* Model Name: the name of the model, the same as the name on HuggingFace -* FMHA: Fused MultiHead Attention (see introduction below) -* WO: Weight Only Quantization (int8 / int4) -* SQ: Smooth Quantization (int8) -* AWQ: Activation Aware Weight Quantization (int4) -* FP8: FP8 Quantization -* TP: Tensor Parallel -* PP: Pipeline Parallel -* ST: Strongly Typed -* C++: C++ Runtime -* benchmark: benchmark by python / C++ Runtime -* IFB: In-flight Batching (see introduction below) - -## Model comparison - -| Name | nL | nAH | nKH | nHW | nH | nF | nMSL | nV | bP2D | bBQKV | bBDense | Comments | -| :--------------: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :----: | :---: | :---: | :-----: | :----------------------------------------------------------------- | -| glm_4_9b | 40 | 32 | 2 | 128 | 4096 | 13696 | 8192 | 151552 | N | Y | N | | -| glm_4_9b_chat | 40 | 32 | 2 | 128 | 4096 | 13696 | 8192 | 151552 | N | Y | N | | - -* nL: number of layers -* nAH: number of attention heads -* nKH: number of kv heads (less than nAH if multi_query_attention is used) -* nHW: head width -* nH: hidden size -* nF: FFN hidden size -* nMSL: max sequence length (input + output) -* nV: vocabulary size -* bP2D: use position_encoding_2d (Y: Yes, N: No) -* bBQKV: use bias for QKV multiplication in self-attention -* bBDense: use bias for Dense multiplication in self-attention - -## Tokenizer and special tokens comparison - -| Name | Tokenizer | bos | eos | pad | cls | startofpiece | endofpiece | mask | smask | gmask | -| :--------------: | :--------------: | :----: | :----: | :---: | :---: | :----------: | :--------: | :----: | :---: | :----: | -| glm_4_9b | ChatGLM4Tokenizer | | 151329 | 151329 | | 151333 | | | | | -| glm_4_9b_chat | ChatGLM4Tokenizer | | 151329 | 151329 | | 151333 | | | | | - -## Usage - -The next section describe how to build the engine and run the inference demo. - -### 1. Download repo and weights from HuggingFace Transformers - -```bash -pip install -r requirements.txt -apt-get update -apt-get install git-lfs - -# clone one or more models we want to build -git clone https://huggingface.co/THUDM/glm-10b glm_10b -git clone https://huggingface.co/THUDM/glm-4-9b glm_4_9b -``` - -### 2. Convert weights from HF Transformers to TensorRT LLM format - -The [`convert_checkpoint.py`](./convert_checkpoint.py) script converts HF weights to TensorRT LLM checkpoints. The number of checkpoint files (in .safetensors format) is same to the number of GPUs used to run inference. - -```bash -# GLM-4-9B: single gpu, dtype float16 -python3 convert_checkpoint.py --model_dir glm_4_9b --output_dir trt_ckpt/glm_4_9b/fp16/1-gpu -``` - -### 3. Build TensorRT engine(s) - -The `trtllm-build` command builds TensorRT LLM engines from TensorRT LLM checkpoints. The number of engine files is also same to the number of GPUs used to run inference. - -Normally, the `trtllm-build` command only requires a single GPU, but you can enable parallel building by passing the number of GPUs to the `--workers` argument. - -```bash -# GLM-4-9B: single-gpu engine with dtype float16, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir trt_ckpt/glm_4_9b/fp16/1-gpu \ - --gemm_plugin bfloat16 \ - --output_dir trt_engines/glm_4_9b/fp16/1-gpu -``` - -If the engines are run successfully, you will see output like (glm-4-9b as the example): - -```txt -...... -[01/26/2024-02:40:36] [TRT] [I] Engine generation completed in 136.52 seconds. -[01/26/2024-02:40:36] [TRT] [I] [MemUsageStats] Peak memory usage of TRT CPU/GPU memory allocators: CPU 1016 MiB, GPU 11909 MiB -[01/26/2024-02:40:36] [TRT] [I] [MemUsageChange] TensorRT-managed allocation in building engine: CPU +0, GPU +11909, now: CPU 0, GPU 11909 (MiB) -[01/26/2024-02:40:40] [TRT] [I] [MemUsageStats] Peak memory usage during Engine building and serialization: CPU: 29706 MiB -[01/26/2024-02:40:40] [TRT-LLM] [I] Total time of building Unnamed Network 0: 00:02:20 -[01/26/2024-02:40:42] [TRT-LLM] [I] Serializing engine to trt_engines/glm_4_9b/fp16/1-gpu/rank0.engine... -[01/26/2024-02:42:29] [TRT-LLM] [I] Engine serialized. Total time: 00:01:47 -[01/26/2024-02:42:30] [TRT-LLM] [I] Total time of building all engines: 00:05:19 -``` - -#### Enable plugins - -* Use `--gpt_attention_plugin ` to configure GPT Attention plugin (default as float16) -* Use `--gemm_plugin ` to configure GEMM plugin (default as float16) -* Use `--context_fmha enable` to enable FMHA kernels, which can provide better performance and low GPU memory occupation. - * `--gpt_attention_plugin float16` must be used when using FMHA. - -#### In-flight batching - -* The engine(s) must be built accordingly if [in-flight batching in C++ runtime](../../docs/in_flight_batching.md) will be used. -* Use `--gpt_attention_plugin float16`, `--paged_kv_cache enable`, `--remove_input_padding enable` to build engine(s) supporting In-flight Batching. - * It is possible to use `--gpt_attention_plugin float32` In-flight Batching. - * The size of the block in paged KV cache can be controlled additionally by using `--tokens_per_block=N`. - -### 4. Run inference - -#### Single node, single GPU - -```bash -# Run the default engine of GLM-4-9B on single GPU, other model name is available if built. -python3 ../../../run.py --input_text "What's new between ChatGLM3-6B and ChatGLM2-6B?" \ - --max_output_len 50 \ - --tokenizer_dir glm_4_9b \ - --engine_dir trt_engines/glm_4_9b/fp16/1-gpu -``` - -#### Single node, multi GPU - -```bash -# Run the Tensor Parallel 2 engine of glm_4_9b on two GPU, other model name is available if built. -mpirun -n 2 \ - python ../../../run.py --input_text "What's new between ChatGLM3-6B and ChatGLM2-6B?" \ - --max_output_len 50 \ - --tokenizer_dir glm_4_9b \ - --engine_dir trt_engines/glm_4_9b/fp16/2-gpu -``` - -* `--allow-run-as-root` might be needed if using `mpirun` as root. -* `trtllm-build` flag `--context_fmha enable` uses FP16 accumulator, which might cause low accuracy. In this case, add `--enable_context_fmha_fp32_acc` to the inference command should be used to protect accuracy at a cost of small performance drop. - -If the engines are run successfully, you will see output like (ChatGLM3-6B as the example): - -```txt -...... -Input [Text 0]: "[gMASK]sop What's new between ChatGLM3-6B and ChatGLM2-6B?" -Output [Text 0 Beam 0]: "There is no new information provided in the official documentation, but I found some differences in the code. The ChatGLM3-6B has an additional parameter called 'config', which is not present in ChatGLM2-6B. Additionally" -``` - -### 5. Run summarization task - -```bash -# Run the summarization of glm_4_9b task, other model name is available if built. -python3 ../../../summarize.py --test_trt_llm \ - --hf_model_dir glm_4_9b \ - --engine_dir trt_engines/glm_4_9b/fp16/1-gpu -``` - -### Weight Only quantization - -Use `--use_weight_only` to enable INT8-Weight-Only quantization, this will significantly lower the latency and memory footprint. Furthermore, use `--weight_only_precision int8` or `--weight_only_precision int4` to configure the data type of the weights. - -```bash -# glm_4_9b: single gpu, int8 weight only quantization -python3 convert_checkpoint.py --model_dir glm_4_9b \ - --use_weight_only \ - --weight_only_precision int8 \ - --output_dir trt_ckpt/glm_4_9b/int8_wo/1-gpu - -# glm_4_9b: single-gpu engine with int8 weight only quantization, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir trt_ckpt/glm_4_9b/int8_wo/1-gpu \ - --gemm_plugin float16 \ - --output_dir trt_engines/glm_4_9b/int8_wo/1-gpu - -# Run inference. -python3 ../../../run.py --input_text "What's new between ChatGLM3-6B and ChatGLM2-6B?" \ - --max_output_len 50 \ - --tokenizer_dir glm_4_9b \ - --engine_dir trt_engines/glm_4_9b/int8_wo/1-gpu -``` - -### Smooth Quantization (SQ) - -Use `--smoothquant` to enable smooth quantization. - -```bash -# glm_4_9b: single gpu, int8 smooth quantization -python3 convert_checkpoint.py --model_dir glm_4_9b \ - --smoothquant 0.5 \ - --per_channel \ - --per_token \ - --output_dir trt_ckpt/glm_4_9b/sq/1-gpu - -# glm_4_9b: single-gpu engine with int8 smooth quantization, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir trt_ckpt/glm_4_9b/sq/1-gpu \ - --gemm_plugin float16 \ - --output_dir trt_engines/glm_4_9b/sq/1-gpu - -# Run inference. -python3 ../../../run.py --input_text "What's new between ChatGLM3-6B and ChatGLM2-6B?" \ - --max_output_len 50 \ - --tokenizer_dir glm_4_9b \ - --engine_dir trt_engines/glm_4_9b/sq/1-gpu -``` - -### Activation-aware Weight Quantization (AWQ) - -The [`quantize.py`](../../../quantization/quantize.py) script can be used to quantize the models and export TensorRT LLM checkpoints. - -```bash -# glm_4_9b: single gpu, int4 awq quantization -python ../../../quantization/quantize.py --model_dir glm_4_9b \ - --dtype float16 \ - --qformat int4_awq \ - --output_dir trt_ckpt/glm_4_9b/int4_awq/1-gpu - -# glm_4_9b: single-gpu engine with int4 awq quantization, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir trt_ckpt/glm_4_9b/int4_awq/1-gpu \ - --gemm_plugin float16 \ - --output_dir trt_engines/glm_4_9b/int4_awq/1-gpu - -# Run inference. -python3 ../../../run.py --input_text "What's new between ChatGLM3-6B and ChatGLM2-6B?" \ - --max_output_len 50 \ - --tokenizer_dir glm_4_9b \ - --engine_dir trt_engines/glm_4_9b/int4_awq/1-gpu -``` - -### FP8 Quantization - -The [`quantize.py`](../../../quantization/quantize.py) script can be used to quantize the models and export TensorRT LLM checkpoints. - -```bash -# glm_4_9b: single gpu, fp8 quantization -python ../../../quantization/quantize.py --model_dir glm_4_9b \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir trt_ckpt/glm_4_9b/fp8/1-gpu - -# glm_4_9b: single-gpu engine with fp8 quantization, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir trt_ckpt/glm_4_9b/fp8/1-gpu \ - --gemm_plugin float16 \ - --output_dir trt_engines/glm_4_9b/fp8/1-gpu - -# Run inference. -python3 ../../../run.py --input_text "What's new between ChatGLM3-6B and ChatGLM2-6B?" \ - --max_output_len 50 \ - --tokenizer_dir glm_4_9b \ - --engine_dir trt_engines/glm_4_9b/fp8/1-gpu -``` diff --git a/examples/models/core/glm-4-9b/convert_checkpoint.py b/examples/models/core/glm-4-9b/convert_checkpoint.py deleted file mode 100644 index 4f7f67e03790..000000000000 --- a/examples/models/core/glm-4-9b/convert_checkpoint.py +++ /dev/null @@ -1,271 +0,0 @@ -import argparse -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import ChatGLMForCausalLM -from tensorrt_llm.models.chatglm.config import GLM_VERSIONS -from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument( - '--chatglm_version', - default=None, - choices=[None] + GLM_VERSIONS, - help= - "By default the script will try to infer the chatglm_version from model_dir. " - "Or users may overwrite chatglm_version by explicitly passing the version." - ) - 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('--cp_size', - type=int, - default=1, - help='N-way context parallelism size') - parser.add_argument( - '--dtype', - type=str, - default='auto', - choices=['auto', 'float16', 'bfloat16', 'float32'], - help= - "The data type for the model weights and activations if not quantized. " - "If 'auto', the data type is automatically inferred from the source model; " - "however, if the source dtype is float32, it is converted to float16.") - parser.add_argument('--logits_dtype', - type=str, - default='float32', - choices=['float16', 'float32']) - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - - parser.add_argument( - '--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( - '--calib_dataset', - type=str, - default='cnn_dailymail', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - 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( - '--per_channel', - action="store_true", - default=False, - 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', - action="store_true", - default=False, - 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( - '--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('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument( - '--device', - help= - "The device to run calibration; effective for HuggingFace model only.", - default='cuda', - choices=['cuda', 'cpu']) - parser.add_argument("--load_model_on_cpu", action="store_true") - args = parser.parse_args() - - return args - - -def args_to_quant_config(args: argparse.Namespace) -> QuantConfig: - '''return config dict with quantization info based on the command line args - ''' - quant_config = QuantConfig() - - if args.use_weight_only: - if args.weight_only_precision == 'int8': - quant_config.quant_algo = QuantAlgo.W8A16 - elif args.weight_only_precision == 'int4': - quant_config.quant_algo = QuantAlgo.W4A16 - elif args.smoothquant: - quant_config.smoothquant_val = args.smoothquant - if args.per_channel: - if args.per_token: - quant_config.quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN - else: - quant_config.quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN - else: - if args.per_token: - quant_config.quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN - else: - quant_config.quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN - - if args.int8_kv_cache: - quant_config.kv_cache_quant_algo = QuantAlgo.INT8 - - return quant_config - - -def args_to_build_options(args): - return { - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'logits_dtype': args.logits_dtype, - } - - -def execute(workers, func): - if workers == 1: - for rank, f in enumerate(func): - f(rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def convert_and_save_hf(args): - world_size = args.tp_size * args.pp_size - quant_config = args_to_quant_config(args) - override_fields = args_to_build_options(args) - - if args.smoothquant is not None or args.int8_kv_cache: - mapping = Mapping(world_size=world_size, - tp_size=args.tp_size, - pp_size=args.pp_size, - cp_size=args.cp_size) - ChatGLMForCausalLM.quantize(args.model_dir, - args.output_dir, - dtype=args.dtype, - mapping=mapping, - quant_config=quant_config, - device=args.device, - calib_dataset=args.calib_dataset, - **override_fields) - else: - - def convert_and_save_rank(rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - glm = ChatGLMForCausalLM.from_hugging_face( - args.model_dir, - args.dtype, - mapping=mapping, - quant_config=quant_config, - chatglm_version=args.chatglm_version, - load_model_on_cpu=args.load_model_on_cpu, - **override_fields) - glm.config.mapping.cp_size = args.cp_size - glm.config.mapping.attn_tp_size = -1 - glm.config.mapping.attn_cp_size = -1 - glm.config.mapping.world_size *= args.cp_size - glm.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del glm - - execute(args.workers, [convert_and_save_rank] * world_size) - release_gc() - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - logger.set_level(args.log_level) - - assert args.pp_size == 1, "Pipeline parallelism is not supported." - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - convert_and_save_hf(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/core/glm-4-9b/requirements.txt b/examples/models/core/glm-4-9b/requirements.txt deleted file mode 100644 index cdc65bf2bb38..000000000000 --- a/examples/models/core/glm-4-9b/requirements.txt +++ /dev/null @@ -1,8 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -protobuf -rouge_score -sentencepiece -tiktoken diff --git a/examples/models/core/glm-4-9b/tokenization_chatglm.py b/examples/models/core/glm-4-9b/tokenization_chatglm.py deleted file mode 100644 index 75495af5f7f7..000000000000 --- a/examples/models/core/glm-4-9b/tokenization_chatglm.py +++ /dev/null @@ -1,343 +0,0 @@ -import base64 -import json -import os -from typing import Any, Dict, List, Optional, Union - -import regex as re -import tiktoken -from torch import TensorType -from transformers import PreTrainedTokenizer -from transformers.tokenization_utils_base import BatchEncoding, EncodedInput -from transformers.utils import PaddingStrategy - - -class ChatGLM4Tokenizer(PreTrainedTokenizer): - vocab_files_names = {"vocab_file": "tokenizer.model"} - model_input_names = ["input_ids", "attention_mask", "position_ids"] - - def __init__(self, - vocab_file, - padding_side="left", - clean_up_tokenization_spaces=False, - encode_special_tokens=False, - **kwargs): - self.name = "GLM4Tokenizer" - self.vocab_file = vocab_file - pat_str = "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+" - self.pat_str = re.compile(pat_str) - self.encode_special_tokens = encode_special_tokens - - mergeable_ranks = {} - with open(vocab_file) as f: - for line in f: - token, rank = line.strip().split() - rank = int(rank) - token = base64.b64decode(token) - mergeable_ranks[token] = rank - - self.mergeable_ranks = mergeable_ranks - - self.tokenizer = tiktoken.Encoding(name="my_tokenizer", - pat_str=pat_str, - mergeable_ranks=mergeable_ranks, - special_tokens={}) - self.decoder = {rank: token for token, rank in mergeable_ranks.items()} - self.n_words = len(self.decoder) - - super().__init__( - padding_side=padding_side, - clean_up_tokenization_spaces=clean_up_tokenization_spaces, - **kwargs) - - @property - def vocab_size(self): - return self.n_words - - def get_vocab(self): - """ Returns vocab as a dict """ - vocab = { - self._convert_id_to_token(i): i - for i in range(self.vocab_size) - } - vocab.update(self.added_tokens_encoder) - return vocab - - def convert_tokens_to_string(self, tokens: List[Union[bytes, str, - int]]) -> str: - """ - Converts a sequence of tokens in a single string. - """ - text = "" - temp = b"" - for t in tokens: - if isinstance(t, int): - t = chr(t) - if isinstance(t, str): - if temp: - text += temp.decode("utf-8", errors="replace") - elif isinstance(t, bytes): - temp += t - else: - raise TypeError( - "token should only be of type int, bytes or str") - if temp: - text += temp.decode("utf-8", errors="replace") - return text - - def _tokenize(self, text, **kwargs): - tokens = [] - ids = self.tokenizer.encode(text) - for t in ids: - tokens.append(self.decoder[t]) - return tokens - - def _convert_token_to_id(self, token): - """ Converts a token (str) in an id using the vocab. """ - return self.mergeable_ranks[token] - - def _convert_id_to_token(self, index): - """Converts an index (integer) in a token (str) using the vocab.""" - return self.decoder.get(index, "") - - def save_vocabulary(self, save_directory, filename_prefix=None): - """ - Save the vocabulary and special tokens file to a directory. - - Args: - save_directory (`str`): - The directory in which to save the vocabulary. - filename_prefix (`str`, *optional*): - An optional prefix to add to the named of the saved files. - - Returns: - `Tuple(str)`: Paths to the files saved. - """ - if os.path.isdir(save_directory): - vocab_file = os.path.join(save_directory, - self.vocab_files_names["vocab_file"]) - else: - vocab_file = save_directory - - with open(self.vocab_file, 'rb') as fin: - proto_str = fin.read() - - with open(vocab_file, "wb") as writer: - writer.write(proto_str) - - return (vocab_file, ) - - def get_prefix_tokens(self): - prefix_tokens = [ - self.convert_tokens_to_ids("[gMASK]"), - self.convert_tokens_to_ids("") - ] - return prefix_tokens - - def build_single_message(self, role, metadata, message, tokenize=True): - assert role in ["system", "user", "assistant", "observation"], role - if tokenize: - role_tokens = [self.convert_tokens_to_ids(f"<|{role}|>") - ] + self.tokenizer.encode(f"{metadata}\n", - disallowed_special=()) - message_tokens = self.tokenizer.encode(message, - disallowed_special=()) - tokens = role_tokens + message_tokens - return tokens - else: - return str(f"<|{role}|>{metadata}\n{message}") - - def apply_chat_template( - self, - conversation: Union[List[Dict[str, str]], List[List[Dict[str, str]]], - "Conversation"], - add_generation_prompt: bool = False, - tokenize: bool = True, - padding: bool = False, - truncation: bool = False, - max_length: Optional[int] = None, - return_tensors: Optional[Union[str, TensorType]] = None, - return_dict: bool = False, - tokenizer_kwargs: Optional[Dict[str, Any]] = None, - add_special_tokens: bool = True, - **kwargs, - ) -> Union[str, List[int], List[str], List[List[int]], BatchEncoding]: - - if return_dict and not tokenize: - raise ValueError( - "`return_dict=True` is incompatible with `tokenize=False`, because there is no dict " - "of tokenizer outputs to return.") - - def handle_single_conversation(conversation): - input_ids = self.get_prefix_tokens() if add_special_tokens else [] - input_message = "[gMASK]" if add_special_tokens else "" - for item in conversation: - if item.get("tools"): - tools = item["tools"] - content = "你是一个名为 GhatGLM 的人工智能助手。你是基于智谱AI训练的语言模型 GLM-4 模型开发的，你的任务是针对用户的问题和要求提供适当的答复和支持。" - content += "\n\n# 可用工具" - for tool in tools: - if tool["type"] == "function": - function = tool["function"] - content += f"\n\n## {function['name']}\n\n{json.dumps(function, ensure_ascii=False, indent=4)}" - content += "\n在调用上述函数时，请使用 Json 格式表示调用的参数。" - elif tool["type"] == "python": - content += "\n\n## python\n\n当你向 `python` 发送包含 Python 代码的消息时，该代码将会在一个有状态的 Jupyter notebook 环境中执行。\n`python` 返回代码执行的输出，或在执行 60 秒后返回超时。\n`/mnt/data` 将会持久化存储你的文件。在此会话中，`python` 无法访问互联网。不要使用 `python` 进行任何网络请求或者在线 API 调用，这些在线内容的访问将不会成功。" - elif tool["type"] == "simple_browser": - content += "\n\n## simple_browser\n\n你可以使用 `simple_browser` 工具。该工具支持以下函数：\n`search(query: str, recency_days: int)`：使用搜索引擎进行查询并显示结果，可以使用 `recency_days` 参数控制搜索内容的时效性。\n`mclick(ids: list[int])`：获取一系列指定 id 的页面内容。每次调用时，须选择3-10个页面。选择多个角度的页面，同时尽可能选择可信任的信息来源。考虑到部分页面是无法加载的，你也可以多打开一些可能有用的页面而不用担心内容过多。\n`open_url(url: str)`：打开指定的 URL。\n\n使用 `【{引用 id}†{引用文本}】` 来引用内容。\n\n操作步骤：1. 使用 `search` 来获得信息列表; 2. 使用 `mclick` 来获取指定 ID 页面的内容; 3. 根据获得的内容进行回复。在回复中应当引用信息来源。\n 如果用户提供了 URL，也可以用 `open_url` 直接打开页面。\n如果初次搜索结果没有找到合适的信息，也可以再次使用 `search` 进行搜索。" - elif tool["type"] == "cogview": - content += "\n\n## cogview\n\n如果用户的请求中包含了对图像的描述，你可以使用 `cogview` 来生成图像并展示给用户。你需要向 `cogview` 发送图像描述，规则：\n- 发送给 `cogview` 的消息必须使用英语。用户的中文描述必须完全翻译为英语。\n- 应当尽可能详细地描述图像生成的需求，需求描述约 100 英文单词。\n- 保持用户原始描述的意图。不要虚构内容或者没见过的人物。\n- 如无特殊说明，所在地为中国，持有中国立场并遵循中国社会主义价值观。" - else: - raise NotImplementedError( - f"Unknown tool type {tool['type']}") - input = self.build_single_message("system", - "", - content, - tokenize=tokenize) - if tokenize: - input_ids.extend(input) - else: - input_message += input - if item["content"]: - input = self.build_single_message(item["role"], - item.get("metadata", ""), - item["content"], - tokenize=tokenize) - if tokenize: - input_ids.extend(input) - else: - input_message += input - if add_generation_prompt: - if tokenize: - input_ids.extend( - [self.convert_tokens_to_ids("<|assistant|>")]) - else: - input_message += "<|assistant|>" - return input_ids if tokenize else input_message - - # Main logic to handle different conversation formats - if isinstance(conversation, list) and all( - isinstance(i, dict) for i in conversation): - result = handle_single_conversation(conversation) - elif isinstance(conversation, list) and all( - isinstance(i, list) for i in conversation): - result = [handle_single_conversation(c) for c in conversation] - elif hasattr(conversation, "messages"): - result = handle_single_conversation(conversation.messages) - else: - raise ValueError("Invalid conversation format") - - if tokenize: - output = self.batch_encode_plus( - [result] if isinstance(result[0], int) else result, - padding=padding, - truncation=truncation, - max_length=max_length, - return_tensors=return_tensors, - is_split_into_words=True, - add_special_tokens=False) - if return_dict: - return output - else: - return output["input_ids"] - else: - return result - - def build_inputs_with_special_tokens( - self, - token_ids_0: List[int], - token_ids_1: Optional[List[int]] = None) -> List[int]: - """ - Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and - adding special tokens. A BERT sequence has the following format: - - - single sequence: `[CLS] X [SEP]` - - pair of sequences: `[CLS] A [SEP] B [SEP]` - - Args: - token_ids_0 (`List[int]`): - List of IDs to which the special tokens will be added. - token_ids_1 (`List[int]`, *optional*): - Optional second list of IDs for sequence pairs. - - Returns: - `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens. - """ - prefix_tokens = self.get_prefix_tokens() - token_ids_0 = prefix_tokens + token_ids_0 - if token_ids_1 is not None: - token_ids_0 = token_ids_0 + token_ids_1 + [ - self.convert_tokens_to_ids("") - ] - return token_ids_0 - - def _pad( - self, - encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding], - max_length: Optional[int] = None, - padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, - pad_to_multiple_of: Optional[int] = None, - return_attention_mask: Optional[bool] = None, - padding_side: str = "left", # Fix for new transformers - ) -> dict: - """ - Pad encoded inputs (on left/right and up to predefined length or max length in the batch) - - Args: - encoded_inputs: - Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`). - max_length: maximum length of the returned list and optionally padding length (see below). - Will truncate by taking into account the special tokens. - padding_strategy: PaddingStrategy to use for padding. - - - PaddingStrategy.LONGEST Pad to the longest sequence in the batch - - PaddingStrategy.MAX_LENGTH: Pad to the max length (default) - - PaddingStrategy.DO_NOT_PAD: Do not pad - The tokenizer padding sides are defined in self.padding_side: - - - 'left': pads on the left of the sequences - - 'right': pads on the right of the sequences - pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value. - This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability - `>= 7.5` (Volta). - return_attention_mask: - (optional) Set to False to avoid returning attention mask (default: set to model specifics) - """ - # Load from model defaults - assert self.padding_side == "left" - - required_input = encoded_inputs[self.model_input_names[0]] - seq_length = len(required_input) - - if padding_strategy == PaddingStrategy.LONGEST: - max_length = len(required_input) - - if max_length is not None and pad_to_multiple_of is not None and ( - max_length % pad_to_multiple_of != 0): - max_length = ( - (max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of - - needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len( - required_input) != max_length - - # Initialize attention mask if not present. - if "attention_mask" not in encoded_inputs: - encoded_inputs["attention_mask"] = [1] * seq_length - - if "position_ids" not in encoded_inputs: - encoded_inputs["position_ids"] = list(range(seq_length)) - - if needs_to_be_padded: - difference = max_length - len(required_input) - - if "attention_mask" in encoded_inputs: - encoded_inputs["attention_mask"] = [ - 0 - ] * difference + encoded_inputs["attention_mask"] - if "position_ids" in encoded_inputs: - encoded_inputs["position_ids"] = [ - 0 - ] * difference + encoded_inputs["position_ids"] - encoded_inputs[self.model_input_names[ - 0]] = [self.pad_token_id] * difference + required_input - - return encoded_inputs diff --git a/examples/models/core/gpt/.gitignore b/examples/models/core/gpt/.gitignore deleted file mode 100644 index 4c73534033eb..000000000000 --- a/examples/models/core/gpt/.gitignore +++ /dev/null @@ -1,5 +0,0 @@ -gpt* -*.log -c-model -*.safetensors -*.json diff --git a/examples/models/core/gpt/README.md b/examples/models/core/gpt/README.md deleted file mode 100644 index f50f45e75d67..000000000000 --- a/examples/models/core/gpt/README.md +++ /dev/null @@ -1,786 +0,0 @@ -# GPT - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document explains how to build the [GPT](https://huggingface.co/gpt2) model using TensorRT LLM and run on a single GPU, a single node with multiple GPUs or multiple nodes with multiple GPUs. - -- [GPT](#gpt) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [1. Download weights from HuggingFace Transformers](#1-download-weights-from-huggingface-transformers) - - [2. Convert weights from HF Transformers to TensorRT LLM format](#2-convert-weights-from-hf-transformers-to-tensorrt-llm-format) - - [3. Build TensorRT engine(s)](#3-build-tensorrt-engines) - - [Fused MultiHead Attention (FMHA)](#fused-multihead-attention-fmha) - - [In-flight batching and paged KV cache](#in-flight-batching-and-paged-kv-cache) - - [4. Build TensorRT engine(s) with Random Weights](#4-build-tensorrt-engines-with-random-weights) - - [5. Run inference](#5-run-inference) - - [Single node, single GPU](#single-node-single-gpu) - - [Single node, multiple GPUs](#single-node-multiple-gpus) - - [Multiple nodes, multiple GPUs using Slurm](#multiple-nodes-multiple-gpus-using-slurm) - - [Quantization](#quantization) - - [SmoothQuant](#smoothquant) - - [Model Transformation](#model-transformation) - - [INT8 inference](#int8-inference) - - [Usage](#usage-1) - - [INT8 KV Cache](#int8-kv-cache) - - [Weight Only Quantization](#weight-only-quantization) - - [FP8 Quantization](#fp8-quantization) - - [Embedding Parallelism](#embedding-parallelism) - - [1. Embedding parallelism](#1-embedding-parallelism) - - [2. The sharding dimension for embedding parallelism](#2-the-sharding-dimension-for-embedding-parallelism) - - [GPT Variant - Granite(20B and 34B)](#gpt-variant---granite) - - [GPT Variant - SantaCoder](#gpt-variant---santacoder) - - [GPT Variant - StarCoder (v1 and v2)](#gpt-variant---starcoder-v1-and-v2) - - [Run StarCoder2 with LoRA](#run-starcoder2-with-lora) - - [GPT-Next](#gpt-next) - - [Prompt-tuning](#prompt-tuning) - - [MultiLoRA with the Nemo checkpoint](#multilora-with-the-nemo-checkpoint) - -## 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/models/core/gpt`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`run.py`](../../../run.py) to run the inference on an input text; -* [`summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 - * FP8 - * Inflight Batching - * PAGED_KV_CACHE - * FP8 KV CACHE - * Tensor Parallel - * Pipeline Parallel - * STRONGLY TYPED - * INT8 SmoothQuant - * INT8 weight only - * INT4 weight only - -## Usage - -The next two sections describe how to convert the weights from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) -format to the TensorRT LLM format. - -### 1. Download weights from HuggingFace Transformers - -Please install required packages first: - -```bash -pip install -r requirements.txt -``` - -```bash -# Download hf gpt2 model -rm -rf gpt2 && git clone https://huggingface.co/gpt2-medium gpt2 -pushd gpt2 && rm pytorch_model.bin model.safetensors && wget -q https://huggingface.co/gpt2-medium/resolve/main/pytorch_model.bin && popd -``` - -### 2. Convert weights from HF Transformers to TensorRT LLM format -The [`convert_checkpoint.py`](./convert_checkpoint.py) script converts HF weights to TensorRT LLM checkpoints. The number of checkpoint files (in .safetensors format) is same to the number of GPUs used to run inference. - -```bash -# single gpu, dtype float16 -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --output_dir gpt2/trt_ckpt/fp16/1-gpu - -# 2-way tensor parallelism -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --tp_size 2 \ - --output_dir gpt2/trt_ckpt/fp16/2-gpu - -# 2-way tensor parallelism and 2-way pipeline parallelism -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --tp_size 2 \ - --pp_size 2 \ - --output_dir gpt2/trt_ckpt/fp16/4-gpu -``` - -### 3. Build TensorRT engine(s) -The `trtllm-build` command builds TensorRT LLM engines from TensorRT LLM checkpoints. The checkpoint directory provides the model's weights and architecture configuration. The number of engine files is also same to the number of GPUs used to run inference. - -`trtllm-build` command has a variety of options. In particular, the plugin-related options have two categories: -* Plugin options that requires a data type (e.g., `gpt_attention_plugin`), you can - * explicitly specify `float16`/`bfloat16`/`float32`, so that the plugins are enabled with the specified precision; - * implicitly specify `auto`, so that the plugins are enabled with the precision automatically inferred from model dtype (i.e., the dtype specified in weight conversion); or - * disable the plugin by `disable`. -* Other features that requires a boolean (e.g., `context_fmha`, `paged_kv_cache`, `remove_input_padding`), you can - * enable/disable the feature by specifying `enable`/`disable`. - -The defaults have been carefully tuned for better performance. For example, `gpt_attention_plugin`, `context_fmha`, `paged_kv_cache` and `remove_input_padding` are enabled by default. See more details by `trtllm-build --help`. - -Normally, the `trtllm-build` command only requires a single GPU, but you can enable parallel building by passing the number of GPUs to the `--workers` argument. - -```bash -# Build a single-GPU float16 engine from TensorRT LLM checkpoint. -# gpt_attention_plugin (the special TensorRT LLM GPT Attention plugin) and remove_input_padding are enabled by default for runtime performance. -trtllm-build --checkpoint_dir gpt2/trt_ckpt/fp16/1-gpu \ - --output_dir gpt2/trt_engines/fp16/1-gpu - -# Build 2-way tensor parallelism engines from TensorRT LLM checkpoint. -trtllm-build --checkpoint_dir gpt2/trt_ckpt/fp16/2-gpu \ - --output_dir gpt2/trt_engines/fp16/2-gpu - -# Build 2-way tensor parallelism and 2-way pipeline parallelism engines from TensorRT LLM checkpoint. -trtllm-build --checkpoint_dir gpt2/trt_ckpt/fp16/4-gpu \ - --output_dir gpt2/trt_engines/fp16/4-gpu -``` - -If the engines are built successfully, you will see output like: -``` -...... -[03/12/2024-10:21:08] [TRT] [I] Engine generation completed in 35.9738 seconds. -[03/12/2024-10:21:08] [TRT] [I] [MemUsageStats] Peak memory usage of TRT CPU/GPU memory allocators: CPU 212 MiB, GPU 775 MiB -[03/12/2024-10:21:08] [TRT] [I] [MemUsageChange] TensorRT-managed allocation in building engine: CPU +0, GPU +775, now: CPU 0, GPU 775 (MiB) -[03/12/2024-10:21:09] [TRT] [I] [MemUsageStats] Peak memory usage during Engine building and serialization: CPU: 6600 MiB -[03/12/2024-10:21:09] [TRT-LLM] [I] Total time of building Unnamed Network 0: 00:00:36 -[03/12/2024-10:21:09] [TRT-LLM] [I] Serializing engine to gpt2/trt_engines/fp16/1-gpu/rank0.engine... -[03/12/2024-10:21:11] [TRT-LLM] [I] Engine serialized. Total time: 00:00:02 -[03/12/2024-10:21:11] [TRT-LLM] [I] Total time of building all engines: 00:00:41 -``` - -#### Fused MultiHead Attention (FMHA) - -`trtllm-build` enables FMHA kernels by default. You may disable it by adding `--context_fmha disable`. - -If you find that the default fp16 accumulation cannot meet the requirement, you can try to enable fp32 accumulation by adding `--enable_context_fmha_fp32_acc` to the inference command (`run.py` or `summarize.py`). However, it is expected to see performance drop. - -Note that the FMHA kernels have to be used together with `gpt_attention_plugin` enabled. - -#### In-flight batching and paged KV cache - -If one wants to use [in-flight batching in C++ runtime](../../docs/in_flight_batching.md), the engine(s) must be built accordingly. In-flight batching in C++ runtime works only with attention plugin, paged KV cache and with packed data. Currently, the `trtllm-build` by default enables `gpt_attention_plugin`, `paged_kv_cache` and `remove_input_padding`, so the built engine(s) can support in-flight batching (unless you explicitly disable one of these options). One can additionally control the size of the block in paged KV cache using `--tokens_per_block=N`. - -### 4. Build TensorRT engine(s) with Random Weights -You can build engine(s) using random weights, which is useful for benchmarking. First, the [`../generate_checkpoint_config.py`](../generate_checkpoint_config.py) script can be used to generate a TensorRT LLM checkpoint config file: - -```bash -# Generate an 8-GPU GPT-175B float16 checkpoint config file. -python3 ../../../generate_checkpoint_config.py --architecture GPTForCausalLM \ - --vocab_size 51200 \ - --hidden_size 12288 \ - --num_hidden_layers 96 \ - --num_attention_heads 96 \ - --dtype float16 \ - --tp_size 8 \ - --output_path gpt_175b/trt_ckpt/fp16/8-gpu/config.json - - -# Generate a 16-GPU GPT-530B float16 checkpoint config file. -python3 ../../../generate_checkpoint_config.py --architecture GPTForCausalLM \ - --vocab_size 51200 \ - --hidden_size 20480 \ - --num_hidden_layers 105 \ - --num_attention_heads 128 \ - --dtype float16 \ - --tp_size 16 \ - --output_path gpt_530b/trt_ckpt/fp16/16-gpu/config.json -``` - -Then, use `trtllm-build` command to build engine(s) with random weights and the model architecture specified by the generated config file. - -```bash -# Build 8-GPU GPT-175B float16 engines using dummy weights, useful for performance tests. -# Enable several TensorRT LLM plugins to increase runtime performance. It also helps with build time. -trtllm-build --model_config gpt_175b/trt_ckpt/fp16/8-gpu/config.json \ - --gemm_plugin auto \ - --max_batch_size 256 \ - --output_dir gpt_175b/trt_engines/fp16/8-gpu \ - --workers 8 - -# Build 16-GPU GPT-530B float16 engines using dummy weights, useful for performance tests. -# Enable several TensorRT LLM plugins to increase runtime performance. It also helps with build time. -trtllm-build --model_config gpt_530b/trt_ckpt/fp16/16-gpu/config.json \ - --gemm_plugin auto \ - --max_batch_size 128 \ - --max_input_len 128 \ - --max_seq_len 148 \ - --output_dir gpt_530b/trt_engines/fp16/16-gpu \ - --workers 8 -``` - -### 5. Run inference -#### Single node, single GPU - -The [`run.py`](../../../run.py) script can be used to run inference with the built engine(s). - -```bash -python3 ../../../run.py --engine_dir gpt2/trt_engines/fp16/1-gpu \ - --tokenizer_dir gpt2 \ - --max_output_len 8 -``` - -If the engines are run successfully, you will see output like: -``` -...... -Input [Text 0]: "Born in north-east France, Soyer trained as a" -Output [Text 0 Beam 0]: " chef before moving to London in the early" -``` - -The [`summarize.py`](../../../summarize.py) script can run the built engines to summarize the articles from the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. -For each summary, the script can compute the -[ROUGE](https://en.wikipedia.org/wiki/ROUGE_(metric)) scores and use the `ROUGE-1` score to validate the implementation. -By passing `--test_trt_llm` flag, the script will evaluate TensorRT LLM engines. You may also pass `--test_hf` flag to evaluate the HF model. - -```bash -python3 ../../../summarize.py --engine_dir gpt2/trt_engines/fp16/1-gpu \ - --hf_model_dir gpt2 \ - --test_trt_llm \ - --test_hf -``` -If the engines are run successfully, you will see output like: -``` -...... -[03/13/2024-05:43:18] [TRT-LLM] [I] TensorRT LLM (total latency: 1.520904541015625 sec) -[03/13/2024-05:43:18] [TRT-LLM] [I] TensorRT LLM (total output tokens: 0) -[03/13/2024-05:43:18] [TRT-LLM] [I] TensorRT LLM (tokens per second: 0.0) -[03/13/2024-05:43:18] [TRT-LLM] [I] TensorRT LLM beam 0 result -[03/13/2024-05:43:18] [TRT-LLM] [I] rouge1 : 21.13474087351942 -[03/13/2024-05:43:18] [TRT-LLM] [I] rouge2 : 6.2641616526063775 -[03/13/2024-05:43:18] [TRT-LLM] [I] rougeL : 16.693574311238077 -[03/13/2024-05:43:18] [TRT-LLM] [I] rougeLsum : 18.477384201634088 -[03/13/2024-05:43:18] [TRT-LLM] [I] Hugging Face (total latency: 8.76440143585205 sec) -[03/13/2024-05:43:18] [TRT-LLM] [I] HF beam 0 result -[03/13/2024-05:43:18] [TRT-LLM] [I] rouge1 : 20.834898522466 -[03/13/2024-05:43:18] [TRT-LLM] [I] rouge2 : 5.6914719275508805 -[03/13/2024-05:43:18] [TRT-LLM] [I] rougeL : 16.297064309934132 -[03/13/2024-05:43:18] [TRT-LLM] [I] rougeLsum : 18.018627021792142 -``` - -#### Single node, multiple GPUs - -To run engines using multiple GPUs on a single node, you can use `mpirun` as: - -```bash -mpirun -np 2 \ - python3 ../../../run.py --engine_dir gpt2/trt_engines/fp16/2-gpu \ - --tokenizer_dir gpt2 \ - --max_output_len 8 - -# Note that GPT-175B is built with random weights, so the output will also be random -mpirun -np 8 \ - python3 ../../../run.py --engine_dir gpt_175b/trt_engines/fp16/8-gpu \ - --max_output_len 8 -``` - -#### Multiple nodes, multiple GPUs using [Slurm](https://slurm.schedmd.com) - -To run engines using multiple nodes, you should use a cluster manager like `Slurm`. The following section shows how to configure TensorRT LLM to execute on two nodes using Slurm. - -We start by preparing an `sbatch` script called `tensorrt_llm_run.sub`. That script contains the following code (you must replace the `` strings with your own values): - -```bash -#!/bin/bash -#SBATCH -o logs/tensorrt_llm.out -#SBATCH -e logs/tensorrt_llm.error -#SBATCH -J -#SBATCH -A -#SBATCH -p -#SBATCH --nodes=2 -#SBATCH --ntasks-per-node=8 -#SBATCH --time=00:30:00 - -sudo nvidia-smi -lgc 1410,1410 - -srun --mpi=pmix \ - --container-image \ - --container-mounts : \ - --container-workdir \ - --output logs/tensorrt_llm_%t.out \ - --error logs/tensorrt_llm_%t.error \ - python3 -u ../../../run.py --engine_dir --max_output_len 8 -``` - -Then, submit the job using: - -```bash -sbatch tensorrt_llm_run.sub -``` - -You might have to contact your cluster's administrator to help you customize the above script. - - -## Quantization - -### SmoothQuant - -This section explains how to use SmoothQuant on GPT models with TensorRT-LLM. - -[SmoothQuant](https://arxiv.org/abs/2211.10438) is a post-training quantization -(PTQ) method to quantize LLM models to INT8 for faster inference. As explained -in the article, SmoothQuant modifies a model to enable INT8 quantization -without significantly altering the accuracy. - -#### Model Transformation - -A LLM model is made of multiple matrix-multiplication operations (or GEMMs): `Y -= XW` where `X` of shape `[n, k]`, holds the activation (produced at run-time) -and `W`, of shape `[k, m]` are the learned weights. `Y`, of shape `[n, m]`, is -the matrix product of `X` and `W`. - -SmoothQuant introduces scaling along the `k` dimension by defining a vector of -strictly positive coefficients `s`. `Y = X diag(s)^{-1} diag(s) W`. We now have -`Y = X'W'` where `X' = X diag(s)^{-1}` and `W' = diag(s) W`. This -transformation is introduced so the quantization behaves better. In *normal* -models, `X` tends to be ill-conditioned: it has mostly small-magnitude -coefficients, but also some outliers that makes quantization difficult. -Conversely, the re-scaled `X'` is better suited for INT8 conversion. - -In this example, we only replace Attention's QKV and MLP's FC1 GEMMs to their -Smoothquant'd version since it is sufficient to maintain the accuracy for the -GPT model. During inference, `X'` is computed by fusing the channel-wise -multiplication by `diag(s)^{-1}` with the preceding layernorm's lambda and beta -parameters. `W'` is pre-computed and doesn't need additional modification -during inference. - -#### INT8 inference - -The INT8 quantization scheme used in TensorRT LLM theoretically works on any -GPT model. However, Smoothquant'd models tend to produce more accurate results -with reduced precision. - -INT8 inference modifies GEMMs `Y = XW` so that both `X` and `W` use INT8. The -matrix-multiplication is sped-up because of smaller weight size and fast matrix -products computation thanks to NVIDIA *Tensor Cores* operating on INT8 inputs. - -During inference, X is transformed from its standard floating point (fp) -values: `X_{i8} <- X_{fp} * s_x`. This scaling puts `X` values in the INT8 -range: `[-128, 127]`. Similarly, W is scaled, `W_{i8} <- W_{fp} * s_w` but that -operation is done at model export time, no need for subsequent operations at -run-time. - -The optimized TensorRT LLM GEMM implementation for SmoothQuant does the integer -matrix-multiplication `Y_{i32} <- X_{i8} W_{i8}` and rescales the result to its -original range `Y_{fp} <- Y_{i32} * (s_x)^{-1} * (s_w)^{-1}`. Note that -`Y_{i32}` isn't stored in memory, the re-scaling happens in the GEMM's epilogue -and only `Y_{fp}` gets saved. - -By default `s_x` and `s_w` are single-value coefficients. This is the -*per-tensor* mode. Values for `s_x` and `s_w` are static, estimated at model -export time. - -TensorRT LLM also supports more elaborate modes: - - per-channel: `s_w` is a fixed vector of size `[1, m]`. For that, - TensorRT LLM loads the adequately scaled version of of `W_{i8}` at model - construction time. - - per-token: `s_x` is a vector of size `[n, 1]` determined at run-time, based - on the per-token (a.k.a. per-row) absolute maximum of `X`. -Users can mix-and-match per-channel and per-token options. Both tend to -increase the accuracy of the model at the cost of a slightly increased latency. - -#### Usage -[`convert_checkpoint.py`](./convert_checkpoint.py) features a -`--smoothquant` option. It must be set to a decimal value in `[0, 1]` and -corresponds to the `alpha` parameter in the [SmoothQuant -paper](https://arxiv.org/abs/2211.10438). Setting `--smoothquant` will smooth the model -as explained in [model transformation](#model-transformation) and export the -scaling factors needed for INT8 inference. - -By default, it will run the model in the per-tensor mode, as explained in [INT8 -inference](#int8-inference). You can add any combination of `--per_token` and `--per_channel` to get the corresponding behaviors. - -```bash -# Per-tensor SmoothQuant -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --smoothquant 0.5 \ - --output_dir gpt2/trt_ckpt/int8-sq/1-gpu - -# Per-token per-channel SmoothQuant -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --smoothquant 0.5 \ - --per_token \ - --per_channel \ - --output_dir gpt2/trt_ckpt/int8-sq-ptpc/1-gpu -``` - -Then, use `trtllm-build` to build engine(s). - -```bash -# Per-tensor SmoothQuant -trtllm-build --checkpoint_dir gpt2/trt_ckpt/int8-sq/1-gpu \ - --output_dir gpt2/trt_engines/int8-sq/1-gpu - -# Per-token per-channel SmoothQuant -trtllm-build --checkpoint_dir gpt2/trt_ckpt/int8-sq-ptpc/1-gpu \ - --output_dir gpt2/trt_engines/int8-sq-ptpc/1-gpu -``` - -Note that GPT attention plugin is required to be enabled for SmoothQuant for now. - -User can also use `ModelOpt` to do INT8 quantization. Especially for gpt variant Starcoder2. -```bash -python3 example/quantization/quantize.py --model_dir starcoder2 \ - --dtype float16 \ - --qformat int8_sq \ - --output_dir starcoder2/trt_ckpt/int8-sq/ -``` -Then, use `trtllm-build` to build engine(s). - -```bash -trtllm-build --checkpoint_dir starcoder2/trt_ckpt/int8-sq/ \ - --output_dir starcoder2/trt_engine/int8-sq/ -``` - - -### INT8 KV Cache - -[`convert_checkpoint.py`](./convert_checkpoint.py) features a -`--int8_kv_cache` option. Setting `--int8_kv_cache` will calibrate the model and export the -scaling factors needed for INT8 KV cache inference. - -```bash -# Int8 KV cache -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --int8_kv_cache \ - --output_dir gpt2/trt_ckpt/int8kv/1-gpu - -trtllm-build --checkpoint_dir gpt2/trt_ckpt/int8kv/1-gpu \ - --output_dir gpt2/trt_engines/int8kv/1-gpu -``` - -INT8 KV cache can be used with or without gpt attention plugin. - -### Weight Only Quantization - -[`convert_checkpoint.py`](./convert_checkpoint.py) features a `--use_weight_only` option that can enable weight-only quantization. You can further set the weight-only precision by passing `int8` or `int4` to the `--weight_only_precision` flag. - -```bash -# Int8 weight-only quantization -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int8 \ - --output_dir gpt2/trt_ckpt/int8-wo/1-gpu - -# Int4 weight-only quantization -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int4 \ - --output_dir gpt2/trt_ckpt/int4-wo/1-gpu -``` - -Then, use `trtllm-build` to build engine(s). - -```bash -# Int8 weight-only quantization -trtllm-build --checkpoint_dir gpt2/trt_ckpt/int8-wo/1-gpu \ - --output_dir gpt2/trt_engines/int8-wo/1-gpu - -# Int4 weight-only quantization -trtllm-build --checkpoint_dir gpt2/trt_ckpt/int4-wo/1-gpu \ - --output_dir gpt2/trt_engines/int4-wo/1-gpu -``` - -### FP8 Quantization - -[`quantize.py`](../../../quantization/quantize.py) can do FP8 quantization and/or FP8 kv cache quantization, and export TensorRT LLM checkpoint. - -```bash -# FP8 quantization with FP8 kv cache -python3 ../../../quantization/quantize.py --model_dir gpt2 \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir gpt2/trt_ckpt/fp8/1-gpu - -trtllm-build --checkpoint_dir gpt2/trt_ckpt/fp8/1-gpu \ - --output_dir gpt2/trt_engines/fp8/1-gpu -``` - -## Embedding Parallelism -Since the embedding lookup table can be several gigabytes in size. We can distribute this weight across multiple GPUs in order to reduce the memory consumption per GPU. - -### 1. Embedding parallelism -To enable this feature, add the flag `--use_parallel_embedding` to `convert_checkpoint.py`. - -### 2. The sharding dimension for embedding parallelism - -Assume the size of embedding lookup table is (vocab\_size \* hidden\_size), we can shard it along the vocab\_size (`--embedding_sharding_dim 0`) or hidden\_size (`--embedding_sharding_dim 1`) dimension. - -2.1 To shard the embedding lookup table along the hidden\_size dimension, set the flag `--use_parallel_embedding --embedding_sharding_dim 1`. Here is an example: - -```bash -# 2-way tensor parallelism with embedding parallelism along hidden dimension -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --tp_size 2 \ - --use_parallel_embedding \ - --embedding_sharding_dim 1 \ - --output_dir gpt2/trt_ckpt/fp16/2-gpu - -trtllm-build --checkpoint_dir gpt2/trt_ckpt/fp16/2-gpu \ - --output_dir gpt2/trt_engines/fp16/2-gpu -``` - -2.2 To shard the embedding lookup table along the vocab\_size dimension, set the flag `--use_parallel_embedding --embedding_sharding_dim 0`. In this case, you can optionally enable the lookup plugin when building the engines. - -```bash -# 2-way tensor parallelism with embedding parallelism along vocab dimension -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --tp_size 2 \ - --use_parallel_embedding \ - --embedding_sharding_dim 0 \ - --output_dir gpt2/trt_ckpt/fp16/2-gpu - -trtllm-build --checkpoint_dir gpt2/trt_ckpt/fp16/2-gpu \ - --output_dir gpt2/trt_engines/fp16/2-gpu -``` - -## GPT Variant - Granite - -For Granite, the steps are similar to StarCoder. - -```bash -# Download hf granite model -git clone https://huggingface.co/ibm-granite/granite-34b-code-instruct granite - -# Convert to TensorRT LLM checkpoint -python3 convert_checkpoint.py --model_dir granite \ - --dtype float16 \ - --gpt_variant starcoder \ - --tp_size 4 \ - --output_dir granite/trt_ckpt/fp16/4-gpu - -# Build TensorRT LLM engines -trtllm-build --checkpoint_dir granite/trt_ckpt/fp16/4-gpu \ - --gemm_plugin auto \ - --output_dir granite/trt_engines/fp16/4-gpu - -# Run inference -mpirun -np 4 \ - python3 ../../../run.py --engine_dir granite/trt_engines/fp16/4-gpu \ - --tokenizer_dir granite \ - --input_text "def print_hello_world():" \ - --max_output_len 20 -``` - -## GPT Variant - SantaCoder - -The SantaCoder extends the existing GPT model with multi-query attention mechanism. The following example shows building a 4-GPU engine and running simple prompt to generate the implementation of `print_hello_world()`. - -```bash -# Download hf santacoder model -git clone https://huggingface.co/bigcode/santacoder - -# Convert to TensorRT LLM checkpoint -python3 convert_checkpoint.py --model_dir santacoder \ - --dtype float16 \ - --tp_size 4 \ - --output_dir santacoder/trt_ckpt/fp16/4-gpu - -# Build TensorRT LLM engines -trtllm-build --checkpoint_dir santacoder/trt_ckpt/fp16/4-gpu \ - --gemm_plugin auto \ - --output_dir santacoder/trt_engines/fp16/4-gpu - -# Run inference -mpirun -np 4 \ - python3 ../../../run.py --engine_dir santacoder/trt_engines/fp16/4-gpu \ - --tokenizer_dir santacoder \ - --input_text "def print_hello_world():" \ - --max_output_len 20 -``` - - -## GPT Variant - StarCoder (v1 and v2) - -For StarCoder, the steps are similar to SantaCoder. - -```bash -# Download hf starcoder model -git clone https://huggingface.co/bigcode/starcoder - -# Convert to TensorRT LLM checkpoint -python3 convert_checkpoint.py --model_dir starcoder \ - --dtype float16 \ - --tp_size 4 \ - --output_dir starcoder/trt_ckpt/fp16/4-gpu - -# Build TensorRT LLM engines -trtllm-build --checkpoint_dir starcoder/trt_ckpt/fp16/4-gpu \ - --gemm_plugin auto \ - --output_dir starcoder/trt_engines/fp16/4-gpu - -# Run inference -mpirun -np 4 \ - python3 ../../../run.py --engine_dir starcoder/trt_engines/fp16/4-gpu \ - --tokenizer_dir starcoder \ - --input_text "def print_hello_world():" \ - --max_output_len 20 -``` - -For StarCoder2, you can use almost the same steps as shown above. - - Note that StarCoder2 hasn't been merged to the official releases of transformers package yet, so remember using the [main branch of transformers repo](https://github.com/huggingface/transformers). - - Add `--max_attention_window_size 4096` when running with run.py or summarization, which enables the sliding window attention. - - the sliding window size comes from the hf model [config.json](https://huggingface.co/bigcode/starcoder2-15b/blob/main/config.json#L23). - -### Run StarCoder2 with LoRA - -TensorRT LLM supports running StarCoder2 models with FP16/BF16/FP32 LoRA. In this section, we use starcoder2-15b as an example to show how to run an FP8 base model with FP16 LoRA module. - -* download the base model and lora model from HF - -```bash -git-lfs clone https://huggingface.co/bigcode/starcoder2-15b -git-lfs clone https://huggingface.co/KaQyn/peft-lora-starcoder2-15b-unity-copilot -``` - -* Quantize the StarCoder2 model to fp8 from HF -```bash -BASE_STARCODER2_MODEL=./starcoder2-15b -python ../../../quantization/quantize.py --model_dir ${BASE_STARCODER2_MODEL} \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir starcoder2-15b/trt_ckpt/fp8/1-gpu \ - --calib_size 512 -``` - -* Build engine and run inference. -```bash -trtllm-build --checkpoint_dir starcoder2-15b/trt_ckpt/fp8/1-gpu \ - --output_dir starcoder2-15b/trt_engines/fp8_lora/1-gpu \ - --gemm_plugin auto \ - --lora_plugin auto \ - --lora_dir ./peft-lora-starcoder2-15b-unity-copilot - -python ../../../run.py --engine_dir starcoder2-15b/trt_engines/fp8_lora/1-gpu \ - --max_output_len 20 \ - --tokenizer_dir ${BASE_STARCODER2_MODEL} \ - --input_text "def print_hello_world():" \ - --lora_task_uids 0 \ - --no_add_special_tokens \ - --use_py_session -``` - -## GPT-Next - -NVIDIA has released a GPT-like model with some architectural improvements, that you can find here: [https://huggingface.co/nvidia/GPT-2B-001](https://huggingface.co/nvidia/GPT-2B-001). This architecture is also supported by TensorRT-LLM. - -Different from Huggingface's checkpoint, you should specify the NeMo checkpoint path using `--nemo_ckpt_path` for `convert_checkpoint.py`. The script also extracts the tokenizer file from the NeMo checkpoint and saves it to the TensorRT LLM checkpoint folder, which can be used in the inference scripts. - -```bash -# Download NeMo checkpoint -wget https://huggingface.co/nvidia/GPT-2B-001/resolve/main/GPT-2B-001_bf16_tp1.nemo - -# Convert to TensorRT LLM checkpoint -# It also extracts the tokenizer file and saves to the TensorRT LLM checkpoint folder -python3 convert_checkpoint.py --nemo_ckpt_path GPT-2B-001_bf16_tp1.nemo \ - --dtype bfloat16 \ - --output_dir gpt-next-2B/trt_ckpt/bf16/1-gpu - -# Build TensorRT LLM engines -# --gpt_attention_plugin must be set for GPT-Next since Rotary positional embeddings (RoPE) is only supported by the gpt attention plugin at this time. -trtllm-build --checkpoint_dir gpt-next-2B/trt_ckpt/bf16/1-gpu \ - --output_dir gpt-next-2B/trt_engines/bf16/1-gpu - -# Run inference -python3 ../../../run.py --engine_dir gpt-next-2B/trt_engines/bf16/1-gpu \ - --vocab_file gpt-next-2B/trt_ckpt/bf16/1-gpu/tokenizer.model \ - --no_add_special_tokens \ - --max_output_len 8 -``` - -### Prompt-tuning - -For efficient fine-tuning, the NeMo framework allows you to learn virtual tokens to accomplish a downstream task. For more details, please read the -NeMo documentation [here](https://docs.nvidia.com/nemo-framework/user-guide/latest/overview.html). - -TensorRT LLM supports inference with those virtual tokens. To enable it, pass the prompt embedding table's maximum size at build time with `--max_prompt_embedding_table_size N`. For example: - -```bash -# Convert to TensorRT LLM checkpoint -python3 convert_checkpoint.py --nemo_ckpt_path megatron_converted_8b_tp4_pp1.nemo \ - --dtype float16 \ - --output_dir gpt-next-8B/trt_ckpt/fp16/1-gpu - -# Build TensorRT LLM engines with prompt-tuning enabled -trtllm-build --checkpoint_dir gpt-next-8B/trt_ckpt/fp16/1-gpu \ - --max_prompt_embedding_table_size 100 \ - --output_dir gpt-next-8B/trt_engines/fp16/1-gpu -``` - -You can now export the learned embedding table with: -```bash -python3 nemo_prompt_convert.py -i email_composition.nemo -o email_composition.npy -``` -It'll give you a summary of the different tasks in the table, that you can specify at runtime. - -Finally, you can run inference on pre-defined tokens: -```bash -python3 ../../../run.py --engine_dir gpt-next-8B/trt_engines/fp16/1-gpu \ - --vocab_file gpt-next-8B/trt_ckpt/fp16/1-gpu/tokenizer.model \ - --no_add_special_tokens \ - --prompt_table_path email_composition.npy \ - --prompt_tasks 0 \ - --max_output_len 8 -``` - - -### MultiLoRA with the Nemo checkpoint - -```bash -# Download NeMo checkpoint -wget https://huggingface.co/nvidia/GPT-2B-001/resolve/main/GPT-2B-001_bf16_tp1.nemo - -# Convert to TensorRT LLM checkpoint -python3 convert_checkpoint.py --nemo_ckpt_path GPT-2B-001_bf16_tp1.nemo \ - --dtype float16 \ - --output_dir gpt-next-2B/trt_ckpt/fp16/1-gpu - -# Build TensorRT LLM engines -trtllm-build --checkpoint_dir gpt-next-2B/trt_ckpt/fp16/1-gpu \ - --lora_plugin auto \ - --lora_dir gpt2b_lora-900.nemo gpt2b_lora-stories.nemo \ - --lora_ckpt_source "nemo" \ - --lora_target_modules attn_qkv \ - --max_batch_size 4 \ - --max_beam_width 2 \ - --max_input_len 512 \ - --max_seq_len 562 \ - --output_dir gpt-next-2B/trt_engines/fp16/1-gpu - -# Run inference directly from NeMo LoRA checkpoint -# --lora_task_ids correspond to the index of the models given with --lora_dir. -1 means no LoRA -python3 ../../../run.py --engine_dir gpt-next-2B/trt_engines/fp16/1-gpu \ - --vocab_file gpt-next-2B/trt_ckpt/fp16/1-gpu/tokenizer.model \ - --no_add_special_tokens \ - --max_output_len 20 \ - --use_py_session \ - --lora_task_uids 0 -1 1 \ - --input_text "After Washington had returned to Williamsburg, Dinwiddie ordered him to lead a larger force to assist Trent in his work. While en route, Washington learned of Trent's retreat. Since Tanaghrisson had promised support to the British, Washington continued toward Fort Duquesne and met with the Mingo leader. Learning of a French scouting party in the area, Washington, with Tanaghrisson and his party, surprised the Canadians on May 28 in what became known as the Battle of Jumonville Glen. They killed many of the Canadians, including their commanding officer, Joseph Coulon de Jumonville, whose head was reportedly split open by Tanaghrisson with a tomahawk. The historian Fred Anderson suggests that Tanaghrisson was acting to gain the support of the British and regain authority over his own people. They had been inclined to support the French, with whom they had long trading relationships. One of Tanaghrisson's men told Contrecoeur that Jumonville had been killed by British musket fire. Question: Upon learning of a French scounting party in the area, what did Washington do? Answer:" "You hold the job title in the Wizarding World of Harry Potter where you say random words looking for spells" "You hold the job title in the Wizarding World of Harry Potter where you say random words looking for spells" -``` - -The output would look like (Note that in this case the adapters have only been trained for a few epochs, so the result quality is poor): - -``` -...... -Input [Text 0]: "After Washington had returned to Williamsburg, Dinwiddie ordered him to lead a larger force to assist Trent in his work. While en route, Washington learned of Trent's retreat. Since Tanaghrisson had promised support to the British, Washington continued toward Fort Duquesne and met with the Mingo leader. Learning of a French scouting party in the area, Washington, with Tanaghrisson and his party, surprised the Canadians on May 28 in what became known as the Battle of Jumonville Glen. They killed many of the Canadians, including their commanding officer, Joseph Coulon de Jumonville, whose head was reportedly split open by Tanaghrisson with a tomahawk. The historian Fred Anderson suggests that Tanaghrisson was acting to gain the support of the British and regain authority over his own people. They had been inclined to support the French, with whom they had long trading relationships. One of Tanaghrisson's men told Contrecoeur that Jumonville had been killed by British musket fire. Question: Upon learning of a French scounting party in the area, what did Washington do? Answer:" -Output [Text 0 Beam 0]: "He surprised the Canadians on May 28 in what became known as the Battle of Jumonville" -Input [Text 1]: "You hold the job title in the Wizarding World of Harry Potter where you say random words looking for spells" -Output [Text 1 Beam 0]: ". - -The game is played with a deck of cards, and the player who has the most" -Input [Text 2]: "You hold the job title in the Wizarding World of Harry Potter where you say random words looking for spells" -Output [Text 2 Beam 0]: ". - -You are a wizard who is a wizard. - -You are a wizard who is" -``` diff --git a/examples/models/core/gpt/convert_checkpoint.py b/examples/models/core/gpt/convert_checkpoint.py deleted file mode 100644 index c8bbebe39925..000000000000 --- a/examples/models/core/gpt/convert_checkpoint.py +++ /dev/null @@ -1,343 +0,0 @@ -import argparse -import json -import os -import shutil -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed -from pathlib import Path - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import GPTForCausalLM -from tensorrt_llm.models.gpt.convert import (UnpackedNemoCheckpointDir, - copy_tokenizer_files, load_hf_gpt, - unpack_nemo_ckpt, - update_tokenizer_paths) -from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument('--nemo_ckpt_path', type=str, default=None) - parser.add_argument( - '--gpt_variant', - default=None, - choices=[ - None, 'gpt2', 'santacoder', 'starcoder', 'starcoder2', 'persimmon', - 'kosmos-2', 'nemotron' - ], - help= - "By default the script will try to infer the gpt_variant from model_dir. " - "Or users may overwrite gpt_variant by explicitly passing the variant.") - 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( - '--dtype', - type=str, - default='auto', - choices=['auto', 'float16', 'bfloat16', 'float32'], - help= - "The data type for the model weights and activations if not quantized. " - "If 'auto', the data type is automatically inferred from the source model; " - "however, if the source dtype is float32, it is converted to float16.") - parser.add_argument("--load_model_on_cpu", action="store_true") - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - - parser.add_argument( - '--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( - '--calib_dataset', - type=str, - default='lambada', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - 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( - '--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( - "--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("--dataset_cache_dir", - type=str, - default=None, - help="cache dir to load the hugging face dataset") - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument( - '--nemo_rename_key', - type=str, - nargs='+', - default=[], - help= - "Change a layer name when loading a NeMo checkpoint. Should follow :" - ) - - args = parser.parse_args() - - tensorrt_llm.logger.set_level(args.log_level) - return args - - -def args_to_quant_config(args: argparse.Namespace) -> QuantConfig: - '''return config dict with quantization info based on the command line args - ''' - quant_config = QuantConfig() - if args.use_weight_only: - if args.weight_only_precision == 'int8': - quant_config.quant_algo = QuantAlgo.W8A16 - elif args.weight_only_precision == 'int4': - quant_config.quant_algo = QuantAlgo.W4A16 - elif args.smoothquant: - quant_config.smoothquant_val = args.smoothquant - if args.per_channel: - if args.per_token: - quant_config.quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN - else: - quant_config.quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN - else: - if args.per_token: - quant_config.quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN - else: - quant_config.quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN - - if args.int8_kv_cache: - quant_config.kv_cache_quant_algo = QuantAlgo.INT8 - - # Check if model ckpt is pre-quantized to fp8. - hf_quant_config_path = f"{args.model_dir}/hf_quant_config.json" - if os.path.exists(hf_quant_config_path): - with open(hf_quant_config_path, 'r') as f: - hf_quant_config = json.load(f) - if hf_quant_config.get("producer", {}).get("name") == "modelopt": - modelopt_quant_config = hf_quant_config.get("quantization", {}) - if modelopt_quant_config.get("quant_algo", None) == QuantAlgo.FP8: - quant_config.quant_algo = QuantAlgo.FP8 - if modelopt_quant_config.get("kv_cache_quant_algo", - None) == QuantAlgo.FP8: - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - - return quant_config - - -def convert_and_save_hf(args): - model_dir = args.model_dir - load_model_on_cpu = args.load_model_on_cpu - world_size = args.tp_size * args.pp_size - - override_fields = { - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'gpt_variant': args.gpt_variant, - } - - quant_config = args_to_quant_config(args) - is_prequantized_to_fp8 = quant_config.quant_algo == QuantAlgo.FP8 - if is_prequantized_to_fp8: - override_fields.update({'is_prequantized_to_fp8': True}) - - if args.smoothquant is not None or args.int8_kv_cache: - mapping = Mapping(world_size=world_size, - tp_size=args.tp_size, - pp_size=args.pp_size) - GPTForCausalLM.quantize( - args.model_dir, - args.output_dir, - dtype=args.dtype, - mapping=mapping, - quant_config=quant_config, - device='cpu' if args.load_model_on_cpu else 'cuda', - calib_dataset=args.calib_dataset, - **override_fields) - else: - # Defer weight loading if checkpoint is prequantized to fp8. - if is_prequantized_to_fp8: - hf_model_or_dir = model_dir - else: - hf_model_or_dir = load_hf_gpt(model_dir, load_model_on_cpu) - - def convert_and_save_rank(args, rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - model = GPTForCausalLM.from_hugging_face(hf_model_or_dir, - args.dtype, - mapping=mapping, - quant_config=quant_config, - **override_fields) - model.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del model - - execute(args.workers, [convert_and_save_rank] * world_size, args) - release_gc() - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def convert_and_save_nemo(args): - world_size = args.tp_size * args.pp_size - quant_config = args_to_quant_config(args) - - override_fields = { - 'use_parallel_embedding': True, - 'embedding_sharding_dim': 0, - } - - nemo_ckpt_dir = os.path.join(args.output_dir, "unpacked") - nemo_ckpt_dir = unpack_nemo_ckpt(args.nemo_ckpt_path, nemo_ckpt_dir) - - def convert_and_save_rank(args, rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - model = GPTForCausalLM.from_nemo( - nemo_ckpt_dir, - dtype=args.dtype, - mapping=mapping, - quant_config=quant_config, - load_model_on_cpu=args.load_model_on_cpu, - nemo_rename_key=args.nemo_rename_key, - **override_fields) - model.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del model - - execute(args.workers, [convert_and_save_rank] * world_size, args) - release_gc() - - # Copy tokenizer files - unpacked_checkpoints_dir = UnpackedNemoCheckpointDir( - nemo_ckpt_dir, load_checkpoints_to_cpu=args.load_model_on_cpu) - nemo_model_config = unpacked_checkpoints_dir.model_config - tokenizer_config = update_tokenizer_paths( - nemo_model_config["tokenizer"], - unpacked_checkpoints_dir.get_all_tokenizer_file_paths()) - copy_tokenizer_files(tokenizer_config, Path(args.output_dir)) - - # Clean up unpacked nemo checkpoint - shutil.rmtree(nemo_ckpt_dir) - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - # TODO(qijun): Currently, the convert script depends on a torch op: - # torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix, - # which is included in tensorrt_llm Python package. Otherwise, the convert - # script does not need to import tensorrt_llm. Will remove it after reimplementing - # the op with PyTorch. - print(tensorrt_llm.__version__) - args = parse_arguments() - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - if args.model_dir is not None: - convert_and_save_hf(args) - elif args.nemo_ckpt_path is not None: - convert_and_save_nemo(args) - else: - raise NotImplementedError("No source model path specified!") - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/core/gpt/merge_ptuning_tables.py b/examples/models/core/gpt/merge_ptuning_tables.py deleted file mode 100644 index 7431541f8425..000000000000 --- a/examples/models/core/gpt/merge_ptuning_tables.py +++ /dev/null @@ -1,44 +0,0 @@ -#! /usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import argparse -from pathlib import Path -from typing import List - -import numpy as np - - -def combine_tables(input_tables: List[Path], output_table: Path): - tables = [np.load(table) for table in input_tables] - max_vocab_size = max(table.shape[1] for table in tables) - padded_tables = [ - np.pad(table, [(0, 0), (0, max_vocab_size - table.shape[1]), (0, 0)]) - for table in tables - ] - merged_table = np.concatenate(padded_tables, axis=0) - np.save(output_table, merged_table) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("input_tables", - type=Path, - nargs="+", - help="paths to tables to merge.") - parser.add_argument("output_table", - type=Path, - help="path where to save the combined table") - args = parser.parse_args() - combine_tables(args.input_tables, args.output_table) diff --git a/examples/models/core/gpt/nemo_lora_convert.py b/examples/models/core/gpt/nemo_lora_convert.py deleted file mode 100644 index 59e675402ec2..000000000000 --- a/examples/models/core/gpt/nemo_lora_convert.py +++ /dev/null @@ -1,206 +0,0 @@ -#! /usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import argparse -import datetime -import logging -import tempfile -from pathlib import Path - -import numpy as np -import torch -import yaml - -from tensorrt_llm._utils import str_dtype_to_torch, to_json_file, torch_to_numpy -from tensorrt_llm.lora_manager import LoraManager, get_all_nemo_lora_weights -from tensorrt_llm.models.gpt.convert import cpu_map_location, unpack_nemo_ckpt - -log_format = "%(asctime)s %(name)s [%(levelname)s] %(message)s" -logging.basicConfig(format=log_format) -LOGGER = logging.getLogger(__name__) - - -def get_lora_keys(layer_idx): - in_key = f'model.language_model.encoder.layers.{layer_idx}.self_attention.adapter_layer.lora_kqv_adapter.linear_in.weight' - out_key = f'model.language_model.encoder.layers.{layer_idx}.self_attention.adapter_layer.lora_kqv_adapter.linear_out.weight' - return in_key, out_key - - -def save_val(val, dir, key, tp_num=None, write_npy=False): - ext = "npy" if write_npy else "bin" - suffix = ext if tp_num is None else f"{tp_num}.{ext}" - if write_npy: - np.save(dir / f"model.{key}.{suffix}", val) - else: - val.tofile(dir / f"model.{key}.{suffix}") - - -def lora_convert(out_dir, lora_config, lora_weights, customization_id, - precision): - saved_dir = Path(out_dir) - saved_dir.mkdir(parents=True, exist_ok=True) - num_layers = int(lora_config["num_layers"]) - config = {"lora_config": {"lora_kqv_adapter": {}}} - config['lora_config']['precision'] = precision - layer_weights = get_all_nemo_lora_weights(lora_weights) - for layer_idx in range(num_layers): - linear_in_weight = layer_weights[layer_idx]['in'] - linear_out_weight = layer_weights[layer_idx]['out'] - config["lora_config"]["lora_kqv_adapter"]["0"] = { - "key": f"{customization_id}", - "low_rank": f"{linear_in_weight.shape[0]}", - } - - # do something else here. just choose some key instead of basing it on the nemo key - in_key, out_key = get_lora_keys(layer_idx) - - save_val( - torch_to_numpy( - linear_in_weight.transpose( - 1, 0).contiguous().to(dtype=str_dtype_to_torch(precision))), - saved_dir, - in_key.replace("lora_kqv_adapter", f"lora_kqv_adapter.{0}")) - save_val( - torch_to_numpy( - linear_out_weight.transpose( - 1, 0).contiguous().to(dtype=str_dtype_to_torch(precision))), - saved_dir, - out_key.replace("lora_kqv_adapter", f"lora_kqv_adapter.{0}")) - - to_json_file(config, saved_dir / "lora_weights.json") - - -def lora_convert_cpp_runtime(out_dir, - lora_config, - lora_weights, - precision='float16'): - saved_dir = Path(out_dir) - saved_dir.mkdir(parents=True, exist_ok=True) - num_layers = int(lora_config["num_layers"]) - weights = [] - weight_config = [] - layer_weights = get_all_nemo_lora_weights(lora_weights) - for layer_idx in range(num_layers): - in_weights = layer_weights[layer_idx]['in'] - out_weights = layer_weights[layer_idx]['out'] - LOGGER.debug(f"layer {layer_idx} in_weights: {in_weights.shape}") - LOGGER.debug(f"layer {layer_idx} out_weights: {out_weights.shape}") - in_out_weights = [] - adapter_size = 0 - for w, inout in ((in_weights, "in"), (out_weights, "out")): - assert len(w.shape) == 2 - # assume that the hidden dim is the larger of the 2 - dim0 = w.shape[0] - dim1 = w.shape[1] - adapter_size = min(dim0, dim1) - # in_weights should have shape [adaper_size, hidden] - if dim1 < dim0 and inout == "in": - adapter_size = dim1 - w = w.transpose(1, 0) - # out_weights should have shape [hidden, adapter_size] - elif dim0 < dim1 and inout == "out": - adapter_size = dim0 - w = w.transpose(1, 0) - - w = w.contiguous().flatten().to(dtype=str_dtype_to_torch(precision)) - in_out_weights.append(w) - in_out_weights = torch.concatenate(in_out_weights).flatten().numpy() - weights.append(in_out_weights) - weight_config.append( - np.array([ - LoraManager.LORA_MODULE_IDS["attn_qkv"], layer_idx, adapter_size - ], - dtype=np.int32)) - all_weights = np.expand_dims(np.stack(weights), 0) - all_configs = np.expand_dims(np.stack(weight_config), 0) - - save_val(all_weights, - saved_dir, - "lora_weights", - tp_num=None, - write_npy=True) - save_val(all_configs, saved_dir, "lora_config", tp_num=None, write_npy=True) - - -def main(args): - start_time = datetime.datetime.now() - with tempfile.TemporaryDirectory() as prompt_out_dir: - prompt_out_dir = Path(prompt_out_dir) - unpack_nemo_ckpt(args.in_file, prompt_out_dir) - LOGGER.info("Spent %s (h:m:s) to unpack NeMo prompt archive", - datetime.datetime.now() - start_time) - - model_weights_ckpt = "model_weights.ckpt" - with open(prompt_out_dir / "model_config.yaml") as f: - prompt_config = yaml.full_load(f) - LOGGER.debug(prompt_config) - - start_time = datetime.datetime.now() - weight_path = prompt_out_dir / model_weights_ckpt - - prompt_weights = torch.load( - weight_path, - map_location=cpu_map_location, - ) - if args.write_cpp_runtime_tensors: - lora_convert_cpp_runtime(args.out_dir, - prompt_config, - prompt_weights, - precision=args.storage_type) - else: - lora_convert(args.out_dir, - prompt_config, - prompt_weights, - args.customization_id, - precision=args.storage_type) - - LOGGER.info("Spent %s (h:m:s) to convert the prompt model", - datetime.datetime.now() - start_time) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - '--out-dir', - '-o', - type=Path, - help='path to output embedding table file in the .npy format', - required=True) - parser.add_argument('--in-file', - '-i', - type=Path, - help='path to input prompt-tuning checkpoint file', - required=True) - parser.add_argument("--verbose", - action="store_true", - help="Provide verbose messages") - parser.add_argument("--customization-id", type=str, default="lora") - parser.add_argument("--write-cpp-runtime-tensors", - action="store_true", - default=False) - parser.add_argument("--storage-type", - type=str, - default="float16", - choices=["float32", "float16", "bfloat16"]) - args = parser.parse_args() - - LOGGER.setLevel(logging.DEBUG if args.verbose else logging.INFO) - - print("\n=============== Argument ===============") - for key in vars(args): - print(f"{key}: {vars(args)[key]}") - print("========================================") - - main(args) diff --git a/examples/models/core/gpt/nemo_prompt_convert.py b/examples/models/core/gpt/nemo_prompt_convert.py deleted file mode 100755 index c55131759c9d..000000000000 --- a/examples/models/core/gpt/nemo_prompt_convert.py +++ /dev/null @@ -1,145 +0,0 @@ -#! /usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import argparse -import datetime -import logging -import tempfile -from pathlib import Path - -import numpy as np -import torch -import yaml - -from tensorrt_llm._utils import torch_to_numpy -from tensorrt_llm.models.gpt.convert import cpu_map_location, unpack_nemo_ckpt - -log_format = "%(asctime)s %(name)s [%(levelname)s] %(message)s" -logging.basicConfig(format=log_format) -LOGGER = logging.getLogger(__name__) - - -def prompt_convert(out_file, prompt_config, prompt_weights): - nemo_type = "peft_tuning" if "peft" in prompt_config else "prompt_learning" - - vtokens_embeddings = [] - vtokens_len = [] - - if nemo_type == "peft_tuning": - ptuning_key = "model.embedding.adapter_layer.ptuning_adapter.inference_table" - if ptuning_key not in prompt_weights: - key_match = "adapter_layer.ptuning_adapter" - for k in prompt_weights.keys(): - if key_match in k: - ptuning_key = k - break - else: - raise ValueError( - "Could not find a suitable ptuning key in Nemo dict." - f" Tried {ptuning_key} or any key matching *{key_match}*") - prompt_task_weights = prompt_weights[ptuning_key] - - if 'hidden_size' in prompt_config: - assert prompt_config['hidden_size'] == prompt_task_weights.shape[ - 1], "P-Tuning hidden size does not match the model's." - - vtokens_embeddings.append(prompt_task_weights) - vtokens_len.append(prompt_task_weights.shape[0]) - else: - prompt_templates = prompt_config["task_templates"] - actual_task_id = 0 - for task_name_id, prompt_task in enumerate(prompt_templates): - prompt_task_name = prompt_task["taskname"] - LOGGER.info(f"Task {actual_task_id}: {prompt_task['taskname']}") - prompt_task_weights = prompt_weights["prompt_table"].get( - f"prompt_table.{prompt_task_name}.prompt_embeddings.weight") - if prompt_task_weights is None: - continue - vtokens_embeddings.append(prompt_task_weights) - vtokens_len.append(prompt_task_weights.shape[0]) - actual_task_id += 1 - - max_vtoken_len = max(vtokens_len) - embedding_dim = vtokens_embeddings[0].shape[1] - - # pad tasks to longest task embedding table - for i, vtoken_emb_table in enumerate(vtokens_embeddings): - padded_table = torch.zeros((max_vtoken_len, embedding_dim)) - padded_table[:vtoken_emb_table.shape[0], :] = vtoken_emb_table - vtokens_embeddings[i] = padded_table - - vtokens_embeddings = torch.stack(vtokens_embeddings) - np.save(out_file, torch_to_numpy(vtokens_embeddings)) - - -def main(args): - start_time = datetime.datetime.now() - with tempfile.TemporaryDirectory() as prompt_out_dir: - prompt_out_dir = Path(prompt_out_dir) - unpack_nemo_ckpt(args.in_file, prompt_out_dir) - LOGGER.info("Spent %s (h:m:s) to unpack NeMo prompt archive", - datetime.datetime.now() - start_time) - - model_weights_ckpt = "model_weights.ckpt" - with open(prompt_out_dir / "model_config.yaml") as f: - prompt_config = yaml.full_load(f) - LOGGER.debug(prompt_config) - - start_time = datetime.datetime.now() - weight_path = prompt_out_dir / model_weights_ckpt - if not weight_path.exists(): - weight_path = prompt_out_dir / "mp_rank_00" / model_weights_ckpt - - prompt_weights = torch.load( - weight_path, - map_location=cpu_map_location, - ) - prompt_convert(args.out_file, prompt_config, prompt_weights) - - LOGGER.info("Spent %s (h:m:s) to convert the prompt model", - datetime.datetime.now() - start_time) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - '--out-file', - '-o', - type=Path, - help='path to output embedding table file in the .npy format', - required=True) - parser.add_argument('--in-file', - '-i', - type=Path, - help='path to input prompt-tuning checkpoint file', - required=True) - parser.add_argument("--storage-type", - "-t", - type=str, - default="fp32", - choices=["fp32", "fp16"]) - parser.add_argument("--verbose", - action="store_true", - help="Provide verbose messages") - args = parser.parse_args() - - LOGGER.setLevel(logging.DEBUG if args.verbose else logging.INFO) - - print("\n=============== Argument ===============") - for key in vars(args): - print(f"{key}: {vars(args)[key]}") - print("========================================") - - main(args) diff --git a/examples/models/core/gpt/requirements.txt b/examples/models/core/gpt/requirements.txt deleted file mode 100644 index 592e01e5ba6d..000000000000 --- a/examples/models/core/gpt/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score -SentencePiece>=0.1.99 diff --git a/examples/models/core/gpt/run_hf.py b/examples/models/core/gpt/run_hf.py deleted file mode 100644 index 850d44b04f03..000000000000 --- a/examples/models/core/gpt/run_hf.py +++ /dev/null @@ -1,162 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import argparse -import csv -from pathlib import Path - -import numpy as np -import torch -from transformers import AutoModelForCausalLM, AutoTokenizer, T5Tokenizer - -import tensorrt_llm - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--max_output_len', type=int, required=True) - parser.add_argument('--log_level', type=str, default='warning') - parser.add_argument('--model_dir', type=str, default='gpt2') - parser.add_argument('--data_type', - type=str, - choices=['fp32', 'fp16'], - default='fp32') - parser.add_argument('--input_text', - type=str, - default='Born in north-east France, Soyer trained as a') - 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('--tokenizer', - dest='tokenizer_path', - help="HF tokenizer config path", - default='gpt2') - parser.add_argument('--vocab_file', - help="Used for sentencepiece tokenizers") - return parser.parse_args() - - -def generate( - max_output_len: int, - log_level: str = 'error', - model_dir: str = 'gpt2', - data_type: str = 'fp32', - 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_path='gpt2', - vocab_file=None, -): - tensorrt_llm.logger.set_level(log_level) - - model = AutoModelForCausalLM.from_pretrained(model_dir, - trust_remote_code=True) - model.cuda() - if data_type == 'fp16': - model.half() - - if vocab_file is not None: - tokenizer = T5Tokenizer(vocab_file=vocab_file) - END_ID = 50256 - else: - tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) - END_ID = tokenizer.eos_token_id - - 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 != END_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, dtype=torch.int32, device='cuda') - input_lengths = torch.tensor([input_ids.size(1)], - dtype=torch.int32, - device='cuda') - max_input_length = torch.max(input_lengths).item() - else: - input_lengths = torch.tensor([len(x) for x in input_tokens], - dtype=torch.int32, - device='cuda') - max_input_length = torch.max(input_lengths).item() - input_ids = np.full((len(input_lengths), max_input_length), END_ID) - for i in range(len(input_lengths)): - input_ids[i][-len(input_tokens[i]):] = input_tokens[i] - input_ids = torch.tensor(input_ids, dtype=torch.int32, device='cuda') - - top_k = 1 - temperature = 1 - output_ids = model.generate(input_ids, - max_length=max_input_length + max_output_len, - top_k=top_k, - temperature=temperature, - eos_token_id=END_ID, - pad_token_id=END_ID) - torch.cuda.synchronize() - - 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}') - outputs = output_ids[b][max_input_length:].tolist() - output_text = tokenizer.decode(outputs) - print(f'Output: {output_text}') - - if output_csv is not None: - output_file = Path(output_csv) - output_file.parent.mkdir(exist_ok=True, parents=True) - outputs = output_ids[:, max_input_length:].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 = output_ids[:, max_input_length:].tolist() - np.save(output_file, np.array(outputs, dtype='int32')) - - -if __name__ == '__main__': - args = parse_arguments() - generate(**vars(args)) diff --git a/examples/models/core/internlm2/.gitignore b/examples/models/core/internlm2/.gitignore deleted file mode 100644 index 7ce339719a3e..000000000000 --- a/examples/models/core/internlm2/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -internlm* -tokenizer.model diff --git a/examples/models/core/internlm2/README.md b/examples/models/core/internlm2/README.md deleted file mode 100644 index 0c1f8de5b908..000000000000 --- a/examples/models/core/internlm2/README.md +++ /dev/null @@ -1,207 +0,0 @@ -# InternLM2 - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run InternLM2 7B / 20B models in TensorRT LLM on both single GPU, single node multi-GPU and multi-node multi-GPU. - -## Overview - -The TensorRT LLM InternLM2 implementation is based on the LLaMA model. The implementation can -be found in [model.py](../../../../tensorrt_llm/models/llama/model.py). -The TensorRT LLM InternLM2 example code lies in [`examples/models/core/internlm2`](./): - -* [`convert_checkpoint.py`](./convert_checkpoint.py) converts the Huggingface Model of InternLM2 into TensorRT LLM checkpoint. - - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`run.py`](../../../run.py) to run the inference on an input text; -* [`summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 / BF16 - * INT8 & INT4 Weight-Only - * Tensor Parallel - -## Usage - -The TensorRT LLM InternLM2 example code locates at [examples/models/core/internlm2](./). It takes HF weights as input, and builds the corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Build TensorRT engine(s) - -Please install required packages first to make sure the example uses matched `tensorrt_llm` version: - -```bash -pip install -r requirements.txt -``` - -TensorRT LLM InternLM2 builds TensorRT engine(s) from HF checkpoint. If no checkpoint directory is specified, TensorRT LLM will build engine(s) with dummy weights. - -InternLM2 has released several checkpoints of different size or capabilities under https://huggingface.co/internlm. Users can pick any one repository and follow instructions to prepare the checkpoint. - -Below examples use [internlm2-chat-7b](https://huggingface.co/internlm/internlm2-chat-7b) and [internlm2-chat-20b](https://huggingface.co/internlm/internlm2-chat-20b) and assume these repositories are cloned or linked under this directory, for example `./internlm2-chat-7b`. - -Normally `trtllm-build` only requires single GPU, but if you've already got all the GPUs needed for inference, you could enable parallel building to make the engine building process faster by adding `--workers` argument. Please note that currently `--workers` feature only supports single node. - -Here're some examples: - -```bash -# Build a single-GPU float16 engine from HF weights. -# gpt_attention_plugin is necessary in InternLM2. -# Try use_gemm_plugin to prevent accuracy issue. -cd examples/models/core/internlm2 - -# Convert the InternLM2 7B model using a single GPU and FP16. -python convert_checkpoint.py --model_dir ./internlm2-chat-7b/ \ - --dtype float16 \ - --output_dir ./internlm2-chat-7b/trt_engines/fp16/1-gpu/ -# Note: setting `--dtype bfloat16` to use bfloat16 precision. - -# BUild the InternLM2 7B model using a single GPU -trtllm-build --checkpoint_dir ./internlm2-chat-7b/trt_engines/fp16/1-gpu/ \ - --output_dir ./engine_outputs \ - --gemm_plugin float16 - -# Convert the InternLM2 7B model using a single GPU and apply INT8 weight-only quantization.. -python convert_checkpoint.py --model_dir ./internlm2-chat-7b/ \ - --dtype float16 \ - --output_dir ./internlm2-chat-7b/trt_engines/int8/1-gpu/ \ - --use_weight_only \ - --weight_only_precision int8 - -trtllm-build --checkpoint_dir ./internlm2-chat-7b/trt_engines/int8/1-gpu/ \ - --output_dir ./engine_outputs \ - --gemm_plugin float16 - -# Note: setting `--weight_only_precision int4` to use INT4 weight-only quantization - -# Build InternLM2 7B using 2-way tensor parallelism. -python convert_checkpoint.py --model_dir ./internlm2-chat-7b/ \ - --dtype float16 \ - --output_dir ./internlm2-chat-7b/trt_engines/fp16/2-gpu/ \ - --tp_size 2 - -trtllm-build --checkpoint_dir ./internlm2-chat-7b/trt_engines/fp16/2-gpu/ \ - --output_dir ./engine_outputs \ - --gemm_plugin float16 - -# Build InternLM2 20B using 2-way tensor parallelism. -python convert_checkpoint.py --model_dir ./internlm2-chat-20b/ \ - --dtype bfloat16 \ - --output_dir ./internlm2-chat-20b/trt_engines/bf16/2-gpu/ \ - --tp_size 2 --workers 2 - -trtllm-build --checkpoint_dir ./internlm2-chat-7b/trt_engines/bf16/2-gpu/ \ - --output_dir ./engine_outputs \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 -``` - -#### INT8 weight only - -Examples: - -```bash -cd examples/models/core/internlm2 - -# For 7B models -python convert_checkpoint.py --model_dir ./internlm2-chat-7b \ - --output_dir ./internlm2-chat-7b/w8a16/ \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int8 - -# Build 7B model with both INT8 weight-only -trtllm-build --checkpoint_dir ./internlm2-chat-7b/w8a16 \ - --output_dir ./engine_outputs \ - --gemm_plugin float16 -``` - - -```bash -cd examples/models/core/internlm2 - -# For 20B models -python convert_checkpoint.py --model_dir ./internlm2-chat-20b \ - --output_dir ./internlm2-chat-20b/w8a16 \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int8 - -# Build 20B model with both INT8 weight-only -trtllm-build --checkpoint_dir ./internlm2-chat-20b/w8a16 \ - --output_dir ./engine_outputs \ - --gemm_plugin float16 \ -``` - -### Run - -To run a TensorRT LLM InternLM2 model using the engines generated by `trtllm-build` - -```bash -# InternLM2 7B with fp16 -python ../../../run.py --max_output_len=120 \ - --input_text 'Tell me about yourself.' \ - --tokenizer_dir ./internlm2-chat-7b/ \ - --engine_dir=./internlm2-chat-7b/trt_engines/fp16/1-gpu/ - -# InternLM2 7B with bf16 -python ../../../run.py --max_output_len=120 \ - --input_text 'Tell me about yourself.' \ - --tokenizer_dir ./internlm2-chat-7b/ \ - --engine_dir=./internlm2-chat-7b/trt_engines/bf16/1-gpu/ - -# InternLM2 7B with int8 weight only quantization -python ../../../run.py --max_output_len=120 \ - --input_text 'Tell me about yourself.' \ - --tokenizer_dir ./internlm2-chat-7b/ \ - --engine_dir=./internlm2-chat-7b/trt_engines/weight_only/1-gpu/ - -# InternLM2 7B with fp16 and tensor parallelism -mpirun -n 2 --allow-run-as-root \ - python ../../../run.py --max_output_len=120 \ - --input_text 'Tell me about yourself.' \ - --tokenizer_dir ./internlm2-chat-7b/ \ - --engine_dir=./internlm2-chat-7b/trt_engines/fp16/2-gpu/ - -# InternLM2 20B with fp16 and tensor parallelism and pipeline parallelism -mpirun -n 4 --allow-run-as-root \ - python ../../../run.py --max_output_len=120 \ - --input_text 'Tell me about yourself.' \ - --tokenizer_dir ./internlm2-chat-7b/ \ - --engine_dir=./internlm2-chat-7b/trt_engines/bf16/4-gpu/ -``` - -### Summarization using the InternLM2 model - -```bash -# Run summarization using the InternLM2 7B model in FP16. -python ../../../summarize.py --test_trt_llm --test_hf \ - --hf_model_dir ./internlm2-chat-7b/ \ - --data_type fp16 \ - --engine_dir ./engine_outputs - -# Run summarization using the InternLM2 7B model quantized to w8a16. -python ../../../summarize.py --test_trt_llm --test_hf \ - --hf_model_dir ./internlm2-chat-7b/ \ - --data_type fp16 \ - --engine_dir ./engine_outputs - -# Run summarization using the InternLM2 7B model in FP16 using two GPUs. -mpirun -n 2 --allow-run-as-root \ - python ../../../summarize.py --test_trt_llm --test_hf \ - --hf_model_dir ./internlm2-chat-7b/ \ - --data_type fp16 \ - --engine_dir ./internlm2-chat-7b/trt_engines/fp16/2-gpu/ - -# Run summarization using the InternLM2 20B model in BF16 using 4 GPUs. -mpirun -n 4 --allow-run-as-root \ - python ../../../summarize.py --test_trt_llm --test_hf \ - --hf_model_dir ./internlm2-chat-20b/ \ - --data_type bf16 \ - --engine_dir ./internlm2-chat-20b/trt_engines/bf16/4-gpu/ -``` diff --git a/examples/models/core/internlm2/convert_checkpoint.py b/examples/models/core/internlm2/convert_checkpoint.py deleted file mode 100644 index 44c80d6d51ff..000000000000 --- a/examples/models/core/internlm2/convert_checkpoint.py +++ /dev/null @@ -1,549 +0,0 @@ -import argparse -import json -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Dict, Optional - -import numpy as np -import safetensors -import torch -from einops import rearrange -from transformers import AutoConfig, AutoModelForCausalLM - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import get_hf_rope_theta, release_gc -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.llama import convert - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - 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('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - - parser.add_argument( - '--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('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument('--log_level', type=str, default='info') - args = parser.parse_args() - - tensorrt_llm.logger.set_level(args.log_level) - return args - - -def get_qkv_weight(weight: torch.Tensor, - hidden_size: int, - num_heads: int, - tp_size: int, - tp_rank: int, - is_bias: bool, - num_kv_heads: Optional[int] = None) -> torch.Tensor: - """ Splits the QKV matrix according to tensor parallelism """ - head_size = hidden_size // num_heads - num_kv_groups = num_heads // num_kv_heads - mha_mode = num_kv_heads == num_heads - weight = rearrange(weight, - '(h gs d) dim -> h gs d dim', - gs=2 + num_kv_groups, - d=head_size) - q_w, k_w, v_w = torch.split(weight, [num_kv_groups, 1, 1], dim=1) - if is_bias: - q_w = q_w.ravel() - k_w = k_w.ravel() - v_w = v_w.ravel() - qkv_w = torch.cat((q_w, k_w, v_w)) - qkv_w = convert.split_qkv_bias_tp(qkv_w, num_heads, hidden_size, - tp_size, tp_rank) - else: - q_w = rearrange(q_w, 'h gs d dim -> (h gs d) dim') - k_w = rearrange(k_w, 'h gs d dim -> (h gs d) dim') - v_w = rearrange(v_w, 'h gs d dim -> (h gs d) dim') - if not mha_mode: - if num_kv_heads < tp_size: - k_w = convert.dup_kv_weight(k_w, num_kv_heads, tp_size) - v_w = convert.dup_kv_weight(v_w, num_kv_heads, tp_size) - assert (k_w.shape[0] % (tp_size * head_size)) == 0 - assert (v_w.shape[0] % (tp_size * head_size)) == 0 - wq = convert.split(q_w, tp_size, tp_rank) - wk = convert.split(k_w, tp_size, tp_rank) - wv = convert.split(v_w, tp_size, tp_rank) - qkv_w = torch.concat((wq, wk, wv)) - - else: - qkv_w = torch.cat([q_w, k_w, v_w], dim=0) - - qkv_w = convert.split_qkv_tp(qkv_w, num_heads, hidden_size, tp_size, - tp_rank) - return qkv_w - - -def get_tllm_linear_weight( - weight: torch.Tensor, - prefix: str, - bias: Optional[torch.Tensor] = None, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[f'{prefix}.weight'] = processed_torch_weights - results[f'{prefix}.per_channel_scale'] = torch_weight_scales - else: - results[f'{prefix}.weight'] = weight.contiguous() - - if bias is not None: - results[f'{prefix}.bias'] = bias - - return results - - -def convert_from_hf(hf_model, - hf_config, - mapping: Mapping, - dtype: str = 'float32', - use_parallel_embedding: bool = False, - sharding_dim: int = 0, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8): - weights = {} - tik = time.time() - - model_params = dict(hf_model.named_parameters()) - #This is for InternVL2 - if hf_config.architectures[0] == 'InternLM2ForCausalLM': - keys_to_rename = [ - key for key in model_params.keys() if 'language_model.' in key - ] - keys_to_delete = [ - key for key in model_params.keys() if 'vision_model.' in key - ] - for key in keys_to_rename: - keys_rename = key.replace('language_model.', '') - model_params[keys_rename] = model_params[key] - del model_params[key] - for key in keys_to_delete: - del model_params[key] - - dtype = getattr(torch, dtype) - num_attention_heads = hf_config.num_attention_heads - hidden_size = hf_config.hidden_size - vocab_size = hf_config.vocab_size - num_kv_heads = hf_config.num_key_value_heads - num_hidden_layers = hf_config.num_hidden_layers - layers_range = mapping.pp_layers(num_hidden_layers) - is_xcomposer2 = ( - hf_config.architectures[0] == 'InternLMXComposer2ForCausalLM') - lora_weights = {} - for l in layers_range: - prefix = f'model.layers.{l}' - tllm_prex = f'transformer.layers.{l - layers_range[0]}' - - qkv_weight = convert.get_weight(model_params, - f'{prefix}.attention.wqkv', dtype) - qkv_w = get_qkv_weight(qkv_weight, - hidden_size, - num_attention_heads, - mapping.tp_size, - mapping.tp_rank, - is_bias=False, - num_kv_heads=num_kv_heads) - - if is_xcomposer2: - lora_prefix = f'base_model.model.model.layers.{l}' - assert num_attention_heads % num_kv_heads == 0 - num_key_value_groups = num_attention_heads // num_kv_heads - - qkv_loraA = convert.get_weight(model_params, - prefix + '.attention.wqkv.Plora_A', - dtype) - qkv_loraB = convert.get_weight(model_params, - prefix + '.attention.wqkv.Plora_B', - dtype) - q_lora_name = f"{lora_prefix}.self_attn.q_proj" - k_lora_name = f"{lora_prefix}.self_attn.k_proj" - v_lora_name = f"{lora_prefix}.self_attn.v_proj" - - #save qkv_loraA to be (q/k/v)_loraA - lora_weights[f"{q_lora_name}.lora_A.weight"] = qkv_loraA - lora_weights[f"{k_lora_name}.lora_A.weight"] = qkv_loraA - lora_weights[f"{v_lora_name}.lora_A.weight"] = qkv_loraA - - qkv_lora_rank = qkv_loraB.shape[-1] - head_size = hidden_size // num_attention_heads - - qkv_loraB = qkv_loraB.reshape(-1, num_key_value_groups + 2, - head_size, qkv_lora_rank) - q_loraB = qkv_loraB[:, :num_key_value_groups, :, :].reshape( - -1, qkv_lora_rank).contiguous() - k_loraB = qkv_loraB[:, - -2, :, :].reshape(-1, - qkv_lora_rank).contiguous() - v_loraB = qkv_loraB[:, - -1, :, :].reshape(-1, - qkv_lora_rank).contiguous() - - #save (q/k/v)_loraB - lora_weights[f"{q_lora_name}.lora_B.weight"] = q_loraB - lora_weights[f"{k_lora_name}.lora_B.weight"] = k_loraB - lora_weights[f"{v_lora_name}.lora_B.weight"] = v_loraB - - wo_loraA = convert.get_weight(model_params, - prefix + '.attention.wo.Plora_A', - dtype) - wo_loraB = convert.get_weight(model_params, - prefix + '.attention.wo.Plora_B', - dtype) - - lora_weights[ - f"{lora_prefix}.self_attn.o_proj.lora_A.weight"] = wo_loraA - lora_weights[ - f"{lora_prefix}.self_attn.o_proj.lora_B.weight"] = wo_loraB - - mlp_gate_loraA = convert.get_weight( - model_params, prefix + '.feed_forward.w3.Plora_A', dtype) - mlp_gate_loraB = convert.get_weight( - model_params, prefix + '.feed_forward.w3.Plora_B', dtype) - lora_weights[ - f"{lora_prefix}.mlp.up_proj.lora_A.weight"] = mlp_gate_loraA - lora_weights[ - f"{lora_prefix}.mlp.up_proj.lora_B.weight"] = mlp_gate_loraB - - mlp_fc_loraA = convert.get_weight( - model_params, prefix + '.feed_forward.w1.Plora_A', dtype) - mlp_fc_loraB = convert.get_weight( - model_params, prefix + '.feed_forward.w1.Plora_B', dtype) - lora_weights[ - f"{lora_prefix}.mlp.gate_proj.lora_A.weight"] = mlp_fc_loraA - lora_weights[ - f"{lora_prefix}.mlp.gate_proj.lora_B.weight"] = mlp_fc_loraB - - mlp_proj_loraA = convert.get_weight( - model_params, prefix + '.feed_forward.w2.Plora_A', dtype) - mlp_proj_loraB = convert.get_weight( - model_params, prefix + '.feed_forward.w2.Plora_B', dtype) - lora_weights[ - f"{lora_prefix}.mlp.down_proj.lora_A.weight"] = mlp_proj_loraA - lora_weights[ - f"{lora_prefix}.mlp.down_proj.lora_B.weight"] = mlp_proj_loraB - - qkv_bias = None - if f'{prefix}.attention.wqkv.bias' in model_params: - qkv_bias = convert.get_bias(model_params, - f'{prefix}.attention.wqkv', dtype) - if qkv_bias is None: - qkv_b = None - else: - qkv_b = get_qkv_weight(qkv_bias, - hidden_size, - num_attention_heads, - mapping.tp_size, - mapping.tp_rank, - is_bias=True, - num_kv_heads=num_kv_heads) - weights.update( - get_tllm_linear_weight( - qkv_w, - f'{tllm_prex}.attention.qkv', - qkv_b, - use_weight_only, - plugin_weight_only_quant_type, - )) - - attn_dense_weight = convert.get_weight(model_params, - f'{prefix}.attention.wo', dtype) - attn_dense_w = convert.split_matrix_tp(attn_dense_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - attn_dense_bias = None - if f'{prefix}.attention.wo.bias' in model_params: - attn_dense_bias = convert.get_bias(model_params, - f'{prefix}.attention.wo', dtype) - - weights.update( - get_tllm_linear_weight( - attn_dense_w, - f'{tllm_prex}.attention.dense', - attn_dense_bias, - use_weight_only, - plugin_weight_only_quant_type, - )) - - mlp_fc_weight = convert.get_weight(model_params, - f'{prefix}.feed_forward.w1', dtype) - mlp_fc_w = convert.split_matrix_tp(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - mlp_fc_b = None - weights.update( - get_tllm_linear_weight(mlp_fc_w, f'{tllm_prex}.mlp.fc', mlp_fc_b, - use_weight_only, - plugin_weight_only_quant_type)) - - mlp_proj_weight = convert.get_weight(model_params, - f'{prefix}.feed_forward.w2', dtype) - mlp_proj_w = convert.split_matrix_tp(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - mlp_proj_bias = None - weights.update( - get_tllm_linear_weight(mlp_proj_w, f'{tllm_prex}.mlp.proj', - mlp_proj_bias, use_weight_only, - plugin_weight_only_quant_type)) - - mlp_gate_weight = convert.get_weight(model_params, - f'{prefix}.feed_forward.w3', dtype) - mlp_gate_w = convert.split_matrix_tp(mlp_gate_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - mlp_gate_bias = None - weights.update( - get_tllm_linear_weight(mlp_gate_w, f'{tllm_prex}.mlp.gate', - mlp_gate_bias, use_weight_only, - plugin_weight_only_quant_type)) - - # Layer norms do not use tensor parallelism - input_ln_weight = convert.get_weight(model_params, - f'{prefix}.attention_norm', dtype) - weights[f'{tllm_prex}.input_layernorm.weight'] = input_ln_weight - - post_ln_weight = convert.get_weight(model_params, f'{prefix}.ffn_norm', - dtype) - weights[f'{tllm_prex}.post_layernorm.weight'] = post_ln_weight - - release_gc() - - if is_xcomposer2: - torch.save(lora_weights, 'adapter_model.bin') - adapter_config = { - "base_model_name_or_path": - "Internlm-xcomposer-2-7b-vl-hf", - "bias": - "none", - "enable_lora": - None, - "fan_in_fan_out": - False, - "inference_mode": - True, - "lora_alpha": - 256.0, - "lora_dropout": - 0.05, - "merge_weights": - False, - "modules_to_save": - None, - "peft_type": - "LORA", - "r": - 256, - "target_modules": [ - "q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", - "down_proj", "up_proj" - ], - "task_type": - "CAUSAL_LM" - } - with open(os.path.join('adapter_config.json'), 'w') as f: - json.dump(adapter_config, f, indent=4) - - embed_w = convert.get_weight(model_params, 'model.tok_embeddings', dtype) - if use_parallel_embedding: - embed_w = convert.split_matrix_tp(embed_w, - mapping.tp_size, - mapping.tp_rank, - dim=sharding_dim) - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = embed_w - lm_head_weights = convert.get_weight(model_params, 'output', dtype) - if mapping.is_last_pp_rank(): - if vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = convert.pad_vocab_size(vocab_size, - mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - - lm_head_weights = torch.from_numpy( - np.pad(lm_head_weights.detach().cpu().numpy(), - ((0, pad_width), (0, 0)), - 'constant', - constant_values=0)) - weights['lm_head.weight'] = convert.split_matrix_tp(lm_head_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - ln_f_w = convert.get_weight(model_params, 'model.norm', dtype) - weights['transformer.ln_f.weight'] = ln_f_w - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -if __name__ == '__main__': - emit_engine_arch_deprecation("convert_checkpoint.py") - args = parse_arguments() - world_size = args.tp_size * args.pp_size - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - quant_algo = None - plugin_weight_only_quant_type = None - if args.use_weight_only and args.weight_only_precision == 'int8': - plugin_weight_only_quant_type = torch.int8 - quant_algo = 'W8A16' - elif args.use_weight_only and args.weight_only_precision == 'int4': - plugin_weight_only_quant_type = torch.quint4x2 - quant_algo = 'W4A16' - - hf_config = AutoConfig.from_pretrained(args.model_dir, - trust_remote_code=True) - #This is for InternVL2 - if hasattr(hf_config, 'llm_config'): - hf_config = hf_config.llm_config - - config = { - 'architecture': hf_config.architectures[0], - 'dtype': args.dtype, - 'logits_dtype': 'float32', - 'num_hidden_layers': hf_config.num_hidden_layers, - 'num_attention_heads': hf_config.num_attention_heads, - 'num_key_value_heads': hf_config.num_key_value_heads, - 'hidden_size': hf_config.hidden_size, - 'intermediate_size': hf_config.intermediate_size, - 'norm_epsilon': hf_config.rms_norm_eps, - 'vocab_size': hf_config.vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'rotary_base': get_hf_rope_theta(hf_config, 10000.0), - 'max_position_embeddings': hf_config.max_position_embeddings, - 'hidden_act': hf_config.hidden_act, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'quantization': { - 'quant_algo': quant_algo, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'has_partial_lora_mask': - hf_config.architectures[0] == 'InternLMXComposer2ForCausalLM', - 'attn_bias': getattr(hf_config, 'bias', False), - 'rotary_scaling': getattr(hf_config, "rope_scaling", None) - } - - with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: - json.dump(config, f, indent=4) - - def covert_and_save(rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - - hf_model = AutoModelForCausalLM.from_pretrained(args.model_dir, - trust_remote_code=True, - dtype="auto") - weights = convert_from_hf( - hf_model, - hf_config, - mapping, - dtype=args.dtype, - use_parallel_embedding=args.use_parallel_embedding, - sharding_dim=args.embedding_sharding_dim, - use_weight_only=args.use_weight_only, - plugin_weight_only_quant_type=plugin_weight_only_quant_type) - del hf_model - save_file = os.path.join(args.output_dir, f'rank{rank}.safetensors') - print(f'Saving to {save_file}') - safetensors.torch.save_file(weights, save_file) - - if args.workers == 1: - for rank in range(world_size): - covert_and_save(rank) - else: - with ThreadPoolExecutor(max_workers=args.workers) as p: - futures = [ - p.submit(covert_and_save, rank) for rank in range(world_size) - ] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') diff --git a/examples/models/core/internlm2/requirements.txt b/examples/models/core/internlm2/requirements.txt deleted file mode 100644 index d27fa26c6812..000000000000 --- a/examples/models/core/internlm2/requirements.txt +++ /dev/null @@ -1 +0,0 @@ -einops diff --git a/examples/models/core/llama/.gitignore b/examples/models/core/llama/.gitignore deleted file mode 100644 index 02915ec1a252..000000000000 --- a/examples/models/core/llama/.gitignore +++ /dev/null @@ -1,5 +0,0 @@ -llama* -tokenizer.model -*output* -*.safetensors -*.json diff --git a/examples/models/core/llama/README.md b/examples/models/core/llama/README.md deleted file mode 100644 index f427053284ce..000000000000 --- a/examples/models/core/llama/README.md +++ /dev/null @@ -1,1592 +0,0 @@ -# LLaMA - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a LLaMA model in TensorRT LLM on both single GPU, single node multi-GPU and multi-node multi-GPU. - -- [LLaMA](#llama) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [Build TensorRT engine(s)](#build-tensorrt-engines) - - [LLaMA v2 Updates](#llama-v2-updates) - - [LLaMA v3 Updates](#llama-v3-updates) - - [Long context length](#long-context-length) - - [Long context evaluation](#long-context-evaluation) - - [1M long context test case](#1m-long-context-test-case) - - [INT8 KV cache](#int8-kv-cache) - - [SmoothQuant](#smoothquant) - - [FP8 Post-Training Quantization](#fp8-post-training-quantization) - - [Groupwise quantization (AWQ/GPTQ)](#groupwise-quantization-awqgptq) - - [AWQ](#awq) - - [GPTQ](#gptq) - - [w4aINT8 quantization (QServe)](#w4aint8-quantization-qserve) - - [NVFP4 quantization](#nvfp4-quantization) - - [Run](#run) - - [Multi-GPU multi-node (MGMN) support](#multi-gpu-multi-node-mgmn-support) - - [Summarization using the LLaMA model](#summarization-using-the-llama-model) - - [Mistral v0.1](#mistral-v01) - - [Mistral Nemo](#mistral-nemo) - - [Running CodeLlama](#running-codellama) - - [Build](#build) - - [Run](#run-1) - - [Run models with LoRA](#run-models-with-lora) - - [Run LLaMa with several lora checkpoints](#run-llama-with-several-lora-checkpoints) - - [Run FP8 Mistral v0.1 with FP16 lora checkpoint](#run-fp8-mistral-v01-with-fp16-lora-checkpoint) - - [Run INT4-AWQ LLaMa with several FP16 lora checkpoints](#run-int4-awq-llama-with-several-fp16-lora-checkpoints) - - [Run LLaMa with StreamingLLM](#run-llama-with-streamingllm) - - [Run LLaMA-3.1 405B Model](#run-llama-31-405b-model) - - [Convert Checkpoint to TensorRT LLM Unified Checkpoint](#convert-checkpoint-to-tensorrt-llm-unified-checkpoint) - - [Build Engine](#build-engine) - - [Run Inference](#run-inference) - - [Run LLaMa-3.3 70B Model on PyTorch Backend](#run-llama-33-70b-model-on-pytorch-backend) - - [Prepare TensorRT LLM extra configs](#prepare-tensorrt-llm-extra-configs) - - [Launch trtllm-serve OpenAI-compatible API server](#launch-trtllm-serve-openai-compatible-api-server) - - [Run performance benchmarks](#run-performance-benchmarks) - -## 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/models/core/llama`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert the LLaMA model into TensorRT LLM checkpoint format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`run.py`](../../../run.py) to run the inference on an input text; -* [`summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * BF16/FP16 - * FP8 - * INT8 & INT4 Weight-Only - * SmoothQuant - * Groupwise quantization (AWQ/GPTQ) - * w4aINT8 quantization (QServe) - * FP8 KV CACHE - * INT8 KV CACHE (+ AWQ/per-channel weight-only) - * Tensor Parallel + Pipeline Parallel, Tensor Parallel + Context Parallel - * STRONGLY TYPED - -## Usage - -The TensorRT LLM LLaMA example code locates at [examples/models/core/llama](./). It takes HF weights as input, and builds the corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Build TensorRT engine(s) - -Please install required packages first to make sure the example uses matched `tensorrt_llm` version: - -```bash -pip install --upgrade -r requirements.txt -``` - -Need to prepare the HF LLaMA checkpoint by following the guides here https://huggingface.co/docs/transformers/main/en/model_doc/llama. - -The `trtllm-build` command builds TensorRT engine(s) from HF checkpoint. If no checkpoint directory is specified, TensorRT LLM will build engine(s) with dummy weights. - -`trtllm-build` command has a variety of options. In particular, the plugin-related options have two categories: -* Plugin options that requires a data type (e.g., `gpt_attention_plugin`), you can - * explicitly specify `float16`/`bfloat16`/`float32`, so that the plugins are enabled with the specified precision; - * implicitly specify `auto`, so that the plugins are enabled with the precision automatically inferred from model dtype (i.e., the dtype specified in weight conversion); or - * disable the plugin by `disable`. -* Other features that requires a boolean (e.g., `context_fmha`, `paged_kv_cache`, `remove_input_padding`), you can - * enable/disable the feature by specifying `enable`/`disable`. - -The defaults have been carefully tuned for better performance. For example, `gpt_attention_plugin`, `context_fmha`, `paged_kv_cache` and `remove_input_padding` are enabled by default. See more details by `trtllm-build --help`. - -Normally `trtllm-build` only requires single GPU, but if you've already got all the GPUs needed for inference, you could enable parallel building to make the engine building process faster by adding `--workers` argument. Please note that currently `workers` feature only supports single node. - -`--use_fused_mlp=enable` enables GEMM horizontal fusion in gated MLP layer, which reduces input traffic and potentially improves performance. For FP8 PTQ, the downside is slight reduction of accuracy because one of the quantization scaling factors are discarded (accuracy 0.45734 vs 0.45755 for LLaMA-v2 7B using modelopt/examples/hf/instruct_eval/mmlu.py). - -`--use_fused_mlp=enable --gemm_swiglu_plugin ` fuses 2 GEMMs without biases and SwiGLU into one kernel. This is a preview feature and is only supported for dtype `fp8`. The supported architecture is SM90. - -Here're some examples: - -```bash -# Build a single-GPU float16 engine from HF weights. -# Try use_gemm_plugin to prevent accuracy issue. - -# Build the LLaMA 7B model using a single GPU and FP16. -python convert_checkpoint.py --model_dir ./tmp/llama/7B/ \ - --output_dir ./tllm_checkpoint_1gpu_fp16 \ - --dtype float16 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_fp16 \ - --output_dir ./tmp/llama/7B/trt_engines/fp16/1-gpu \ - --gemm_plugin auto - -# Build the LLaMA 7B model using a single GPU and BF16. -python convert_checkpoint.py --model_dir ./tmp/llama/7B/ \ - --output_dir ./tllm_checkpoint_1gpu_bf16 \ - --dtype bfloat16 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_bf16 \ - --output_dir ./tmp/llama/7B/trt_engines/bf16/1-gpu \ - --gemm_plugin auto - -# Build the LLaMA 7B model using a single GPU and apply INT8 weight-only quantization. -python convert_checkpoint.py --model_dir ./tmp/llama/7B/ \ - --output_dir ./tllm_checkpoint_1gpu_fp16_wq \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int8 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_fp16_wq \ - --output_dir ./tmp/llama/7B/trt_engines/weight_only/1-gpu/ \ - --gemm_plugin auto - -# Build LLaMA 7B using 2-way tensor parallelism. -python convert_checkpoint.py --model_dir ./tmp/llama/7B/ \ - --output_dir ./tllm_checkpoint_2gpu_tp2 \ - --dtype float16 \ - --tp_size 2 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_2gpu_tp2 \ - --output_dir ./tmp/llama/7B/trt_engines/fp16/2-gpu/ \ - --gemm_plugin auto - -# Build LLaMA 7B using 2-way tensor parallelism and 2-way pipeline parallelism. -python convert_checkpoint.py --model_dir ./tmp/llama/7B/ \ - --output_dir ./tllm_checkpoint_4gpu_tp2_pp2 \ - --dtype float16 \ - --tp_size 2 \ - --pp_size 2 -trtllm-build --checkpoint_dir ./tllm_checkpoint_4gpu_tp2_pp2 \ - --output_dir ./tmp/llama/7B/trt_engines/fp16/4-gpu/ \ - --gemm_plugin auto - -# Build LLaMA 7B using 2-way tensor parallelism and 2-way context parallelism. -python convert_checkpoint.py --model_dir ./tmp/llama/7B/ \ - --output_dir ./tllm_checkpoint_4gpu_tp2_cp2 \ - --dtype float16 \ - --tp_size 2 \ - --cp_size 2 -trtllm-build --checkpoint_dir ./tllm_checkpoint_4gpu_tp2_cp2 \ - --output_dir ./tmp/llama/7B/trt_engines/fp16/4-gpu/ \ - --gemm_plugin auto - -# Build LLaMA 30B using 2-way tensor parallelism. -python convert_checkpoint.py --model_dir ./tmp/llama/30B/hf/ \ - --output_dir ./tllm_checkpoint_2gpu_tp2 \ - --dtype float16 \ - --tp_size 2 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_2gpu_tp2 \ - --output_dir ./tmp/llama/30B/trt_engines/fp16/2-gpu/ \ - --gemm_plugin auto -``` - -#### LLaMA v2 Updates -The LLaMA v2 models with 7B and 13B are compatible with the LLaMA v1 implementation. The above -commands still work. - - -For LLaMA v2 70B, there is a restriction on tensor parallelism that the number of KV heads -must be **divisible by the number of GPUs**. For example, since the 70B model has 8 KV heads, you can run it with -2, 4 or 8 GPUs (1 GPU as well for FP8). - - -```bash -# Build LLaMA 70B using 8-way tensor parallelism. -python convert_checkpoint.py --model_dir ./tmp/llama/70B/hf/ \ - --output_dir ./tllm_checkpoint_8gpu_tp8 \ - --dtype float16 \ - --tp_size 8 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_8gpu_tp8 \ - --output_dir ./tmp/llama/70B/trt_engines/fp16/8-gpu/ \ - --gemm_plugin auto - -# Build LLaMA 70B using 4-way tensor parallelism and 2-way pipeline parallelism. -python convert_checkpoint.py --model_dir ./tmp/llama/70B/hf/ \ - --output_dir ./tllm_checkpoint_8gpu_tp4_pp2 \ - --dtype float16 \ - --tp_size 4 \ - --pp_size 2 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_8gpu_tp4_pp2 \ - --output_dir ./tmp/llama/70B/trt_engines/fp16/8-gpu/ \ - --gemm_plugin auto - -# Build LLaMA 70B TP=8 using Meta checkpoints directly. -python convert_checkpoint.py --meta_ckpt_dir ./tmp/llama/70B/ \ - --output_dir ./tllm_checkpoint_8gpu_tp8 \ - --dtype float16 \ - --tp_size 8 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_8gpu_tp8 \ - --output_dir ./tmp/llama/70B/trt_engines/fp16/8-gpu/ \ - --gemm_plugin auto -``` - -Same instructions can be applied to fine-tuned versions of the LLaMA v2 models (e.g. 7Bf or llama-2-7b-chat). - -#### LLaMA v3 Updates -The LLaMA 3.0 models with 8B and 70b are compatible with the LLaMA v2 implementation. The above -commands still work. - -Note that the `rope_theta` and `vocab_size` are larger in LLaMA v3 models and these values are now inferred -or pickup up from the `params.json` when using the `meta_ckpt_dir`. - -LLaMA 3.2 models are also supported now. For text only model like [Llama-3.2-1B](https://huggingface.co/meta-llama/Llama-3.2-1B), the steps are same to v3.0. For vision model like [Llama-3.2-11B-Vision](https://huggingface.co/meta-llama/Llama-3.2-11B-Vision), please refer to the [examples/models/core/mllama/README.md](../mllama/README.md) - -```bash -# Build LLaMA v3 8B TP=1 using HF checkpoints directly. -python convert_checkpoint.py --model_dir ./tmp/llama/8B/hf/ \ - --output_dir ./tllm_checkpoint_1gpu_tp1 \ - --dtype float16 \ - --tp_size 1 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_tp1 \ - --output_dir ./tmp/llama/8B/trt_engines/fp16/1-gpu/ \ - --gemm_plugin auto - -# Build LLaMA v3 8B TP=1 using Meta checkpoints directly. -python convert_checkpoint.py --meta_ckpt_dir ./tmp/llama/8B/ \ - --output_dir ./tllm_checkpoint_1gpu_tp1 \ - --dtype float16 \ - --tp_size 1 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_tp1 \ - --output_dir ./tmp/llama/8B/trt_engines/fp16/1-gpu/ \ - --gemm_plugin auto - -# Build LLaMA v3 70B using 8-way tensor parallelism. -python convert_checkpoint.py --model_dir ./tmp/llama/70B/hf/ \ - --output_dir ./tllm_checkpoint_8gpu_tp8 \ - --dtype float16 \ - --tp_size 8 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_8gpu_tp8 \ - --output_dir ./tmp/llama/70B/trt_engines/fp16/8-gpu/ \ - --gemm_plugin auto - -# Build LLaMA v3 70B using 4-way tensor parallelism and 2-way pipeline parallelism. -python convert_checkpoint.py --model_dir ./tmp/llama/70B/hf/ \ - --output_dir ./tllm_checkpoint_8gpu_tp4_pp2 \ - --dtype float16 \ - --tp_size 4 \ - --pp_size 2 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_8gpu_tp4_pp2 \ - --output_dir ./tmp/llama/70B/trt_engines/fp16/8-gpu/ \ - --gemm_plugin auto - -# Build LLaMA v3 70B TP=8 using Meta checkpoints directly. -python convert_checkpoint.py --meta_ckpt_dir ./tmp/llama/70B/ \ - --output_dir ./tllm_checkpoint_8gpu_tp8 \ - --dtype float16 \ - --tp_size 8 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_8gpu_tp8 \ - --output_dir ./tmp/llama/70B/trt_engines/fp16/8-gpu/ \ - --gemm_plugin auto -``` - -Same instructions can be applied to fine-tuned versions of the LLaMA v2 models (e.g. 7Bf or llama-2-7b-chat). - -### Long context length -With long context lengths, multi_block_mode is turned on by default to enable faster decoding in multi-head attention. To disable this feature, add `--multi_block_mode=False` to the runtime command. - - -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 convert_checkpoint.py --model_dir ./tmp/LongAlpaca-70B/ \ - --output_dir ./tllm_checkpoint_8gpu_tp8 \ - --dtype float16 \ - --tp_size 8 \ - -trtllm-build --checkpoint_dir ./tllm_checkpoint_8gpu_tp8 \ - --output_dir ./tmp/llama/70B/trt_engines/fp16/8-gpu/ \ - --gemm_plugin auto - -# Get the long text data from Gutenberg Project -wget https://www.gutenberg.org/cache/epub/64317/pg64317.txt - -# Replace the line breaks with special character '\n' and append "Summarize this story:" at end of text -awk '{printf "%s\\n", $0} END {printf "\\nSummarize this story:"}' pg64317.txt > pg64317_sanitized.txt - -# Run with 8 GPUs -# Notice, `--max_input_length ` 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 \ - --max_input_length 32768 \ - --input_file pg64317_sanitized.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. - -#### Long context evaluation - -* Download dataset and model - -```bash -git-lfs clone https://huggingface.co/datasets/DKYoon/SlimPajama-6B -git-lfs clone https://huggingface.co/gradientai/Llama-3-8B-Instruct-Gradient-1048k/ -``` - -* Run examples with max_input_len 16384 - -To evaluate the PPL of very long context, we need to enable `use_paged_context_fmha` and setup `max_num_tokens` to enable the chunked context inference, reducing the activation memory requirement. Also, we need to enable `gather_context_logits` to return the logits to compute the PPL. - -```bash -python examples/models/core/llama/convert_checkpoint.py --model_dir ./Llama-3-8B-Instruct-Gradient-1048k/ \ - --output_dir /tmp/llama-3-8B-1048k/trt_ckpts \ - --dtype float16 - -python -m tensorrt_llm.commands.build --checkpoint_dir /tmp/llama-3-8B-1048k/trt_ckpts \ - --output_dir /tmp/llama-3-8B-1048k/trt_engines \ - --gemm_plugin float16 \ - --gather_context_logits \ - --max_num_tokens 4096 \ - --max_input_len 16384 \ - --max_seq_len 16394 \ - --use_paged_context_fmha enable - -python ./examples/summarize.py --test_trt_llm \ - --tokenizer_dir ./Llama-3-8B-Instruct-Gradient-1048k/ \ - --data_type fp16 \ - --engine_dir /tmp/llama-3-8B-1048k/trt_engines \ - --eval_task eval_context_ppl \ - --max_input_len 16384 \ - --use_py_session \ - --dataset_dir ./SlimPajama-6B/ -``` - -* Run evaluation on passkey task - -To evaluate the accuracy of very long context on `needle in haystack`, we need to enable `use_paged_context_fmha` and setup `max_num_tokens` to enable the chunked context inference, reducing the activation memory requirement. To save memory, we don't enable the `gather_context_logits` here because we don't need logits. - -```bash -python3 examples/infinitebench/construct_synthetic_dataset.py --test_case build_passkey --test_level 4 - -python -m tensorrt_llm.commands.build --checkpoint_dir /tmp/llama-3-8B-1048k/trt_ckpts \ - --output_dir /tmp/llama-3-8B-1048k/trt_engines \ - --gemm_plugin float16 \ - --max_num_tokens 4096 \ - --max_input_len 131072 \ - --max_seq_len 131082 \ - --use_paged_context_fmha enable - -python examples/eval_long_context.py --task passkey \ - --engine_dir /tmp/llama-3-8B-1048k/trt_engines \ - --tokenizer_dir ./Llama-3-8B-Instruct-Gradient-1048k/ \ - --stop_idx 10 \ - --max_input_length 131072 \ - --enable_chunked_context \ - --max_tokens_in_paged_kv_cache 131136 -``` - -* Run evaluation on kv_retrieval - -`kv_retrieval` is harder than `passkey` and is helpful to distinguish the model capability. - -To run the kv_retrieval, we need a third-party repo to prepare the keys. - -```bash -git clone git@github.com:nelson-liu/lost-in-the-middle.git -pip install -r lost-in-the-middle/requirements.txt -python -u lost-in-the-middle/scripts/make_kv_retrieval_data.py --num-keys 3000 --num-examples 500 --output-path kv-retrieval-3000_keys.jsonl.gz -gzip -d kv-retrieval-3000_keys.jsonl.gz -``` - -Prepare input data and run evaluation. - -```bash -python examples/infinitebench/construct_synthetic_dataset.py --test_case build_kv_retrieval --test_level 0 - -python examples/models/core/llama/convert_checkpoint.py --model_dir ./Llama-3-8B-Instruct-Gradient-1048k/ \ - --output_dir /tmp/llama-3-8B-1048k/trt_ckpts \ - --dtype float16 \ - --tp_size 1 - -python -m tensorrt_llm.commands.build --checkpoint_dir /tmp/llama-3-8B-1048k/trt_ckpts \ - --output_dir /tmp/llama-3-8B-1048k/trt_engines \ - --gemm_plugin float16 \ - --max_num_tokens 4096 \ - --max_input_len 131072 \ - --max_seq_len 131082 \ - --use_paged_context_fmha enable - -python examples/eval_long_context.py --task kv_retrieval \ - --engine_dir /tmp/llama-3-8B-1048k/trt_engines \ - --tokenizer_dir ./Llama-3-8B-Instruct-Gradient-1048k/ \ - --stop_idx 10 \ - --max_input_length 131072 \ - --enable_chunked_context \ - --max_tokens_in_paged_kv_cache 131136 \ - --tensorrt_llm_accuracy_threshold 0.6 -``` - -expected results: - -```bash -[05/28/2024-03:31:43] [TRT-LLM] [I] ==== Evaluation ==== -[05/28/2024-03:31:43] [TRT-LLM] [I] # examples: 500 -[05/28/2024-03:31:43] [TRT-LLM] [I] Start index: 0 -[05/28/2024-03:31:43] [TRT-LLM] [I] Stop index: 10 -[05/28/2024-03:31:43] [TRT-LLM] [I] Max tokens: 50 -[05/28/2024-03:34:50] [TRT-LLM] [I] Compute the score -10it [00:00, 131072.00it/s] -[05/28/2024-03:34:51] [TRT-LLM] [I] Evaluation takes: 187.19733428955078 sec. -[05/28/2024-03:34:51] [TRT-LLM] [I] accuracy of 10 examples: 0.6 -``` - -#### 1M long context test case - -- Prepare 1M needle-in-a-haystack datasets - -```bash -python examples/infinitebench/construct_synthetic_dataset.py --test_case build_passkey --test_level 7 -``` - -- Llama-3-8B example - -```bash -git-lfs clone https://huggingface.co/gradientai/Llama-3-8B-Instruct-Gradient-1048k/ - -python examples/models/core/llama/convert_checkpoint.py --model_dir ./Llama-3-8B-Instruct-Gradient-1048k/ \ - --output_dir /tmp/llama-3-8B-1048k/trt_ckpts \ - --dtype float16 \ - --tp_size 4 - -python -m tensorrt_llm.commands.build --checkpoint_dir /tmp/llama-3-8B-1048k/trt_ckpts \ - --output_dir /tmp/llama-3-8B-1048k/trt_engines \ - --gemm_plugin float16 \ - --max_num_tokens 4096 \ - --max_batch_size 1 \ - --max_seq_len 1048576 \ - --use_paged_context_fmha enable \ - --workers 4 - -mpirun -n 4 --allow-run-as-root python examples/eval_long_context.py --task passkey \ - --engine_dir /tmp/llama-3-8B-1048k/trt_engines \ - --tokenizer_dir ./Llama-3-8B-Instruct-Gradient-1048k/ \ - --stop_idx 1 \ - --max_input_length 1048566 \ - --enable_chunked_context \ - --max_tokens_in_paged_kv_cache 1100000 -``` - -- Llama-3-70B example - -For the 70B model, at least 8 A100 80GB GPUs are required. - -```bash -git-lfs clone https://huggingface.co/gradientai/Llama-3-70B-Instruct-Gradient-1048k/ - -python examples/models/core/llama/convert_checkpoint.py --model_dir ./Llama-3-70B-Instruct-Gradient-1048k/ \ - --output_dir /tmp/llama-3-70B-1048k/trt_ckpts \ - --dtype float16 \ - --tp_size 8 - -python -m tensorrt_llm.commands.build --checkpoint_dir /tmp/llama-3-70B-1048k/trt_ckpts \ - --output_dir /tmp/llama-3-70B-1048k/trt_engines \ - --gemm_plugin float16 \ - --max_num_tokens 4096 \ - --max_batch_size 1 \ - --max_seq_len 1048576 \ - --use_paged_context_fmha enable \ - --workers 8 - -mpirun -n 8 --allow-run-as-root python examples/eval_long_context.py --task passkey \ - --engine_dir /tmp/llama-3-70B-1048k/trt_engines \ - --tokenizer_dir ./Llama-3-70B-Instruct-Gradient-1048k/ \ - --stop_idx 1 \ - --max_input_length 1048566 \ - --enable_chunked_context \ - --max_tokens_in_paged_kv_cache 1100000 -``` - -expected result: - -```bash -[05/27/2024-10:30:45] [TRT-LLM] [I] Compute the score -1it [00:00, 4215.38it/s] -[05/27/2024-10:30:45] [TRT-LLM] [I] accuracy of 1 examples: 1.0 -``` - -### INT8 KV cache -INT8 KV cache could be enabled to reduce memory footprint. It will bring more performance gains when batch size gets larger. - -For INT8 KV cache, [`convert_checkpoint.py`](./convert_checkpoint.py) features a -`--int8_kv_cache` option. Setting `--int8_kv_cache` will calibrate the model, -and then export the scaling factors needed for INT8 KV cache inference. - -Example: - -```bash -python convert_checkpoint.py --model_dir ./llama-models/llama-7b-hf \ - --output_dir ./llama-models/llama-7b-hf/int8_kv_cache/ \ - --dtype float16 \ - --int8_kv_cache -``` - -[`convert_checkpoint.py`](./convert_checkpoint.py) add new options for the support of INT8 KV cache. - - -**INT8 KV cache + per-channel weight-only quantization** - -INT8 KV cache could be combined with per-channel 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 convert_checkpoint.py --model_dir ./llama-models/llama-7b-hf \ - --output_dir ./tllm_checkpoint_1gpu_int8_kv_wq \ - --dtype float16 \ - --int8_kv_cache \ - --use_weight_only \ - --weight_only_precision int8 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_int8_kv_wq \ - --output_dir ./tmp/llama/7B/trt_engines/int8_kv_cache_weight_only/1-gpu \ - --gemm_plugin auto -``` - -Test with `summarize.py`: - -```bash -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** - -In addition, you can enable INT8 KV cache together with AWQ (per-group INT4 weight-only quantization)like the following command. - -```bash -python ../../../quantization/quantize.py --model_dir /tmp/llama-7b-hf \ - --output_dir ./tllm_checkpoint_1gpu_awq_int8_kv_cache \ - --dtype float16 \ - --qformat int4_awq \ - --awq_block_size 128 \ - --kv_cache_dtype int8 \ - --calib_size 32 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_awq_int8_kv_cache \ - --output_dir ./tmp/llama/7B/trt_engines/int8_kv_cache_int4_AWQ/1-gpu/ \ - --gemm_plugin auto \ -``` - -Test with `summarize.py`: - -```bash -python ../../../summarize.py --test_trt_llm \ - --hf_model_dir /tmp/llama-7b-hf \ - --data_type fp16 \ - --engine_dir ./tmp/llama/7B/trt_engines/int8_kv_cache_int4_AWQ/1-gpu \ - --test_hf -``` - -### SmoothQuant - -The smoothquant supports both LLaMA v1 and LLaMA 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 convert_checkpoint.py --model_dir /llama-models/llama-7b-hf --output_dir /tmp/tllm_checkpoint_1gpu_sq --dtype float16 --smoothquant 0.5 -trtllm-build --checkpoint_dir /tmp/tllm_checkpoint_1gpu_sq \ - --output_dir ./engine_outputs \ - --gemm_plugin auto -``` - -[`convert_checkpoint.py`](./convert_checkpoint.py) add new options for the support of INT8 inference of SmoothQuant models. - -`--smoothquant` is the starting point of INT8 inference. By default, it -will run the model in the _per-tensor_ mode. - -Then, you can add any combination of `--per-token` and `--per-channel` to get the corresponding behaviors. - -Examples of build invocations: - -```bash -# Build model for SmoothQuant in the _per_token_ + _per_channel_ mode -python3 convert_checkpoint.py --model_dir /llama-models/llama-7b-hf \ - --output_dir /tmp/tllm_checkpoint_1gpu_sq \ - --dtype float16 \ - --smoothquant 0.5 \ - --per_token \ - --per_channel - -trtllm-build --checkpoint_dir /tmp/tllm_checkpoint_1gpu_sq \ - --output_dir ./engine_outputs \ - --gemm_plugin auto -``` - -### FP8 Post-Training Quantization - -The examples below uses the NVIDIA Modelopt (AlgorithMic Model Optimization) toolkit for the model quantization process. - -First make sure Modelopt toolkit is installed (see [examples/quantization/README.md](/examples/quantization/README.md#preparation)) - - -```bash -# Quantize HF LLaMA 70B into FP8 and export trtllm checkpoint -python ../../../quantization/quantize.py --model_dir ./tmp/llama/70B \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir ./tllm_checkpoint_2gpu_fp8 \ - --calib_size 512 \ - --tp_size 2 - -# Build trtllm engines from the trtllm checkpoint -# Enable fp8 context fmha to get further acceleration by setting `--use_fp8_context_fmha enable` -trtllm-build --checkpoint_dir ./tllm_checkpoint_2gpu_fp8 \ - --output_dir ./engine_outputs \ - --gemm_plugin auto \ - --workers 2 -``` - -**Note**: A LLaMA 70B model with BF16 is about 140GB, a LLaMA 70B model with FP8 is about 70GB. -The peak GPU memory consumption when doing FP8 quantizaton is more than 210GB (there is also some activation memory occupation when doing calibration). -So you need a node with at least 4 H100(A100) to run the quantization command. After quantization, 2 GPUs are okay to for building and run. - -Note: use FP8 GEMV to optimize performance in FP8 small-batch-size cases. - -```bash -# Quantize HF LLaMA 7B into FP8 and export trtllm checkpoint -python ../../../quantization/quantize.py --model_dir /tmp/llama-7b-hf \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir ./tllm_checkpoint_1gpu_fp8 \ - --calib_size 512 - -# Build trtllm engines from the trtllm checkpoint -# Enable fp8 gemm plugin to get acceleration in small-batch-size cases -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_fp8 \ - --output_dir ./engine_outputs \ - --gemm_plugin fp8 -``` - -**Note**: FP8 gemv plugin uses CUDA cores to compute, by contrast to Tensor Core gemm kernel within cuBLAS. Over last year, as cuBLAS have improved their performance by a lot under small M case for Hopper(sm90), FP8 gemv kernel may or may not surpass cuBLAS, depending on specific gemm problem shape. Nonetheless, we still strongly recommend FP8 gemv kernel for Ada (sm89) as cuBLAS still falls behind gemv on it. - -### Groupwise quantization (AWQ/GPTQ) -One can enable AWQ/GPTQ INT4 weight only quantization with these options when building engine with `trtllm-build`: - -- `--use_weight_only` enables weight only GEMMs in the network. -- `--per_group` enable groupwise weight only quantization, for GPT-J 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`. -- `--quant_ckpt_path` passes the quantized checkpoint to build the engine. - -AWQ/GPTQ examples below involves 2 steps: -1. Weight quantization -2. Build TRT-LLM engine - -#### AWQ -1. Weight quantization: - - NVIDIA Modelopt toolkit is used for AWQ weight quantization. Please see [examples/quantization/README.md](/examples/quantization/README.md#preparation) for Modelopt installation instructions. - - ```bash - # Quantize HF LLaMA 7B checkpoint into INT4 AWQ format - python ../../../quantization/quantize.py --model_dir ./tmp/llama-7b-hf \ - --dtype float16 \ - --qformat int4_awq \ - --awq_block_size 128 \ - --output_dir ./quantized_int4-awq \ - --calib_size 32 - ``` - HF checkpoints generated with [AutoAWQ](https://github.com/casper-hansen/AutoAWQ) are also supported through the following conversion script: - - ```bash - # Convert AutoAWQ HF checkpoints into TRT-LLM checkpoint - python convert_checkpoint.py --model_dir ./tmp/Llama-2-7B-AWQ \ - --output_dir ./quantized_int4-awq - ``` - -2. Build TRT-LLM engine: - - ```bash - trtllm-build --checkpoint_dir ./quantized_int4-awq \ - --output_dir ./tmp/llama/7B/trt_engines/int4_AWQ/1-gpu/ \ - --gemm_plugin auto - ``` - -#### GPTQ -To run the GPTQ LLaMa example, the following steps are required: - -1. Weight quantization: - - Quantized weights for GPTQ are generated using [AutoGPTQ](https://github.com/AutoGPTQ/AutoGPTQ) as follow: - - ```bash - git clone https://github.com/AutoGPTQ/AutoGPTQ - cd AutoGPTQ - pip install . - - # Download the quant_autogptq script - wget https://gist.githubusercontent.com/TheBloke/b47c50a70dd4fe653f64a12928286682/raw/ebcee019d90a178ee2e6a8107fdd7602c8f1192a/quant_autogptq.py - - # Quantize weights into INT4 and save as safetensors - # Quantized weight with parameter "--act-order" is not supported in TRT-LLM - python quant_autogptq.py ./tmp/llama/7B ./llama-7b-4bit-gs128.safetensors wikitext --bits 4 --group_size 128 --desc_act 0 --damp 0.1 --dtype float16 --seqlen 4096 --num_samples 3 --use_fast - ``` - Then we can convert the saved `./llama-7b-4bit-gs128.safetensors` into TRT-LLM checkpoints by: - ```bash - # Build the LLaMA 7B model using 2-way tensor parallelism and apply INT4 GPTQ quantization. - # Compressed checkpoint safetensors are generated separately from GPTQ. - python convert_checkpoint.py --model_dir /tmp/llama-7b-hf \ - --output_dir ./tllm_checkpoint_2gpu_gptq \ - --dtype float16 \ - --quant_ckpt_path ./llama-7b-4bit-gs128.safetensors \ - --use_weight_only \ - --weight_only_precision int4_gptq \ - --per_group \ - --tp_size 2 - ``` - HF checkpoints generated with AutoGPTQ are also supported through the following conversion script: - - ```bash - # Convert AutoGPTQ HF checkpoints into 2-way tensor parallelism TRT-LLM checkpoint - python convert_checkpoint.py --model_dir ./tmp/Llama-2-7B-GPTQ \ - --output_dir ./tllm_checkpoint_2gpu_gptq \ - --tp_size 2 - ``` - -2. Build TRT-LLM engine: - - ```bash - # Build the LLaMA 7B model using 2-way tensor parallelism and apply INT4 GPTQ quantization. - # Compressed checkpoint safetensors are generated separately from GPTQ. - python convert_checkpoint.py --model_dir /tmp/llama-7b-hf \ - --output_dir ./tllm_checkpoint_2gpu_gptq \ - --dtype float16 \ - --quant_ckpt_path ./llama-7b-4bit-gs128.safetensors \ - --use_weight_only \ - --weight_only_precision int4_gptq \ - --per_group \ - --tp_size 2 - - trtllm-build --checkpoint_dir ./tllm_checkpoint_2gpu_gptq \ - --output_dir ./tmp/llama/7B/trt_engines/int4_GPTQ/2-gpu/ \ - --gemm_plugin auto - ``` - -### w4aINT8 quantization (QServe) - -TensorRT LLM integrates the quantized GEMM from [QServe](https://arxiv.org/abs/2405.04532), which employs 4-bit quantization for weights and 8-bit quantization for activations. This technique offers versatile performance benefits across different scenarios. When the GEMM's m dimension is small, as in small batch-size decoding, it achieves performance comparable to w4a16 by reducing the memory bandwidth required for weight access. Conversely, for larger m dimensions, such as during prefilling or large batch-size decoding, it matches the performance of w8a8 by leveraging INT8 Tensor Cores. - -Please follow the steps to run the model using QServe w4aINT8: - -1. Weight quantization: - - Currently we rely on the 3rd-party repo [deepcompressor](https://github.com/mit-han-lab/deepcompressor) to prepare the fake-quantized checkpoint. Follow the [instructions](https://github.com/mit-han-lab/deepcompressor/blob/main/examples/llm/README.md#usage) to quantize the model. Please use the `configs/qoq-g128.yaml` for per-group quantization, and `configs/qoq-gchn.yaml` for the per-channel quantization. Do not forget to add the flag `--save-model path/to/deepcompressor/ckpt` so that the quantized model is dumped to the disk. - - After quantization, the weights in the original Hugging Face checkpoint (assume under `path/to/huggingface/ckpt/`) will be quantized and the following files are obtained under your `path/to/deepcompressor/ckpt`: - - - `model.pt`: fake-quantized fp16 weights. - - `scale.pt`: quantization scales and zeros. - -2. Checkpoint conversion: - - Convert the DeepCompressor checkpoint into TensorRT LLM checkpoint, potentially with tensor parallelism: - - ```bash - export TRTLLM_DISABLE_UNIFIED_CONVERTER=1 # The current checkpoint conversion code requires legacy path - python convert_checkpoint.py --model_dir path/to/huggingface/ckpt/ \ - --output_dir path/to/trtllm/ckpt/ \ - --dtype float16 \ - --quant_ckpt_path path/to/deepcompressor/ckpt \ - --use_qserve \ - --per_group \ # Add this option if using per-group quantization - --tp_size 2 - ``` - -3. Build engine: - - ```bash - trtllm-build --checkpoint_dir path/to/trtllm/ckpt/ \ - --output_dir path/to/trtllm/engine \ - --gemm_plugin auto - ``` - -### NVFP4 quantization - -TRTLLM supports NVFP4 precision with blocksize=16 for both activations and GEMM weights. - -Please follow the steps to run the model using: - -1. Weight quantization and activation calibration using modelopt: - - ```bash - python example/quantization/quantize.py --model_dir path/to/huggingface/ckpt/ \ - --output_dir path/to/trtllm/ckpt/ \ - --dtype float16 \ - --qformat nvfp4 \ - --kv_cache_dtype fp8 \ - --tp_size 1 - ``` - -2. Build engine: - - ```bash - trtllm-build --checkpoint_dir path/to/trtllm/ckpt/ \ - --output_dir path/to/trtllm/engine - - # with FP8 paged context FMHA for better performance - trtllm-build --checkpoint_dir path/to/trtllm/ckpt/ \ - --output_dir path/to/trtllm/engine \ - --use_paged_context_fmha enable \ - --use_fp8_context_fmha enable - ``` - -### Run - -To run a TensorRT LLM LLaMA model using the engines generated by `trtllm-build` - -```bash -# With fp16 inference -python3 ../../../run.py --max_output_len=50 \ - --tokenizer_dir ./tmp/llama/7B/ \ - --engine_dir=./tmp/llama/7B/trt_engines/fp16/1-gpu/ - -# With bf16 inference -python3 ../../../run.py --max_output_len=50 \ - --tokenizer_dir ./tmp/llama/7B/ \ - --engine_dir=./tmp/llama/7B/trt_engines/bf16/1-gpu/ -``` - -### Multi-GPU multi-node (MGMN) support - -In MGMN case, you can still convert and build engines on a single node and then run the model on a multi-node environment, such as [Slurm](https://slurm.schedmd.com/documentation.html). - -For example, to build LLaMA 70B for 2 nodes with 8 GPUs per node, we can use 8-way tensor parallelism and 2-way pipeline parallelism: - -```bash -python convert_checkpoint.py --model_dir ./tmp/llama/70B/hf/ \ - --output_dir ./tllm_checkpoint_16gpu_tp8_pp2 \ - --dtype float16 \ - --tp_size 8 \ - --pp_size 2 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_16gpu_tp8_pp2 \ - --output_dir ./tmp/llama/70B/trt_engines/fp16/16-gpu/ \ - --workers 8 \ - --gemm_plugin auto -``` - -Note that `–-workers` is still set to 8 to build all engines within a single node. - -To run the LLaMA 70B model on 2 nodes via Slurm, you need to prepare a Slurm script to submit the task, the script contains the following lines: - -```bash -#SBATCH -N 2 -#SBATCH --ntasks-per-node=8 -#SBATCH -p -# more sbatch options here... - -srun --container-image= \ - --mpi=pmix \ - ... \ # more srun options here - python3 ../../../run.py --max_output_len=50 \ - --tokenizer_dir ./tmp/llama/70B/hf/ \ - --engine_dir=./tmp/llama/70B/trt_engines/fp16/16-gpu/ -``` - -Finally, you can submit the task with `sbatch .sh`. - -Considering the Slurm or other cluster management systems may be highly customized and the task-submit command may be variant, the forementioned example is for reference only. The key point is to submit the Python script with the MPI runtime, and TensorRT LLM will take care of the rest. - -### Summarization using the LLaMA model - -```bash -# Run summarization using the LLaMA 7B model in FP16. -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_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_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_dir ./tmp/llama/30B/ \ - --data_type fp16 \ - --engine_dir ./tmp/llama/30B/trt_engines/fp16/2-gpu/ -``` - -#### 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_attention_window_size` 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. - -```bash -# Build Mistral 7B with max input length 32256 -python convert_checkpoint.py --model_dir ./mistral-7b-v0.1 \ - --output_dir ./tllm_checkpoint_1gpu_mistral \ - --dtype float16 -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_mistral \ - --output_dir ./tmp/mistral/7B/trt_engines/fp16/1-gpu/ \ - --gemm_plugin auto \ - --max_input_len 32256 - -# Run Mistral 7B fp16 inference with sliding window/cache size 4096 -python ../../../run.py --max_output_len=50 \ - --tokenizer_dir ./mistral-7b-v0.1 \ - --engine_dir=./tmp/mistral/7B/trt_engines/fp16/1-gpu/ \ - --max_attention_window_size=4096 -``` - -Note that if you are comparing TRT-LLM with Huggingface, -you should install `transformers` with version >= 4.34.1 in order to have Mistral model supported. -And upgrade `flash-attn` package by `pip install --upgrade flash-attn` or you may see wrong results generated by the huggingface implementation. - -#### Mistral Nemo -[Mistral Nemo](https://mistral.ai/news/mistral-nemo/) is compatible with LLaMA interface and can be built and run using the same instructions. -Please upgrade the [transformers](https://pypi.org/project/transformers/) to 4.43.0.dev0 or higher release for running this model. - -```bash -# Build Mistral Nemo with max input length 10240 -python convert_checkpoint.py --model_dir ./Mistral-Nemo-Instruct-2407 \ - --output_dir ./tllm_checkpoint_1gpu_mistral_nemo \ - --dtype bfloat16 \ - --smoothquant 0.5 \ - --per_channel \ - --per_token - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_mistral_nemo \ - --output_dir ./tmp/mistral_nemo/trt_engines/bf16/1-gpu/ \ - --gemm_plugin bfloat16 \ - --max_input_len 10240 - -# Run summarization using the Mistral Nemo model quantized to INT8. -python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./Mistral-Nemo-Instruct-2407 \ - --data_type bf16 \ - --engine_dir ./tmp/mistral_nemo/trt_engines/bf16/1-gpu// -``` - -## Running CodeLlama -Those examples can be used to build and run the CodeLlama models. All 7b, 13b, and 34b sizes and variants are supported. - -There are a couple of differences in CodeLlama in comparison to LLaMA v1/v2 models: rotary_base (`theta=1000000.0f`) and vocabulary size (`32016` (1)). - -_(1): Only applicable to 7b and 13b model sizes_. 34b model variants use `32000`. - -### Build -Use the following command to build `CodeLlama-7b-Instruct`: -```bash -python convert_checkpoint.py --model_dir /tmp/CodeLlama-7b-Instruct-hf \ - --output_dir ./tllm_checkpoint_1gpu_codellama \ - --dtype float16 - - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_codellama \ - --output_dir ./tmp/codellama/trt_engines/fp16/1-gpu/ \ - --gemm_plugin auto -``` -The example below uses the NVIDIA ModelOpt (AlgorithMic Model Optimization) toolkit for the model quantization process. -First make sure Modelopt toolkit is installed (see [examples/quantization/README.md](/examples/quantization/README.md#preparation)) - -```bash -# Quantize HF CodeLlama 7B into FP8 and export trtllm checkpoint -python ../../../quantization/quantize.py --model_dir /tmp/CodeLlama-7b-Instruct-hf \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir ./tllm_checkpoint_1gpu_fp8 \ - --calib_size 512 - -# Build trtllm engines from the trtllm checkpoint -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_fp8 \ - --output_dir ./engine_outputs \ - --gemm_plugin auto -``` - -Use the following command to build `CodeLlama-34b-Instruct` for 4 GPUs (TP=4): -```bash -python convert_checkpoint.py --model_dir /tmp/CodeLlama-34b-Instruct-hf \ - --output_dir ./tllm_checkpoint_4gpu_codellama \ - --dtype float16 \ - --tp_size 4 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_4gpu_codellama \ - --output_dir ./tmp/codellama/trt_engines/fp16/4-gpu/ \ - --gemm_plugin auto -``` - -NOTE: CodeLlama uses the `max_position_embeddings` of 16K. -To build the engine for running similarly long input/output, you need to specify that during build. - -Use `--max_input_len` (which defaults to `1024`) and `--max_seq_len` (which by default is deduced from `max_position_embeddings`) according to your use case, e.g.: -```bash -python convert_checkpoint.py --model_dir /tmp/CodeLlama-34b-Instruct-hf \ - --output_dir ./tllm_checkpoint_4gpu_codellama \ - --dtype float16 \ - --tp_size 8 \ - --use_parallel_embedding - -trtllm-build --checkpoint_dir ./tllm_checkpoint_4gpu_codellama \ - --output_dir ./tmp/codellama/trt_engines/fp16/4-gpu/ \ - --gemm_plugin auto \ - --max_input_len 15360 \ - --max_seq_len 16384 \ - --max_batch_size 4 -``` - -### Run -Use the following command to run the 7b engine from above: -```bash -python ../run.py --max_output_len=40 --tokenizer_dir . --engine_dir codellama_7b --input_text "In Bash, how do I list all text files?" -``` -Use the following command to run the 34b engine with long input/output from above: -```bash -mpirun -n 8 --allow-run-as-root \ - python ../../../run.py --max_output_len=160 --tokenizer_dir ./CodeLlama-34b-Instruct \ - --engine_dir codellama_34b --input_text "In python, write a function for binary searching an element in an integer array." -``` - -## Run models with LoRA - -Download the base model and lora model from HF: - -```bash -git-lfs clone https://huggingface.co/meta-llama/Llama-2-13b-hf -git-lfs clone https://huggingface.co/hfl/chinese-llama-2-lora-13b -``` - -Build engine, setting `--lora_plugin` and `--lora_dir`. If lora has separate lm_head and embedding, they will replace lm_head and embedding of base model. - -```bash -python convert_checkpoint.py --model_dir Llama-2-13b-hf \ - --output_dir ./tllm_checkpoint_2gpu \ - --dtype float16 \ - --tp_size 2 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_2gpu \ - --output_dir /tmp/new_lora_13b/trt_engines/fp16/2-gpu/ \ - --gemm_plugin auto \ - --lora_plugin auto \ - --max_batch_size 1 \ - --max_input_len 512 \ - --max_seq_len 562 \ - --lora_dir chinese-llama-2-lora-13b -``` - -Run inference. Remember to use lora tokenizer because lora model has larger vocab size. - -```bash -mpirun -n 2 python ../../../run.py --engine_dir "/tmp/new_lora_13b/trt_engines/fp16/2-gpu/" \ - --max_output_len 50 \ - --tokenizer_dir "chinese-llama-2-lora-13b/" \ - --input_text "今天天气很好，我到公园的时候，" \ - --lora_task_uids 0 \ - --no_add_special_tokens \ - --use_py_session - - Input: "今天天气很好，我到公园的时候，" -Output: "发现公园里到处都是人，有的在跑步，有的在打羽毛球，还有的在跳绳，我和妈妈一起在公园里散步，我和妈妈在公园里散步的时候，看见了一位老爷爷在打羽毛球" -``` - -Users who want to skip LoRA module may pass uid -1 with `--lora_task_uids -1`. -In that case, the model will not run the LoRA module and the results will be -different. Since the LoRA tokenizer, embedding and LM head are still used, -the results will also be different with vanilla LLaMA and significantly degrade compared with `--lora_task_uids 0`. - -```bash -mpirun -n 2 python ../../../run.py --engine_dir "/tmp/new_lora_13b/trt_engines/fp16/2-gpu/" \ - --max_output_len 50 \ - --tokenizer_dir "chinese-llama-2-lora-13b/" \ - --input_text "今天天气很好，我到公园的时候，" \ - --lora_task_uids -1 \ - --no_add_special_tokens \ - --use_py_session - - Input: "今天天气很好，我到公园的时候，" -Output: "看见好多人们都看书，看书书看书书，看书书看书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书" -``` - -### Run LLaMa with several lora checkpoints - -In this section, we show how to run a model with multiple LoRA modules at the same time. Note that if one of the LoRA module has a -fine-tuned embedding table or logit GEMM, users should guarantee that all the instances of the model can use the same fine-tuned -embedding table or logit GEMM. -Here, we use two LoRA checkpoints as examples. These two LoRA checkponits add LoRA modules to `q_proj` and `v_proj`. Because we only -support adding lora modules on `q`, `k` and `v` at the same time, we need to add `--lora_target_modules "attn_q" "attn_k" "attn_v"`. -In this case, we assign null pointers for the `k` LoRA module in TensorRT LLM and skip the computation at runtime. - -As the rank of the LoRA modules of both checkpoints is 8, we can set `--max_lora_rank 8` to reduce the memory requirement for the LoRA plugin. - -In this example, we use a LoRA checkpoint fine-tuned on the Chinese dataset `luotuo-lora-7b-0.1` and a LoRA checkpoint fine-tuned on -the Japanese dataset `Japanese-Alpaca-LoRA-7b-v0`. For the `lora_manager` to load several checkpoints, we pass several directories -of LoRA checkpoints at the same time: `--lora_dir "luotuo-lora-7b-0.1/" "Japanese-Alpaca-LoRA-7b-v0/"`. -Then, `lora_manager` will assign `lora_task_uids` to these checkpoints. `lora_task_uids -1` is a predefined value, which corresponds to -the base model. If we pass `lora_task_uids 0 1`, this means we want to use the first LoRA checkpoint on first sentence and use the second LoRA checkpoint on the second sentence. - -To verify the correctness, we pass the same Chinese input `美国的首都在哪里? \n答案:` three times as well as the same Japanese input `アメリカ合衆国の首都はどこですか? \n答え:` three times (both inputs mean `Where is the capital of America? \nAnswer`). We run on base model, `luotuo-lora-7b-0.1` and `Japanese-Alpaca-LoRA-7b-v0/`. - -```bash -git-lfs clone https://huggingface.co/qychen/luotuo-lora-7b-0.1 -git-lfs clone https://huggingface.co/kunishou/Japanese-Alpaca-LoRA-7b-v0 -BASE_LLAMA_MODEL=llama-7b-hf/ - -python convert_checkpoint.py --model_dir ${BASE_LLAMA_MODEL} \ - --output_dir ./tllm_checkpoint_1gpu \ - --dtype float16 -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu \ - --output_dir /tmp/llama_7b_with_lora_qkv/trt_engines/fp16/1-gpu/ \ - --gemm_plugin auto \ - --lora_plugin auto \ - --max_batch_size 8 \ - --max_input_len 512 \ - --max_seq_len 562 \ - --lora_dir "luotuo-lora-7b-0.1/" "Japanese-Alpaca-LoRA-7b-v0/" \ - --max_lora_rank 8 \ - --lora_target_modules attn_q attn_k attn_v - -python ../../../run.py --engine_dir "/tmp/llama_7b_with_lora_qkv/trt_engines/fp16/1-gpu/" \ - --max_output_len 10 \ - --tokenizer_dir ${BASE_LLAMA_MODEL} \ - --input_text "美国的首都在哪里? \n答案:" "美国的首都在哪里? \n答案:" "美国的首都在哪里? \n答案:" "アメリカ合衆国の首都はどこですか? \n答え:" "アメリカ合衆国の首都はどこですか? \n答え:" "アメリカ合衆国の首都はどこですか? \n答え:" \ - --lora_task_uids -1 0 1 -1 0 1 \ - --use_py_session --top_p 0.5 --top_k 0 -``` - -The results would be like - -```bash -Input [Text 0]: " 美国的首都在哪里? \n答案:" -Output [Text 0 Beam 0]: "Washington, D.C. -What is the" - -Input [Text 1]: " 美国的首都在哪里? \n答案:" -Output [Text 1 Beam 0]: "华盛顿。 -" - -Input [Text 2]: " 美国的首都在哪里? \n答案:" -Output [Text 2 Beam 0]: "Washington D.C.�����" - -Input [Text 3]: " アメリカ合衆国の首都はどこですか? \n答え:" -Output [Text 3 Beam 0]: "Washington, D.C. -Which of" - -Input [Text 4]: " アメリカ合衆国の首都はどこですか? \n答え:" -Output [Text 4 Beam 0]: "华盛顿。 -" - -Input [Text 5]: " アメリカ合衆国の首都はどこですか? \n答え:" -Output [Text 5 Beam 0]: "ワシントン D.C." -``` - -We can observe that `luotuo-lora-7b-0.1` produces correct answers on the first sentence and the fifth sentence (in Chinese), `Japanese-Alpaca-LoRA-7b-v0` produces correct answers on the sixth sentence (in Japanese). - -### Run FP8 Mistral v0.1 with FP16 lora checkpoint - -In this section, we use Mistral v0.1 as an example show how to run an FP8 base model with FP16 LoRA module. - -* download the base model and lora model from HF - -```bash -git-lfs clone https://huggingface.co/davidkim205/komt-mistral-7b-v1 -git-lfs clone https://huggingface.co/davidkim205/komt-mistral-7b-v1-lora -``` - -* Quantize the Mistral v0.1 model to fp8 from HF -```bash -BASE_MISTRAL_MODEL=komt-mistral-7b-v1/ -python ../../../quantization/quantize.py --model_dir ${BASE_MISTRAL_MODEL} \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir ./tllm_checkpoint_1gpu_fp8 \ - --calib_size 512 -``` - -* Build engine and run inference with sliding window/cache size 4096. -```bash -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_fp8 \ - --output_dir /tmp/mistral_komt_lora/7B/trt_engines/fp8/1-gpu/ \ - --gemm_plugin auto \ - --lora_plugin auto \ - --max_batch_size 8 \ - --max_input_len 32256 \ - --max_seq_len 33280 \ - --lora_dir ./komt-mistral-7b-v1-lora - -python ../../../run.py --max_output_len=1024 \ - --tokenizer_dir ./komt-mistral-7b-v1 \ - --engine_dir=/tmp/mistral_komt_lora/7B/trt_engines/fp8/1-gpu/ \ - --input_text "[INST]오늘은 날씨가 아주 좋다 내가 공원에 갔을 때 [/INST]" \ - --max_attention_window_size=4096 \ - --lora_task_uids 0 \ - --use_py_session \ - --temperature 0.8 \ - --top_p 0.8 \ - --top_k 100 -``` - -The results would be like - -```bash -Input [Text 0]: " [INST]오늘은 날씨가 아주 좋다 내가 공원에 갔을 때 [/INST]" -Output [Text 0 Beam 0]: "날씨가 아주 좋은 날에 공원에 갔을 때는 산책이나 운동을 즐기는 것이 좋습니다. 공원에서 걷거나 조깅을 하면서 신선한 공기를 마시고 자연 속에서 휴식을 취할 수 있습니다. 또한, 공원에서 가족이나 친구와 함께 피크닉을 즐기거나 야외 스포츠를 즐길 수도 있습니다. 날씨가 좋을 때 공원에 가는 것은 건강과 웰빙에 좋은 방법입니다." -``` - - -### Run INT4-AWQ LLaMa with several FP16 lora checkpoints - -TensorRT LLM can also support Quantized base model + FP16/BF16 LoRA. We can first quantize the base model and build engine with the quantized checkpoint and different LoRA adapters. In this section, we show how to run an INT4-AWQ llama model with multiple FP16 LoRA modules. - -* Quantize the llama model to INT4-AWQ from HF -```bash -BASE_LLAMA_MODEL=llama-7b-hf/ -python ../../../quantization/quantize.py --model_dir ${BASE_LLAMA_MODEL} \ - --output_dir ./tllm_checkpoint_1gpu_awq \ - --dtype float16 \ - --qformat int4_awq \ - --awq_block_size 128 \ - --calib_size 32 -``` - -* Download the lora model, build engine, and run inference. -```bash -git-lfs clone https://huggingface.co/qychen/luotuo-lora-7b-0.1 -git-lfs clone https://huggingface.co/kunishou/Japanese-Alpaca-LoRA-7b-v0 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_awq \ - --output_dir /tmp/llama_7b_with_lora_qkv/trt_engines/int4_AWQ/1-gpu/ \ - --gemm_plugin auto \ - --lora_plugin auto \ - --max_batch_size 8 \ - --max_input_len 512 \ - --max_seq_len 562 \ - --lora_dir "luotuo-lora-7b-0.1/" "Japanese-Alpaca-LoRA-7b-v0/" \ - --max_lora_rank 8 \ - --lora_target_modules attn_q attn_k attn_v - -python ../../../run.py --engine_dir "/tmp/llama_7b_with_lora_qkv/trt_engines/int4_AWQ/1-gpu/" \ - --max_output_len 10 \ - --tokenizer_dir ${BASE_LLAMA_MODEL} \ - --input_text "美国的首都在哪里? \n答案:" "美国的首都在哪里? \n答案:" "美国的首都在哪里? \n答案:" "アメリカ合衆国の首都はどこですか? \n答え:" "アメリカ合衆国の首都はどこですか? \n答え:" "アメリカ合衆国の首都はどこですか? \n答え:" \ - --lora_task_uids -1 0 1 -1 0 1 \ - --use_py_session --top_p 0.5 --top_k 0 -``` - -The results would be like - -```bash -Input [Text 0]: " 美国的首都在哪里? \n答案:" -Output [Text 0 Beam 0]: "Washington, D.C. -What is the" - -Input [Text 1]: " 美国的首都在哪里? \n答案:" -Output [Text 1 Beam 0]: "华盛顿。 -" - -Input [Text 2]: " 美国的首都在哪里? \n答案:" -Output [Text 2 Beam 0]: "洛爱睿" - -Input [Text 3]: " アメリカ合衆国の首都はどこですか? \n答え:" -Output [Text 3 Beam 0]: "Washington, D.C. -Copyright " - -Input [Text 4]: " アメリカ合衆国の首都はどこですか? \n答え:" -Output [Text 4 Beam 0]: "华盛顿。 -" - -Input [Text 5]: " アメリカ合衆国の首都はどこですか? \n答え:" -Output [Text 5 Beam 0]: "ワシントン、D.C" -``` - -## Run LLaMa with StreamingLLM - -* Build engine. Set `--streamingllm enable` to enable StreamingLLM. - -```bash -# Build the LLaMA 7B model with StreamingLLM feature using a single GPU and FP16. -python convert_checkpoint.py --model_dir ./tmp/llama/7B/ \ - --output_dir ./tllm_checkpoint_1gpu_streamlingllm \ - --dtype float16 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_streamlingllm \ - --output_dir ./tmp/llama/7B/trt_engines/fp16_StreamingLLM/1-gpu/ \ - --gemm_plugin auto \ - --streamingllm enable - -``` - -* Run inference. Use `--sink_token_length` to set the number of sink tokens, and use `--max_attention_window_size` to set the `sliding_window` value. - -```bash -# Run LLaMA 7B fp16 inference with sliding window/cache size 2048 and sink token length 4. -python3 ../../../run.py --max_output_len=50 \ - --tokenizer_dir ./tmp/llama/7B/ \ - --engine_dir=./tmp/llama/7B/trt_engines/fp16_StreamingLLM/1-gpu/ \ - --max_attention_window_size=2048 \ - --sink_token_length=4 -``` - -Note that the sink tokens is included in the sliding attention tokens, and there are at most `max_attention_window_size` tokens are stored in the KV cache. - -## Run LLaMA-3.1 405B Model - -Currently, TensorRT LLM supports Meta checkpoint and Huggingface checkpoint for LLaMA-3.1. In this section, we demonstrate how to run the LLaMA-3.1 405B model via TensorRT-LLM. Here, we assume users have downloaded the checkpoints and placed them at `llama_3.1_405B_meta_model/` (Meta BF16 checkpoint), `llama_3.1_405B_HF_model/` (HF BF16 checkpoint) and `llama_3.1_405B_HF_FP8_model/` (HF FP8 checkpoint). Before converting the checkpoints to TensorRT LLM unified checkpoints, **please check that `{"rope_scaling": {"rope_type": "llama3"}}` is set in the configuration file**. With this flag, TensorRT LLM will enable the rope scaling of LLaMA-3.1. If not, please add it to the config file. - -Users can run the LLaMA-3.1 model with higher precision (bf16/fp16) or fp8. Here, to prevent accuracy drop, we perform per-channel per-token fp8 quantization (leveraged from https://github.com/pytorch/FBGEMM) on MLP layers, keeping other layers at higher precision. Note that per-channel per-token fp8 quantization is only supported on Huggingface checkpoint now. We will support it on Meta checkpoint soon. Note that this feature only supports SM90. - -### Convert Checkpoint to TensorRT LLM Unified Checkpoint - -To use the fp8 quantization, please add the `--use_fp8_rowwise` flag during the checkpoint conversion. In this demonstration, we convert the Meta checkpoint to bfloat16 with TP8-PP2 and the HF checkpoint to FP8 with TP8. - -Note: You may need to update your transformers installation via `pip install --upgrade transformers`. -Note: For 405B HF model cloned before 09 Aug 2024, there are duplicated kv head weights (The `num_key_value_heads` in `config.json` is 16). Users could use `--remove_duplicated_kv_heads` to remove them. The new checkpoint without duplicated kv heads is uploaded on 09 Aug 2024 and the `num_key_value_heads` is 8 now. For the new checkpoint, adding the flag `--remove_duplicated_kv_heads` would lead to error. - -```bash -# Run BF16 model by BF16 -python examples/models/core/llama/convert_checkpoint.py --meta_ckpt_dir llama_3.1_405B_meta_model/ \ - --output_dir llama_3.1_405B_meta_model/trt_ckpts/tp8-pp2/ \ - --dtype bfloat16 \ - --tp_size 8 \ - --pp_size 2 \ - --load_by_shard \ - --workers 2 - -# Run BF16 model by FP8 -python examples/models/core/llama/convert_checkpoint.py --model_dir llama_3.1_405B_HF_model/ \ - --output_dir llama_3.1_405B_HF_model/trt_ckpts/tp8-pp1/ \ - --dtype bfloat16 \ - --use_fp8_rowwise \ - --tp_size 8 \ - --pp_size 1 \ - --load_by_shard \ - --workers 8 \ - -# Run FP8 model by FP8 -# The source HF model is in FP8 format, so --use_fp8_rowwise is enabled automatically -# Optionally enable --use_meta_fp8_rowwise_recipe to strictly follow the original Meta's LLaMA 3.1 recipe: -# (1) Skip quantization for the first and last Transformer layers -# (2) Skip quantization for the Attention layers -python examples/models/core/llama/convert_checkpoint.py --model_dir llama_3.1_405B_HF_FP8_model/ \ - --output_dir llama_3.1_405B_HF_FP8_model/trt_ckpts/tp8-pp1/ \ - --dtype bfloat16 \ - --tp_size 8 \ - --pp_size 1 \ - --use_meta_fp8_rowwise_recipe \ - --load_by_shard \ - --workers 8 -``` - -### Build Engine - -```bash -trtllm-build --checkpoint_dir llama_3.1_405B_meta_model/trt_ckpts/tp8-pp2/ \ - --output_dir llama_3.1_405B_meta_model/trt_engines/tp8-pp2/ \ - --max_num_tokens 4096 \ - --max_input_len 255000 \ - --max_seq_len 256000 \ - --use_paged_context_fmha enable \ - --workers 8 - -trtllm-build --checkpoint_dir llama_3.1_405B_HF_model/trt_ckpts/tp8-pp1/ \ - --output_dir llama_3.1_405B_HF_model/trt_engines/tp8-pp1/ \ - --max_num_tokens 4096 \ - --max_input_len 64000 \ - --max_seq_len 65000 \ - --use_paged_context_fmha enable \ - --workers 8 - -trtllm-build --checkpoint_dir llama_3.1_405B_HF_FP8_model/trt_ckpts/tp8-pp1/ \ - --output_dir llama_3.1_405B_HF_FP8_model/trt_engines/tp8-pp1/ \ - --max_num_tokens 4096 \ - --max_input_len 64000 \ - --max_seq_len 65000 \ - --use_paged_context_fmha enable \ - --workers 8 -``` - -### Run Inference - -To run inference on the 405B model, we often need to use multi-node to accommodate the entire model. Here, we use slurm to launch the job on multiple nodes. - -Notes: -* For convenience, we use the Huggingface tokenizer for tokenization. - -The following script shows how to run evaluation on long context: - -```bash -# Long context test for 128k -python ./examples/infinitebench/construct_synthetic_dataset.py --test_case build_passkey --test_level 4; mkdir -p 128k_context; mv passkey.jsonl 128k_context; - -srun --mpi pmi2 -N 2 -n 16 --ntasks-per-node 8 --container-image \ ---container-mounts \ ---container-name llama-3.1-405b \ ---container-workdir \ -bash -c 'python ./examples/eval_long_context.py --task passkey \ - --engine_dir llama_3.1_405B_meta_model/trt_engines/tp8-pp2/ \ - --tokenizer_dir llama_3.1_405B_HF_model/ \ - --stop_idx 6 \ - --max_input_length 255000 \ - --enable_chunked_context \ - --kv_cache_free_gpu_memory_fraction 0.999 \ - --max_tokens_in_paged_kv_cache 256064 \ - --data_dir 128k_context \ - --output_dir 128k_context_tp8' - -# output would be like -# 6it [00:00, 15354.38it/s] -# [07/22/2024-19:09:54] [TRT-LLM] [I] Evaluation takes: 372.16685914993286 sec. -# [07/22/2024-19:09:54] [TRT-LLM] [I] accuracy of 6 examples: 1.0 - -# Long context test for 64k -python ./examples/infinitebench/construct_synthetic_dataset.py --test_case build_passkey --test_level 3; mkdir -p 64k_context; mv passkey.jsonl 64k_context; - -srun --mpi pmi2 -N 1 -n 8 --ntasks-per-node 8 --container-image \ ---container-mounts \ ---container-name llama-3.1-405b \ ---container-workdir \ -bash -c 'python ./examples/eval_long_context.py --task passkey \ - --engine_dir llama_3.1_405B_HF_model/trt_engines/tp8-pp1/ \ - --tokenizer_dir llama_3.1_405B_HF_model/ \ - --stop_idx 6 \ - --max_input_length 64000 \ - --enable_chunked_context \ - --kv_cache_free_gpu_memory_fraction 0.999 \ - --max_tokens_in_paged_kv_cache 65064 \ - --data_dir 64k_context \ - --output_dir 64k_context_tp8' - -# Long context test for 64k -srun --mpi pmi2 -N 1 -n 8 --ntasks-per-node 8 --container-image \ ---container-mounts \ ---container-name llama-3.1-405b \ ---container-workdir \ -bash -c 'python ./examples/eval_long_context.py --task passkey \ - --engine_dir llama_3.1_405B_HF_FP8_model/trt_engines/tp8-pp1/ \ - --tokenizer_dir llama_3.1_405B_HF_FP8_model/ \ - --stop_idx 6 \ - --max_input_length 64000 \ - --enable_chunked_context \ - --kv_cache_free_gpu_memory_fraction 0.999 \ - --max_tokens_in_paged_kv_cache 65064 \ - --data_dir 64k_context \ - --output_dir 64k_context_tp8' -``` - -The following script shows how to run evaluation on MMLU tasks: - -```bash -srun --mpi pmi2 -N 2 -n 16 --ntasks-per-node 8 --container-image \ ---container-mounts \ ---container-name llama-3.1-405b \ ---container-workdir \ -bash -c 'python ./examples/mmlu.py --test_trt_llm \ - --engine_dir llama_3.1_405B_meta_model/trt_engines/tp8-pp2/ \ - --tokenizer_dir llama_3.1_405B_HF_model/ \ - --enable_chunked_context \ - --kv_cache_free_gpu_memory_fraction 0.999 \ - --max_tokens_in_paged_kv_cache 256064' - -srun --mpi pmi2 -N 1 -n 8 --ntasks-per-node 8 --container-image \ ---container-mounts \ ---container-name llama-3.1-405b \ ---container-workdir \ -bash -c 'python ./examples/mmlu.py --test_trt_llm \ - --engine_dir llama_3.1_405B_HF_model/trt_engines/tp8-pp1/ \ - --tokenizer_dir llama_3.1_405B_HF_model/ \ - --enable_chunked_context \ - --kv_cache_free_gpu_memory_fraction 0.999 \ - --max_tokens_in_paged_kv_cache 65064' - -srun --mpi pmi2 -N 1 -n 8 --ntasks-per-node 8 --container-image \ ---container-mounts \ ---container-name llama-3.1-405b \ ---container-workdir \ -bash -c 'python ./examples/mmlu.py --test_trt_llm \ - --engine_dir llama_3.1_405B_HF_FP8_model/trt_engines/tp8-pp1/ \ - --tokenizer_dir llama_3.1_405B_HF_FP8_model/ \ - --enable_chunked_context \ - --kv_cache_free_gpu_memory_fraction 0.999 \ - --max_tokens_in_paged_kv_cache 65064' -``` - -## Run LLaMa-3.3 70B Model on PyTorch Backend -This section provides the steps to run LLaMa-3.3 70B model FP8 precision on PyTorch backend by launching TensorRT LLM server and run performance benchmarks. - -### Prepare TensorRT LLM extra configs -```bash -cat >./config.yml <