| import asyncio |
| import json |
| import math |
| import os |
| import threading |
| import time |
| from dataclasses import asdict, dataclass, field |
| from pathlib import Path |
| from typing import Iterable, List, Optional |
|
|
| import numpy as np |
| import torch |
|
|
| from tensorrt_llm import logger |
|
|
| from .._utils import release_gc |
| from ..profiler import device_memory_info, host_memory_info |
| from .llm import LLM, SamplingParams |
| from .utils import is_directory_empty, print_colored |
|
|
|
|
| @dataclass |
| class PerfItem: |
| start: float |
| end: float |
| num_out_tokens: int = 0 |
|
|
| @property |
| def lantency(self) -> float: |
| return self.end - self.start |
|
|
|
|
| @dataclass |
| class Report: |
| num_samples: int |
| total_latency: float |
| avg_latency: float |
| seq_throughput: float |
| token_throughput: float |
| ave_tokens_per_sample: float |
| memory_usage_samples: "MemoryContinuousMonitorThread.RecordList" = field( |
| default_factory=list) |
|
|
| def display(self): |
| print_colored("Performance Report:\n", color="bold_green") |
| print(f"num_samples: {self.num_samples}") |
| print(f"total_latency: {self.total_latency:.2f}") |
| print(f"avg_latency: {self.avg_latency:.2f}") |
| print(f"seq_throughput: {self.seq_throughput:.2f}") |
| print(f"token_throughput: {self.token_throughput:.2f}") |
| print(f"average tokens per sample: {self.ave_tokens_per_sample:.2f}") |
| if self.memory_usage_samples: |
| print("Memory Usage:\n") |
| min_record, max_record, average_record = self.memory_usage_samples.get_min( |
| ), self.memory_usage_samples.get_max( |
| ), self.memory_usage_samples.get_average() |
| print( |
| f" host memory usage: (min: {min_record[0]}, max: {max_record[0]}, average: {average_record[0]})" |
| ) |
| print( |
| f" gpu memory usage: (min: {min_record[1]}, max: {max_record[1]}, average: {average_record[1]})" |
| ) |
| print_colored("__________________________________\n", color="green") |
|
|
| def save_json(self, path: Path, **kwargs): |
| with open(path, 'w') as f: |
| data = self.__dict__.copy() |
| data.update(kwargs) |
| data["memory_usage_samples"] = [ |
| asdict(record) for record in self.memory_usage_samples |
| ] |
| json.dump(data, f) |
|
|
|
|
| @dataclass |
| class Sample: |
| input_ids: List[int] |
| output_len: int |
|
|
|
|
| class MemoryContinuousMonitorThread(threading.Thread): |
| ''' Monitor the host memory and GPU memory usage. ''' |
|
|
| @dataclass |
| class Record: |
| time: float |
| host_memory: float |
| gpu_memory: List[float] |
|
|
| def to_dict(self): |
| return dict(time=self.time, |
| host_memory=self.host_memory, |
| gpu_memory=self.gpu_memory) |
|
|
| class RecordList(list): |
|
|
| def __init__(self, *args, **kwargs): |
| super(MemoryContinuousMonitorThread.RecordList, |
| self).__init__(*args, **kwargs) |
|
|
| def get_min(self): |
| return self._get_memory_usage('min') |
|
|
| def get_max(self): |
| return self._get_memory_usage('max') |
|
|
| def get_average(self): |
| return self._get_memory_usage('average') |
|
|
| def _get_memory_usage(self, op: str): |
| host_memory = [record.host_memory for record in self] |
| gpu_memory = [record.gpu_memory for record in self] |
|
|
| ops = dict(min=np.min, max=np.max, average=np.mean) |
| theop = ops[op] |
|
|
| return theop(host_memory), [theop(gpu) for gpu in zip(*gpu_memory)] |
|
|
| def __init__(self, sampling_interval: float = 1): |
| super(MemoryContinuousMonitorThread, self).__init__() |
| self.sampling_interval = sampling_interval |
| self._stop_event = threading.Event() |
| self.memory_samples = MemoryContinuousMonitorThread.RecordList() |
|
|
| def run(self): |
| while not self._stop_event.is_set(): |
| record = self.monitor() |
| logger.info(f'record: {record}') |
| self.memory_samples.append(record) |
| time.sleep(self.sampling_interval) |
|
|
| def monitor(self) -> "MemoryContinuousMonitorThread.Record": |
| return self.Record(time.time(), get_host_memory_usage(), |
| list(get_gpu_memory_usage())) |
|
|
| def stop(self): |
| self._stop_event.set() |
|
|
|
|
| def get_host_memory_usage() -> float: |
| return host_memory_info(os.getpid())[0] / 1024**3 |
|
|
|
|
| def get_gpu_memory_usage() -> Iterable[float]: |
| for device in range(torch.cuda.device_count()): |
| yield device_memory_info(device)[0] / 1024**3 |
|
|
|
|
| class LLMPerfEvaluator: |
|
|
| @classmethod |
| def create(cls, |
| model: str, |
| samples_path: Path, |
| num_samples: int = -1, |
| warmup: int = 100, |
| batch_size: int = -1, |
| engine_cache_path: Optional[Path] = None, |
| memory_monitor_interval: Optional[int] = None, |
| **kwargs) -> 'LLMPerfEvaluator': |
| ''' |
| Args: |
| model: The model name or a local path to the model directory. |
| samples_path: path to the input data samples |
| num_samples: number of the heading samples to run, if set to -1, all samples will be used |
| warmup: number of samples for warmup runs |
| batch_size: batch size for the runs, if left default, the batch size will be the same as the number of samples |
| engine_cache_path: path to the engine file, if provided, the engine will save the built engine to the path and reuse it for the next runs |
| memory_monitor_interval: the interval to monitor the host and GPU memory usage, if set to None, the memory monitor will be disabled |
| kwargs: the additional arguments are for the LLM constructor |
| ''' |
|
|
| from_cache = False |
| if engine_cache_path and Path.exists( |
| engine_cache_path |
| ) and not is_directory_empty(engine_cache_path): |
| print(f"Loading engine from {engine_cache_path}\n") |
| from_cache = True |
| model = engine_cache_path |
|
|
| memory_monitor_thread = None |
| if memory_monitor_interval is not None: |
| memory_monitor_thread = MemoryContinuousMonitorThread( |
| sampling_interval=memory_monitor_interval) |
| memory_monitor_thread.start() |
|
|
| |
| try: |
| llm = LLM(model, skip_tokenizer_init=True, **kwargs) |
| except Exception as e: |
| logger.error(f"Failed to create LLM with {model} and {kwargs}") |
| raise e |
|
|
| if engine_cache_path is not None and not from_cache: |
| print_colored(f"Saving engine to {engine_cache_path}\n", "green") |
| llm.save(engine_cache_path) |
|
|
| samples: List[Sample] = list(cls.load_dataset(samples_path)) |
| assert len( |
| samples |
| ) >= num_samples, f"num_samples {num_samples} is too large. The dataset only has {len(samples)} samples." |
| samples = samples[:num_samples] if num_samples > 0 else samples |
|
|
| return cls(llm, |
| warmup=warmup, |
| samples=samples, |
| max_num_samples=num_samples, |
| batch_size=batch_size, |
| memory_monitor_thread=memory_monitor_thread) |
|
|
| def __init__(self, |
| llm: LLM, |
| samples: List[Sample], |
| warmup: int, |
| max_num_samples: int, |
| batch_size: int, |
| memory_monitor_thread: Optional[ |
| MemoryContinuousMonitorThread] = None): |
| self.llm = llm |
| self.samples = samples |
| self.warmup = warmup |
| self.max_num_samples = max_num_samples |
| self.perf_items = [] |
| self.batch_size = batch_size if batch_size > 0 else len(self.samples) |
| self.memory_monitor_thread = memory_monitor_thread |
|
|
| self.start = None |
| self.end = None |
|
|
| def run(self, end_id: int = -1, beam_width: int = 1) -> Report: |
| |
| self.perf_items = [] |
| sample_offset = 0 |
|
|
| sampling_params = SamplingParams( |
| end_id=end_id, |
| pad_id=end_id, |
| beam_width=beam_width, |
| ) |
|
|
| async def lane(num_tasks: int, |
| sampling_params: SamplingParams, |
| warmup=False): |
| nonlocal sample_offset |
|
|
| for i in range(num_tasks): |
| sample = self.samples[sample_offset] |
| sample_offset += 1 |
| sampling_params.max_new_tokens = sample.output_len |
| sampling_params.end_id = -2 |
| sampling_params.pad_id = -2 |
|
|
| start = time.time() |
| output = self.llm.generate_async( |
| sample.input_ids, sampling_params=sampling_params) |
| output = await output.aresult() |
| end = time.time() |
|
|
| perf_item = PerfItem(start=start, |
| end=end, |
| num_out_tokens=sum( |
| beam_output.length |
| for beam_output in output.outputs)) |
| if not warmup: |
| self.perf_items.append(perf_item) |
|
|
| if self.warmup > 0: |
| logger.warning("warming up ...") |
| for i in range(math.ceil(self.warmup / len(self.samples))): |
| asyncio.run( |
| lane(min(self.warmup, len(self.samples)), |
| sampling_params, |
| warmup=True)) |
| sample_offset = 0 |
|
|
| logger.warning("running ...") |
| self.start = time.time() |
|
|
| async def run_lanes(): |
| num_tasks = len(self.samples) // self.batch_size |
| lanes = [ |
| lane(num_tasks, sampling_params) for _ in range(self.batch_size) |
| ] |
| await asyncio.gather(*lanes) |
|
|
| asyncio.run(run_lanes()) |
| self.end = time.time() |
|
|
| return self._generate_report() |
|
|
| @staticmethod |
| def load_dataset(path: Path) -> Iterable[Sample]: |
| with open(path, 'r') as f: |
| dataset = json.load(f) |
|
|
| for sample in dataset["samples"]: |
| yield Sample(input_ids=sample["input_ids"], |
| output_len=sample["output_len"]) |
|
|
| def _generate_report(self) -> Report: |
| num_samples = len(self.perf_items) |
| total_latency = self.end - self.start |
| avg_latency = total_latency / num_samples |
| seq_throughput = num_samples / total_latency |
| token_throughput = sum( |
| [perf_item.num_out_tokens |
| for perf_item in self.perf_items]) / total_latency |
| total_tokens = sum( |
| [perf_item.num_out_tokens for perf_item in self.perf_items]) |
|
|
| return Report( |
| num_samples=num_samples, |
| total_latency=total_latency, |
| avg_latency=avg_latency, |
| seq_throughput=seq_throughput, |
| token_throughput=token_throughput, |
| ave_tokens_per_sample=total_tokens / num_samples, |
| memory_usage_samples=self.memory_monitor_thread.memory_samples |
| if self.memory_monitor_thread else []) |
|
|
| def __enter__(self): |
| return self |
|
|
| def __exit__(self, exc_type, exc_value, traceback): |
| self.llm.__exit__(exc_type, exc_value, traceback) |
| del self.llm |
| release_gc() |
|
|
| if self.memory_monitor_thread: |
| self.memory_monitor_thread.stop() |
| self.memory_monitor_thread.join() |
| del self.memory_monitor_thread |
|
|
