mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
108 lines
3.1 KiB
Python
108 lines
3.1 KiB
Python
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
|
|
from tensorrt_llm.hlapi._perf_evaluator import (LLMPerfEvaluator,
|
|
MemoryContinuousMonitorThread)
|
|
from tensorrt_llm.hlapi.llm import KvCacheConfig, ModelConfig
|
|
|
|
try:
|
|
from .grid_searcher import GridSearcher
|
|
except:
|
|
from grid_searcher import GridSearcher
|
|
|
|
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
|
|
from utils.llm_data import llm_models_root
|
|
from utils.util import force_ampere, skip_pre_ampere
|
|
|
|
|
|
def get_model_path(model_name):
|
|
return str(llm_models_root() / model_name)
|
|
|
|
|
|
llama_model_path = get_model_path("llama-models/llama-7b-hf")
|
|
|
|
|
|
def test_memory_thread():
|
|
thread = MemoryContinuousMonitorThread(0.5)
|
|
thread.start()
|
|
time.sleep(3)
|
|
thread.stop()
|
|
print(thread.memory_samples)
|
|
print('max', thread.memory_samples.get_max())
|
|
print('min', thread.memory_samples.get_min())
|
|
print('ave', thread.memory_samples.get_average())
|
|
|
|
|
|
def gen_fake_samples(samples_path: str, num_samples: int, sample_length: int):
|
|
data = {
|
|
"samples": [{
|
|
"input_ids": [20] * sample_length,
|
|
"output_len": sample_length
|
|
} for _ in range(num_samples)]
|
|
}
|
|
with open(samples_path, "w") as f:
|
|
json.dump(data, f)
|
|
|
|
|
|
@force_ampere
|
|
def test_perf_evaluator():
|
|
config = ModelConfig(llama_model_path)
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
workspace = Path(temp_dir)
|
|
samples_path = workspace / "data.json"
|
|
gen_fake_samples(samples_path, 10, 5)
|
|
|
|
# try to set some flags
|
|
kvcache_config = KvCacheConfig(enable_block_reuse=True)
|
|
|
|
evaluator = LLMPerfEvaluator.create(
|
|
config,
|
|
num_samples=10,
|
|
samples_path=samples_path,
|
|
warmup=10,
|
|
kv_cache_config=kvcache_config,
|
|
)
|
|
assert evaluator
|
|
report = evaluator.run()
|
|
report.display()
|
|
report.save_json(workspace / "report.json")
|
|
|
|
|
|
@skip_pre_ampere
|
|
def test_grid_search_tester(sample_length: int = 16,
|
|
report_root: Path = Path("./")):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
workspace = Path(temp_dir)
|
|
samples_path = workspace / "data.json"
|
|
gen_fake_samples(samples_path, 10, sample_length)
|
|
|
|
grid_searcher = GridSearcher(prune_space_for_debug=1)
|
|
|
|
report_path = workspace / "report.json"
|
|
|
|
model_config = ModelConfig(llama_model_path)
|
|
|
|
input_len = int(sample_length * 2)
|
|
output_len = int(sample_length * 2)
|
|
max_num_tokens = 1024
|
|
model_config._set_additional_options(max_output_len=output_len,
|
|
max_input_len=input_len,
|
|
max_num_tokens=max_num_tokens)
|
|
|
|
grid_searcher.evaluate(
|
|
model_config=model_config,
|
|
samples_path=samples_path,
|
|
report_dir=report_path,
|
|
memory_monitor_interval=1,
|
|
)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_perf_evaluator()
|
|
test_grid_search_tester()
|