Add EXAONE-Deep (NVIDIA#3054)

Signed-off-by: yechank <[email protected]>
Co-authored-by: QI JUN <[email protected]>

Author: yechank-nvidia

examples/exaone/README.md

@@ -7,7 +7,9 @@ See the LLaMA example [`examples/llama`](../llama) for details.
 
 - [EXAONE](#exaone)
   - [Support Matrix](#support-matrix)
-  - [Download model checkpoints](#download-model-checkpoints)
+  - [Supported Models](#supported-models)
+    - [EXAONE-3.0](#exaone-30)
+    - [EXAONE-Deep](#exaone-deep)
   - [Usage](#usage)
     - [Convert checkpoint and build TensorRT engine(s)](#convert-checkpoint-and-build-tensorrt-engines)
     - [FP8 Post-Training Quantization](#fp8-post-training-quantization)
@@ -25,12 +27,23 @@ See the LLaMA example [`examples/llama`](../llama) for details.
   * INT8 SmoothQuant
   * INT4 AWQ & W4A8 AWQ
 
-## Download model checkpoints
+## Supported Models
+### EXAONE-3.0
 
-First, download the HuggingFace FP32 checkpoints of EXAONE model.
+Download the HuggingFace FP32 checkpoints of EXAONE-3.0 model. We support EXAONE-3.0 families but here, we only use the `EXAONE-3.0-7.8B-Instruct` model for the example.
 
 ```bash
-git clone https://huggingface.co/LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct hf_models/exaone
+export HF_MODEL_DIR=hf_models/exaone
+git clone https://huggingface.co/LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct $HF_MODEL_DIR
+```
+
+### EXAONE-Deep
+
+Download the HuggingFace BF16 checkpoints of EXAONE-Deep model. Here, we only use the `EXAONE-Deep-2.4B` model for the example. We can use the same procedure as EXAONE-3.0 to convert the weights and build the TensorRT engine.
+
+```bash
+export HF_MODEL_DIR=hf_models/exaone_deep
+git clone https://huggingface.co/LGAI-EXAONE/EXAONE-Deep-2.4B $HF_MODEL_DIR
 ```
 
 ## Usage
@@ -43,7 +56,7 @@ The next section describe how to convert the weights from the [HuggingFace (HF)
 
 # Build the EXAONE model using a single GPU and FP16.
 python ../llama/convert_checkpoint.py \
-    --model_dir hf_models/exaone \
+    --model_dir $HF_MODEL_DIR \
     --output_dir trt_models/exaone/fp16/1-gpu \
     --dtype float16
 
@@ -54,7 +67,7 @@ trtllm-build \
 
 # Build the EXAONE model using a single GPU and and apply INT8 weight-only quantization.
 python ../llama/convert_checkpoint.py \
-    --model_dir hf_models/exaone \
+    --model_dir $HF_MODEL_DIR \
     --output_dir trt_models/exaone/int8_wq/1-gpu \
     --use_weight_only \
     --weight_only_precision int8 \
@@ -67,7 +80,7 @@ trtllm-build \
 
 # Build the EXAONE model using a single GPU and and apply INT4 weight-only quantization.
 python ../llama/convert_checkpoint.py \
-    --model_dir hf_models/exaone \
+    --model_dir $HF_MODEL_DIR \
     --output_dir trt_models/exaone/int4_wq/1-gpu \
     --use_weight_only \
     --weight_only_precision int4 \
@@ -78,9 +91,9 @@ trtllm-build \
     --output_dir trt_engines/exaone/int4_wq/1-gpu \
     --gemm_plugin auto
 
-# Build the EXAONE model using using 2-way tensor parallelism and FP16.
+# Build the EXAONE model using 2-way tensor parallelism and FP16.
 python ../llama/convert_checkpoint.py \
-    --model_dir hf_models/exaone \
+    --model_dir $HF_MODEL_DIR \
     --output_dir trt_models/exaone/fp16/2-gpu \
     --tp_size 2 \
     --dtype float16
@@ -101,7 +114,7 @@ First make sure Modelopt toolkit is installed (see [examples/quantization/README
 ```bash
 # Build the EXAONE model using a single GPU and and apply FP8 quantization.
 python ../quantization/quantize.py \
-    --model_dir hf_models/exaone \
+    --model_dir $HF_MODEL_DIR \
     --dtype float16 \
     --qformat fp8 \
     --kv_cache_dtype fp8 \
@@ -122,7 +135,7 @@ First make sure Modelopt toolkit is installed (see [examples/quantization/README
 ```bash
 # Build the EXAONE model using a single GPU and and apply INT8 SmoothQuant.
 python ../quantization/quantize.py \
-    --model_dir hf_models/exaone \
+    --model_dir $HF_MODEL_DIR \
     --dtype float16 \
     --qformat int8_sq \
     --output_dir trt_models/exaone/int8_sq/1-gpu
@@ -142,7 +155,7 @@ First make sure Modelopt toolkit is installed (see [examples/quantization/README
 ```bash
 # Build the EXAONE model using a single GPU and and apply INT4 AWQ.
 python ../quantization/quantize.py \
-    --model_dir hf_models/exaone \
+    --model_dir $HF_MODEL_DIR \
     --dtype float16 \
     --qformat int4_awq \
     --output_dir trt_models/exaone/int4_awq/1-gpu
@@ -161,7 +174,7 @@ Please make sure your system contains a Hopper GPU before trying the commands be
 ```bash
 # Build the EXAONE model using a single GPU and and apply W4A8 AWQ.
 python ../quantization/quantize.py \
-    --model_dir hf_models/exaone \
+    --model_dir $HF_MODEL_DIR \
     --dtype float16 \
     --qformat w4a8_awq \
     --output_dir trt_models/exaone/w4a8_awq/1-gpu
@@ -180,21 +193,21 @@ Test your engine with the [run.py](../run.py) script:
 python3 ../run.py \
     --input_text "When did the first world war end?" \
     --max_output_len=100 \
-    --tokenizer_dir hf_models/exaone \
+    --tokenizer_dir $HF_MODEL_DIR \
     --engine_dir trt_engines/exaone/fp16/1-gpu
 
 # Run with 2 GPUs
 mpirun -n 2 --allow-run-as-root \
     python3 ../run.py \
     --input_text "When did the first world war end?" \
     --max_output_len=100 \
-    --tokenizer_dir hf_models/exaone \
+    --tokenizer_dir $HF_MODEL_DIR \
     --engine_dir trt_engines/exaone/fp16/2-gpu
 
 python ../summarize.py \
     --test_trt_llm \
     --data_type fp16 \
-    --hf_model_dir hf_models/exaone \
+    --hf_model_dir $HF_MODEL_DIR \
     --engine_dir trt_engines/exaone/fp16/1-gpu
 ```
 

tensorrt_llm/models/llama/model.py

@@ -478,7 +478,7 @@ def from_hugging_face(
                 }
             elif "vila" in model_name:
                 hf_model_dir += "/llm"
-            elif "exaone" in model_name:
+            elif "exaone" in model_name.lower():
                 custom_dict = {
                     "transformer": "transformer",
                     "layers": "h",

tests/integration/defs/.test_durations

@@ -168,8 +168,8 @@
    "examples/test_draft_target_model.py::test_llm_draft_target_model_1gpu[streaming-gpt2-use_cpp_session-use_logits-draft_len_4-float16-bs2]": 222.54111004807055,
    "examples/test_enc_dec.py::test_llm_enc_dec_general[compare_hf-bart-large-cnn-float32-enable_gemm_plugin-enable_attention_plugin-enable_paged_kv_cache-tp:1-pp:1-nb:2-disable_fp8]": 203.55354792065918,
    "examples/test_enc_dec.py::test_llm_enc_dec_general[compare_hf-t5-small-float32-enable_gemm_plugin-enable_attention_plugin-enable_paged_kv_cache-tp:1-pp:1-nb:1-disable_fp8]": 189.6864925120026,
-   "examples/test_exaone.py::test_llm_exaone_1gpu[disable_weight_only-exaone-float16-nb:1]": 473.8068177103996,
-   "examples/test_exaone.py::test_llm_exaone_1gpu[disable_weight_only-exaone-float16-nb:4]": 205.28752172738314,
+   "examples/test_exaone.py::test_llm_exaone_1gpu[disable_weight_only-exaone_3.0_7.8b_instruct-float16-nb:1]": 473.8068177103996,
+   "examples/test_exaone.py::test_llm_exaone_1gpu[disable_weight_only-exaone_3.0_7.8b_instruct-float16-nb:4]": 205.28752172738314,
    "examples/test_multimodal.py::test_llm_multimodal_general[deplot-pp:1-tp:1-float16-bs:1-cpp_e2e:False-nb:1]": 179.15185776166618,
    "examples/test_prompt_lookup.py::test_llm_prompt_lookup_1gpu[no_streaming-gpt2-use_cpp_session-use_tokens-max_matching_ngram_size_2-prompt_lookup_num_tokens_8-float16-bs1]": 233.80333462916315,
    "examples/test_qwen.py::test_llm_qwen_single_gpu_summary[qwen2_0.5b_instruct-enable_paged_kv_cache-enable_remove_input_padding-enable_weight_only-enable_fmha]": 123.65003899484873,

tests/integration/defs/conftest.py

@@ -453,10 +453,13 @@ def llm_exaone_model_root(request) -> str:
     "Get EXAONE model root"
     models_root = llm_models_root()
     assert models_root, "Did you set LLM_MODELS_ROOT?"
-    assert request.param == "exaone", "Is the name of model root is exaone?"
 
-    exaone_model_root = os.path.join(models_root, request.param)
-    assert exists(exaone_model_root), f"{exaone_model_root} does not exist!"
+    exaone_model_root = os.path.join(models_root, "exaone")
+    if hasattr(request, "param"):
+        if request.param == "exaone_3.0_7.8b_instruct":
+            exaone_model_root = os.path.join(models_root, "exaone")
+        elif request.param == "exaone_deep_2.4b":
+            exaone_model_root = os.path.join(models_root, "EXAONE-Deep-2.4B")
 
     return exaone_model_root
 

tests/integration/defs/examples/test_exaone.py

@@ -23,7 +23,9 @@
 @pytest.mark.parametrize("num_beams", [1, 2, 4],
                          ids=lambda num_beams: f'nb:{num_beams}')
 @pytest.mark.parametrize("data_type", ['bfloat16', 'float16'])
-@pytest.mark.parametrize("llm_exaone_model_root", ['exaone'], indirect=True)
+@pytest.mark.parametrize("llm_exaone_model_root",
+                         ['exaone_3.0_7.8b_instruct', 'exaone_deep_2.4b'],
+                         indirect=True)
 @pytest.mark.parametrize("use_weight_only", [True, False],
                          ids=["enable_weight_only", "disable_weight_only"])
 def test_llm_exaone_1gpu(data_type, exaone_example_root, llm_exaone_model_root,
@@ -44,13 +46,11 @@ def test_llm_exaone_1gpu(data_type, exaone_example_root, llm_exaone_model_root,
         data_type=data_type,
         use_weight_only=use_weight_only)
 
-    # TODO: Should we add use_weight_only_groupwise_quant_matmul_plugin?
-
     build_cmd = [
-        "trtllm-build", f"--checkpoint_dir={model_dir}",
-        f"--output_dir={engine_dir}", f"--gpt_attention_plugin={data_type}",
-        f"--gemm_plugin={data_type}", f"--max_beam_width={num_beams}",
-        "--max_batch_size=256"
+        "trtllm-build",
+        f"--checkpoint_dir={model_dir}",
+        f"--output_dir={engine_dir}",
+        f"--max_beam_width={num_beams}",
     ]
     check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env)
 
@@ -80,7 +80,9 @@ def test_llm_exaone_1gpu(data_type, exaone_example_root, llm_exaone_model_root,
 @pytest.mark.parametrize("num_beams", [1],
                          ids=lambda num_beams: f'nb:{num_beams}')
 @pytest.mark.parametrize("data_type", ['float16'])
-@pytest.mark.parametrize("llm_exaone_model_root", ['exaone'], indirect=True)
+@pytest.mark.parametrize("llm_exaone_model_root",
+                         ['exaone_3.0_7.8b_instruct', 'exaone_deep_2.4b'],
+                         indirect=True)
 def test_llm_exaone_2gpu(data_type, exaone_example_root, llm_exaone_model_root,
                          llama_example_root, llm_datasets_root, llm_rouge_root,
                          llm_venv, cmodel_dir, engine_dir, num_beams):
@@ -102,8 +104,7 @@ def test_llm_exaone_2gpu(data_type, exaone_example_root, llm_exaone_model_root,
 
     build_cmd = [
         "trtllm-build", f"--checkpoint_dir={model_dir}",
-        f"--output_dir={engine_dir}", f"--gpt_attention_plugin={data_type}",
-        f"--gemm_plugin={data_type}", f"--max_beam_width={num_beams}"
+        f"--output_dir={engine_dir}", f"--max_beam_width={num_beams}"
     ]
     check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env)
 

tests/integration/test_lists/qa/examples_test_list.txt

@@ -38,10 +38,11 @@ examples/test_enc_dec.py::test_llm_enc_dec_general[compare_hf-t5-small-float32-e
 examples/test_enc_dec.py::test_llm_enc_dec_general[compare_hf-t5-small-float32-enable_gemm_plugin-enable_attention_plugin-enable_paged_kv_cache-tp:2-pp:1-nb:1-enable_fp8]
 examples/test_enc_dec.py::test_llm_enc_dec_general[no_compare_hf-byt5-small-float32-enable_gemm_plugin-enable_attention_plugin-enable_paged_kv_cache-tp:1-pp:1-nb:1-enable_fp8]
 examples/test_enc_dec.py::test_llm_enc_dec_general[no_compare_hf-byt5-small-float32-enable_gemm_plugin-enable_attention_plugin-enable_paged_kv_cache-tp:2-pp:1-nb:1-disable_fp8]
-examples/test_exaone.py::test_llm_exaone_1gpu[disable_weight_only-exaone-float16-nb:1]
-examples/test_exaone.py::test_llm_exaone_1gpu[disable_weight_only-exaone-float16-nb:4]
-examples/test_exaone.py::test_llm_exaone_1gpu[enable_weight_only-exaone-float16-nb:1]
-examples/test_exaone.py::test_llm_exaone_2gpu[exaone-float16-nb:1]
+examples/test_exaone.py::test_llm_exaone_1gpu[disable_weight_only-exaone_3.0_7.8b_instruct-float16-nb:1]
+examples/test_exaone.py::test_llm_exaone_1gpu[disable_weight_only-exaone_3.0_7.8b_instruct-float16-nb:4]
+examples/test_exaone.py::test_llm_exaone_1gpu[disable_weight_only-exaone_3.0_7.8b_instruct-float16-nb:4]
+examples/test_exaone.py::test_llm_exaone_1gpu[enable_weight_only-exaone_deep_2.4b-float16-nb:1]
+examples/test_exaone.py::test_llm_exaone_2gpu[exaone_3.0_7.8b_instruct-float16-nb:1]
 examples/test_gemma.py::test_llm_gemma_1gpu_summary[gemma-2-27b-it-other-bfloat16-8]
 examples/test_gemma.py::test_llm_hf_gemma_quantization_1gpu[gemma-2-27b-it-fp8-bfloat16-8]
 examples/test_gemma.py::test_hf_gemma_fp8_base_bf16_multi_lora[gemma-2-9b-it]

tests/integration/test_lists/qa/llm_sanity_test.txt

@@ -16,8 +16,8 @@ examples/test_enc_dec.py::test_llm_enc_dec_general[compare_hf-flan-t5-small-floa
 examples/test_enc_dec.py::test_llm_enc_dec_general[compare_hf-mbart-large-50-many-to-one-mmt-float16-enable_gemm_plugin-enable_attention_plugin-enable_paged_kv_cache-tp:2-pp:2-nb:1-enable_fp8]
 examples/test_enc_dec.py::test_llm_enc_dec_general[compare_hf-t5-small-float32-enable_gemm_plugin-enable_attention_plugin-enable_paged_kv_cache-tp:2-pp:1-nb:1-enable_fp8]
 examples/test_enc_dec.py::test_llm_enc_dec_general[no_compare_hf-byt5-small-float32-enable_gemm_plugin-enable_attention_plugin-enable_paged_kv_cache-tp:2-pp:1-nb:1-disable_fp8]
-examples/test_exaone.py::test_llm_exaone_1gpu[enable_weight_only-exaone-float16-nb:1]
-examples/test_exaone.py::test_llm_exaone_2gpu[exaone-float16-nb:1]
+examples/test_exaone.py::test_llm_exaone_1gpu[enable_weight_only-exaone_3.0_7.8b_instruct-float16-nb:1]
+examples/test_exaone.py::test_llm_exaone_2gpu[exaone_3.0_7.8b_instruct-float16-nb:1]
 examples/test_gemma.py::test_llm_gemma_1gpu_summary[gemma-2-27b-it-other-bfloat16-8]
 examples/test_gpt.py::test_llm_gpt_starcoder_lora_1gpu[peft-lora-starcoder2-15b-unity-copilot-starcoder2-lora_fp16-base_fp8]
 examples/test_gpt.py::test_streaming_beam[batch_size_3-return_all_generated_tokens-num_beams_4]

tests/integration/test_lists/test-db/l0_a30.yml

@@ -103,9 +103,10 @@ l0_a30:
   - examples/test_mistral.py::test_llm_mistral_v1_1gpu[mistral-7b-v0.1-float16-max_attention_window_size_4096-summarization] # 5 mins
   - examples/test_mistral.py::test_llm_mistral_v1_1gpu[mistral-7b-v0.1-float16-max_attention_window_size_4096-summarization_long] # 6 mins
   - examples/test_enc_dec.py::test_llm_enc_dec_general[compare_hf-t5-small-float32-enable_gemm_plugin-enable_attention_plugin-enable_paged_kv_cache-tp:1-pp:1-nb:2-disable_fp8]
-  - examples/test_exaone.py::test_llm_exaone_1gpu[disable_weight_only-exaone-float16-nb:1] # ? mins
-  - examples/test_exaone.py::test_llm_exaone_1gpu[disable_weight_only-exaone-float16-nb:4] # ? mins
-  - examples/test_exaone.py::test_llm_exaone_2gpu[exaone-float16-nb:1] # ? mins
+  - examples/test_exaone.py::test_llm_exaone_1gpu[disable_weight_only-exaone_3.0_7.8b_instruct-float16-nb:1] # ? mins
+  - examples/test_exaone.py::test_llm_exaone_1gpu[disable_weight_only-exaone_3.0_7.8b_instruct-float16-nb:4] # ? mins
+  - examples/test_exaone.py::test_llm_exaone_1gpu[disable_weight_only-exaone_deep_2.4b-float16-nb:4]
+  - examples/test_exaone.py::test_llm_exaone_2gpu[exaone_3.0_7.8b_instruct-float16-nb:1] # ? mins
   - examples/test_granite.py::test_llm_granite[granite-3.0-2b-instruct-bfloat16] # 5 mins
   - examples/test_draft_target_model.py::test_llm_draft_target_model_1gpu[no_streaming-gpt2-use_cpp_session-use_tokens-draft_len_4-float16-bs2] # 1 min
   - examples/test_draft_target_model.py::test_llm_draft_target_model_1gpu[no_streaming-gpt2-use_cpp_session-use_logits-draft_len_4-float16-bs2] # 1 min

tests/integration/test_lists/waives.txt

@@ -111,7 +111,7 @@ full:B200_PCIe/examples/test_enc_dec.py::test_llm_enc_dec_general[compare_hf-bar
 full:B200_PCIe/examples/test_enc_dec.py::test_llm_enc_dec_general[compare_hf-bart-large-cnn-float16-disable_gemm_plugin-enable_attention_plugin-enable_paged_kv_cache-tp:1-pp:1-nb:1] SKIP (Disable for Blackwell)
 full:B200_PCIe/examples/test_enc_dec.py::test_llm_enc_dec_general[compare_hf-bart-large-cnn-float16-disable_gemm_plugin-enable_attention_plugin-disable_paged_kv_cache-tp:1-pp:1-nb:1] SKIP (Disable for Blackwell)
 full:B200_PCIe/examples/test_enc_dec.py::test_llm_enc_dec_general[compare_hf-mbart-large-50-many-to-one-mmt-float16-enable_gemm_plugin-enable_attention_plugin-disable_paged_kv_cache-tp:1-pp:1-nb:1] SKIP (Disable for Blackwell)
-full:B200_PCIe/examples/test_exaone.py::test_llm_exaone_1gpu[enable_weight_only-exaone-float16-nb:1] SKIP (Disable for Blackwell)
+full:B200_PCIe/examples/test_exaone.py::test_llm_exaone_1gpu[enable_weight_only-exaone_3.0_7.8b_instruct-float16-nb:1] SKIP (Disable for Blackwell)
 full:B200_PCIe/examples/test_gpt.py::test_llm_gpt2_next_prompt_tuning_1gpu SKIP (Disable for Blackwell)
 full:B200_PCIe/examples/test_gpt.py::test_llm_gpt_starcoder_lora_1gpu[peft-lora-starcoder2-15b-unity-copilot-starcoder2-lora_fp16-base_fp16] SKIP (Disable for Blackwell)
 full:B200_PCIe/examples/test_gpt.py::test_llm_gpt_starcoder_lora_1gpu[peft-lora-starcoder2-15b-unity-copilot-starcoder2-lora_fp16-base_fp8] SKIP (Disable for Blackwell)
@@ -224,7 +224,7 @@ full:B200/examples/test_enc_dec.py::test_llm_enc_dec_general[compare_hf-bart-lar
 full:B200/examples/test_enc_dec.py::test_llm_enc_dec_general[compare_hf-bart-large-cnn-float16-disable_gemm_plugin-enable_attention_plugin-enable_paged_kv_cache-tp:1-pp:1-nb:1] SKIP (Disable for Blackwell)
 full:B200/examples/test_enc_dec.py::test_llm_enc_dec_general[compare_hf-bart-large-cnn-float16-disable_gemm_plugin-enable_attention_plugin-disable_paged_kv_cache-tp:1-pp:1-nb:1] SKIP (Disable for Blackwell)
 full:B200/examples/test_enc_dec.py::test_llm_enc_dec_general[compare_hf-mbart-large-50-many-to-one-mmt-float16-enable_gemm_plugin-enable_attention_plugin-disable_paged_kv_cache-tp:1-pp:1-nb:1] SKIP (Disable for Blackwell)
-full:B200/examples/test_exaone.py::test_llm_exaone_1gpu[enable_weight_only-exaone-float16-nb:1] SKIP (Disable for Blackwell)
+full:B200/examples/test_exaone.py::test_llm_exaone_1gpu[enable_weight_only-exaone_3.0_7.8b_instruct-float16-nb:1] SKIP (Disable for Blackwell)
 full:B200/examples/test_gpt.py::test_llm_gpt2_next_prompt_tuning_1gpu SKIP (Disable for Blackwell)
 full:B200/examples/test_gpt.py::test_llm_gpt_starcoder_lora_1gpu[peft-lora-starcoder2-15b-unity-copilot-starcoder2-lora_fp16-base_fp16] SKIP (Disable for Blackwell)
 full:B200/examples/test_gpt.py::test_llm_gpt_starcoder_lora_1gpu[peft-lora-starcoder2-15b-unity-copilot-starcoder2-lora_fp16-base_fp8] SKIP (Disable for Blackwell)