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 <