File size: 3,105 Bytes
4a28d4d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | import os
import subprocess
import fire
import yaml
def get_cmd(model_path, backend, engine_config, data_config):
assert backend in ['turbomind', 'pytorch']
current_dir = os.path.dirname(os.path.abspath(__file__))
dataset_path = data_config.pop('dataset_path')
cmd = ['python3', f'{current_dir}/profile_throughput.py', dataset_path, model_path, '--backend', backend]
for key, value in engine_config.items():
# profile_throughput.py uses "--concurrency" to pass the "max_batch_size" value
if key == 'max_batch_size':
key = 'concurrency'
# change the key like 'cache_max_entry_count' to 'cache-max-entry-count' to suit the optional
# arguments in "python3 benchmark/profile_throughput.py"
key = key.replace('_', '-')
cmd.append(f'--{key}')
cmd.append(str(value))
for key, value in data_config.items():
# change the key like 'sharegpt_output_len' to 'sharegpt-output-len' to suit the optional
# arguments in "python3 benchmark/profile_throughput.py"
key = key.replace('_', '-')
cmd.append(f'--{key}')
cmd.append(str(value))
return cmd
def benchmark(model_path, backend, engine_config, data_config):
"""Benchmark the performance with the given configuration.
Args:
model_path: Path to the model.
:param backend: Backend to use.
:param engine_config: Configuration for the inference engine.
:param data_config: Configuration for the data.
"""
model_name = os.path.basename(model_path)
bs = engine_config['max_batch_size']
cach_ratio = engine_config.get('cache_max_entry_count', 0.8)
tp = engine_config.get('tp', 1)
output_file = f'benchmark_throughput_{model_name}_{backend}_bs{bs}_tp{tp}_cache{cach_ratio}.csv'
try:
if isinstance(data_config, dict):
data_config = [data_config]
assert isinstance(data_config, list) and all(isinstance(d, dict) for d in data_config)
for _data_config in data_config:
_data_config['csv'] = output_file
cmd = get_cmd(model_path, backend, engine_config, _data_config)
print(f"Running command: {' '.join(cmd)}")
subprocess.run(cmd, check=True)
except Exception as e:
print(f'exception happened, {e}')
def main(model_path=None, backend=None, config_path=None):
with open(config_path) as f:
config = yaml.safe_load(f)
engine_configs = config['engine']
data_config = config['data']
if isinstance(engine_configs, dict):
engine_configs = [engine_configs]
assert isinstance(engine_configs, list) and all(isinstance(s, dict) for s in engine_configs)
for engine_config in engine_configs:
# The model_path provided by the user will override the model_path in the config file.
model_path = model_path or engine_config.pop('model_path')
engine_config.pop('model_path', '')
benchmark(model_path, backend, engine_config, data_config)
if __name__ == '__main__':
fire.Fire(main)
|