mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
105 lines
4.4 KiB
Python
105 lines
4.4 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
import pytest
|
|
from defs.common import convert_weights, parse_mpi_cmd, venv_mpi_check_call
|
|
from defs.conftest import get_device_memory, get_sm_version
|
|
from defs.trt_test_alternative import check_call
|
|
|
|
# skip trt flow cases on post-Blackwell-Ultra
|
|
if get_sm_version() >= 103:
|
|
pytest.skip(
|
|
"TRT workflow tests are not supported on post Blackwell-Ultra architecture",
|
|
allow_module_level=True)
|
|
|
|
|
|
# @pytest.mark.skip_less_device(2)
|
|
@pytest.mark.parametrize("num_beams", [1, 2, 4],
|
|
ids=lambda num_beams: f'nb:{num_beams}')
|
|
@pytest.mark.parametrize(
|
|
"use_gpt_attention_plugin", [True, False],
|
|
ids=["enable_attention_plugin", "disable_attention_plugin"])
|
|
@pytest.mark.parametrize("use_gemm_plugin", [True, False],
|
|
ids=["enable_gemm_plugin", "disable_gemm_plugin"])
|
|
@pytest.mark.parametrize("context_fmha_type", [
|
|
"enable_context_fmha", "enable_context_fmha_fp32_acc",
|
|
"disable_context_fmha"
|
|
])
|
|
@pytest.mark.parametrize("dtype", ['float16', 'bfloat16'])
|
|
def test_llm_internlm2_7b_1node_1gpu(internlm2_example_root,
|
|
llm_internlm2_7b_model_root,
|
|
llm_datasets_root, llm_rouge_root,
|
|
llm_venv, cmodel_dir, engine_dir,
|
|
use_gpt_attention_plugin, use_gemm_plugin,
|
|
context_fmha_type, dtype, num_beams):
|
|
"Build & Run internlm2-7b with 1 gpu"
|
|
if dtype == "bfloat16" and not use_gemm_plugin:
|
|
pytest.skip("Please use gemm plugin when dtype is bfloat16.")
|
|
if num_beams == 4 and get_device_memory() < 50000:
|
|
pytest.skip("device memory is insufficient.")
|
|
|
|
model_dir = convert_weights(llm_venv=llm_venv,
|
|
example_root=f"{internlm2_example_root}",
|
|
cmodel_dir=cmodel_dir,
|
|
model="internlm2-7b",
|
|
model_path=llm_internlm2_7b_model_root,
|
|
data_type=dtype,
|
|
gpus=1,
|
|
tp_size=1)
|
|
|
|
build_cmd = [
|
|
"python3 -m tensorrt_llm.commands.build",
|
|
f"--checkpoint_dir={model_dir}",
|
|
f"--output_dir={engine_dir}",
|
|
f"--max_beam_width={num_beams}",
|
|
f"--max_batch_size=1",
|
|
]
|
|
|
|
if use_gpt_attention_plugin:
|
|
build_cmd.append("--remove_input_padding=enable")
|
|
build_cmd.append(f"--gpt_attention_plugin={dtype}")
|
|
else:
|
|
build_cmd.append("--gpt_attention_plugin=disable")
|
|
build_cmd.append("--remove_input_padding=disable")
|
|
build_cmd.append("--paged_kv_cache=disable")
|
|
|
|
if use_gemm_plugin:
|
|
build_cmd.append(f"--gemm_plugin={dtype}")
|
|
else:
|
|
build_cmd.append("--gemm_plugin=disable")
|
|
|
|
if context_fmha_type == "enable_context_fmha":
|
|
build_cmd.append("--context_fmha=enable")
|
|
elif context_fmha_type == "disable_context_fmha":
|
|
build_cmd.append("--context_fmha=disable")
|
|
|
|
print("Building engines...")
|
|
check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env)
|
|
|
|
print('Run internlm2-7b...')
|
|
data_type = "fp16" if dtype == "float16" else "bf16"
|
|
summary_cmd = [
|
|
f"{internlm2_example_root}/../../../summarize.py", "--test_trt_llm",
|
|
"--hf_model_dir", llm_internlm2_7b_model_root, "--engine_dir",
|
|
engine_dir, "--data_type", data_type, "--check_accuracy",
|
|
f"--num_beams={num_beams}", f"--dataset_dir={llm_datasets_root}",
|
|
f"--rouge_dir={llm_rouge_root}"
|
|
]
|
|
if context_fmha_type == "enable_context_fmha_fp32_acc":
|
|
summary_cmd.append("--enable_context_fmha_fp32_acc")
|
|
|
|
venv_mpi_check_call(
|
|
llm_venv, parse_mpi_cmd(["mpirun", "-n", "1", "--allow-run-as-root"]),
|
|
summary_cmd)
|