# 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 os import time from collections import OrderedDict import tensorrt as trt import torch from allowed_configs import get_build_config from base_benchmark import BaseBenchmark, serialize_engine import tensorrt_llm from tensorrt_llm._utils import str_dtype_to_trt, trt_dtype_to_torch from tensorrt_llm.builder import Builder from tensorrt_llm.network import net_guard from tensorrt_llm.plugin.plugin import ContextFMHAType from tensorrt_llm.runtime import TensorInfo class BERTBenchmark(BaseBenchmark): def __init__(self, engine_dir, model_name, mode, batch_sizes, in_lens, dtype, output_dir, n_positions=None, max_input_len=None, max_output_len=None, max_batch_size=None, **kwargs): super().__init__(engine_dir, model_name, dtype, output_dir) self.batch_sizes = batch_sizes self.in_lens = in_lens self.build_time = 0 if engine_dir is not None: # Deserialize engine from engine directory self.serialize_path = os.path.join(engine_dir, self.engine_name) with open(self.serialize_path, 'rb') as f: engine_buffer = f.read() else: # Build engine self.use_bert_attention_plugin = False self.use_gemm_plugin = False self.use_layernorm_plugin = False self.enable_qk_half_accum = False self.enable_context_fmha = False if mode == 'plugin': self.use_bert_attention_plugin = dtype self.use_gemm_plugin = dtype self.use_layernorm_plugin = dtype for key, value in get_build_config(model_name).items(): setattr(self, key, value) # Override the n_positions/max_input_len/max_output_len/max_batch_size to value from cmd line if that's specified. if n_positions is not None: assert isinstance( n_positions, int ) and n_positions > 0, f"n_positions should be a valid int number, got {n_positions}" self.n_positions = n_positions if max_input_len is not None: assert isinstance( max_input_len, int ) and max_input_len > 0, f"max_input_len should be a valid int number, got {max_input_len}" self.max_input_len = max_input_len if max_output_len is not None: assert isinstance( max_output_len, int ) and max_output_len > 0, f"max_output_len should be a valid int number, got {max_output_len}" self.max_output_len = max_output_len if max_batch_size is not None: assert isinstance( max_batch_size, int ) and max_batch_size > 0, f"max_batch_size should be a valid int number, got {max_batch_size}" self.max_batch_size = max_batch_size if kwargs.get('force_num_layer_1', False): self.num_layers = 1 engine_buffer = self.build() assert engine_buffer is not None self.session = tensorrt_llm.runtime.Session.from_serialized_engine( engine_buffer) def get_config(self): for inlen in self.in_lens: if inlen > self.max_input_len: continue for batch_size in self.batch_sizes: if batch_size > self.max_batch_size: continue yield (batch_size, inlen) def prepare_inputs(self, config): batch_size, inlen = config[0], config[1] input_ids = torch.randint(100, (batch_size, inlen)).int().cuda() input_lengths = inlen * torch.ones( (batch_size, ), dtype=torch.int32, device='cuda') inputs = {'input_ids': input_ids, 'input_lengths': input_lengths} output_info = self.session.infer_shapes([ TensorInfo('input_ids', trt.DataType.INT32, input_ids.shape), TensorInfo('input_lengths', trt.DataType.INT32, input_lengths.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 return (inputs, outputs, stream) def build(self): bs_range = [1, (self.max_batch_size + 1) // 2, self.max_batch_size] inlen_range = [1, (self.max_input_len + 1) // 2, self.max_input_len] builder = Builder() builder_config = builder.create_builder_config( name=self.model_name, precision=self.dtype, timing_cache=None, tensor_parallel=self.world_size, # TP only parallel_build=True, num_layers=self.num_layers, num_heads=self.num_heads, num_kv_heads=self.num_heads, hidden_size=self.hidden_size, vocab_size=self.vocab_size, hidden_act=self.hidden_act, max_position_embeddings=self.n_positions, max_batch_size=self.max_batch_size, max_input_len=self.max_input_len, opt_level=self.builder_opt) # Initialize model tensorrt_llm_bert = tensorrt_llm.models.BertModel( num_layers=self.num_layers, num_heads=self.num_heads, hidden_size=self.hidden_size, vocab_size=self.vocab_size, hidden_act=self.hidden_act, max_position_embeddings=self.n_positions, type_vocab_size=self.type_vocab_size, mapping=tensorrt_llm.Mapping(world_size=self.world_size, tp_size=self.world_size)) # Module -> Network network = builder.create_network() if self.use_bert_attention_plugin: network.plugin_config.set_bert_attention_plugin( dtype=self.use_bert_attention_plugin) if self.use_gemm_plugin: network.plugin_config.set_gemm_plugin(dtype=self.use_gemm_plugin) if self.use_layernorm_plugin: network.plugin_config.set_layernorm_plugin( dtype=self.use_layernorm_plugin) if self.enable_qk_half_accum: network.plugin_config.enable_qk_half_accum() if self.enable_context_fmha: network.plugin_config.set_context_fmha(ContextFMHAType.enabled) if self.world_size > 1: network.plugin_config.set_nccl_plugin(self.dtype) with net_guard(network): # Prepare network.set_named_parameters(tensorrt_llm_bert.named_parameters()) # Forward input_ids = tensorrt_llm.Tensor( name='input_ids', dtype=trt.int32, shape=[-1, -1], dim_range=OrderedDict([('batch_size', [bs_range]), ('input_len', [inlen_range])]), ) input_lengths = tensorrt_llm.Tensor(name='input_lengths', dtype=trt.int32, shape=[-1], dim_range=OrderedDict([ ('batch_size', [bs_range]) ])) hidden_states = tensorrt_llm_bert(input_ids=input_ids, input_lengths=input_lengths) # Mark outputs hidden_states_dtype = str_dtype_to_trt(self.dtype) hidden_states.mark_output('hidden_states', hidden_states_dtype) # Network -> Engine start = time.time() engine = builder.build_engine(network, builder_config) end = time.time() self.build_time = round(end - start, 2) if self.output_dir is not None: if not os.path.exists(self.output_dir): os.makedirs(self.output_dir) self.serialize_path = os.path.join(self.output_dir, self.engine_name) serialize_engine(engine, self.serialize_path) if self.runtime_rank == 0: config_path = os.path.join(self.output_dir, 'config.json') builder_config.plugin_config = network.plugin_config builder.save_config(builder_config, config_path) return engine def run(self, inputs, config): ok = self.session.run(*inputs) assert ok, "Runtime execution failed" torch.cuda.synchronize() def report(self, config, latency, percentile95, percentile99, peak_gpu_used): if self.runtime_rank == 0: line = '[BENCHMARK] ' + ( f'model_name {self.model_name} world_size {self.world_size} precision {self.dtype} ' f'batch_size {config[0]} input_length {config[1]} gpu_peak_mem(gb) {peak_gpu_used} ' f'build_time(s) {self.build_time} percentile95(ms) {percentile95} ' f'percentile99(ms) {percentile99} latency(ms) {latency}') print(line) def report(self, config, latency, percentile95, percentile99, peak_gpu_used, csv): report_dict = super().get_report_dict() batch_size, inlen = config[0], config[1] report_dict["num_heads"] = self.num_heads report_dict["num_kv_heads"] = self.num_heads report_dict["num_layers"] = self.num_layers report_dict["hidden_size"] = self.hidden_size report_dict["vocab_size"] = self.vocab_size report_dict["batch_size"] = batch_size report_dict["input_length"] = inlen report_dict["output_length"] = "n/a" report_dict["latency(ms)"] = latency report_dict["build_time(s)"] = self.build_time report_dict["tokens_per_sec"] = "n/a" report_dict["percentile95(ms)"] = percentile95 report_dict["percentile99(ms)"] = percentile99 report_dict["gpu_peak_mem(gb)"] = peak_gpu_used if self.runtime_rank == 0: if csv: line = ",".join([str(v) for v in report_dict.values()]) print(line) with open(self.get_csv_filename(), "a") as file: file.write(line + "\n") else: kv_pairs = [f"{k} {v}" for k, v in report_dict.items()] line = '[BENCHMARK] ' + " ".join(kv_pairs) print(line)