codekingpro/portable-devtools
114k
1# -------------------------------------------------------------------------
2# Copyright (c) Microsoft Corporation. All rights reserved.
3# Licensed under the MIT License. See License.txt in the project root for
4# license information.
5# --------------------------------------------------------------------------
6
7import csv
8import logging
9import os
10import random
11import sys
12import time
13import timeit
14from abc import ABC, abstractmethod
15from concurrent.futures import ThreadPoolExecutor
16from datetime import datetime
17from enum import Enum
18from time import sleep
19from typing import Any
20
21import numpy
22import torch
23import transformers
24from packaging import version
25
26import onnxruntime
27
28logger = logging.getLogger(__name__)
29
30
31class Precision(Enum):
32 FLOAT32 = "fp32"
33 FLOAT16 = "fp16"
34 INT8 = "int8"
35 INT4 = "int4"
36
37 def __str__(self):
38 return self.value
39
40
41class OptimizerInfo(Enum):
42 # no_opt means using the raw ONNX model, but OnnxRuntime might still apply optimization as long as
43 # graph optimization level is not 0 (disable all).
44 NOOPT = "no_opt"
45 BYORT = "by_ort"
46 BYSCRIPT = "by_script"
47
48 def __str__(self):
49 return self.value
50
51
52class ConfigModifier:
53 def __init__(self, num_layers):
54 self.num_layers = num_layers
55
56 def modify(self, config):
57 if self.num_layers is None:
58 return
59 if hasattr(config, "num_hidden_layers"):
60 config.num_hidden_layers = self.num_layers
61 logger.info(f"Modifying pytorch model's number of hidden layers to: {self.num_layers}")
62 if hasattr(config, "encoder_layers"):
63 config.encoder_layers = self.num_layers
64 logger.info(f"Modifying pytorch model's number of encoder layers to: {self.num_layers}")
65 if hasattr(config, "decoder_layers "):
66 config.decoder_layers = self.num_layers
67 logger.info(f"Modifying pytorch model's number of decoder layers to: {self.num_layers}")
68
69 def get_layer_num(self):
70 return self.num_layers
71
72
73IO_BINDING_DATA_TYPE_MAP = {
74 "float32": numpy.float32,
75 # TODO: Add more.
76}
77
78
79def create_onnxruntime_session(
80 onnx_model_path,
81 use_gpu,
82 provider=None,
83 enable_all_optimization=True,
84 num_threads=-1,
85 enable_profiling=False,
86 verbose=False,
87 enable_mlas_gemm_fastmath_arm64_bfloat16=False,
88 provider_options={}, # map execution provider name to its option # noqa: B006
89):
90 sess_options = onnxruntime.SessionOptions()
91
92 if enable_all_optimization:
93 sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
94 else:
95 sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_BASIC
96
97 if enable_profiling:
98 sess_options.enable_profiling = True
99
100 if num_threads > 0:
101 sess_options.intra_op_num_threads = num_threads
102 logger.debug(f"Session option: intra_op_num_threads={sess_options.intra_op_num_threads}")
103
104 if verbose:
105 sess_options.log_severity_level = 0
106 else:
107 sess_options.log_severity_level = 4
108
109 if provider in onnxruntime.get_available_providers():
110 providers = [provider]
111 elif use_gpu:
112 if provider == "dml":
113 providers = ["DmlExecutionProvider", "CPUExecutionProvider"]
114 elif provider == "migraphx":
115 providers = [
116 "MIGraphXExecutionProvider",
117 "CPUExecutionProvider",
118 ]
119 elif provider == "cuda" or provider is None:
120 providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
121 elif provider == "tensorrt":
122 providers = [
123 "TensorrtExecutionProvider",
124 "CUDAExecutionProvider",
125 "CPUExecutionProvider",
126 ]
127 else:
128 raise RuntimeError(f"The execution provider is not supported: {provider}")
129 else:
130 providers = ["CPUExecutionProvider"]
131
132 if provider_options:
133 providers = [(name, provider_options[name]) if name in provider_options else name for name in providers]
134
135 if enable_mlas_gemm_fastmath_arm64_bfloat16:
136 sess_options.add_session_config_entry("mlas.enable_gemm_fastmath_arm64_bfloat16", "1")
137
138 session = None
139 try:
140 session = onnxruntime.InferenceSession(onnx_model_path, sess_options, providers=providers)
141 except Exception:
142 logger.exception(f"Failed to create session for {onnx_model_path} with providers={providers}")
143
144 return session
145
146
147def setup_logger(verbose=True):
148 if verbose:
149 logging.basicConfig(
150 format="[%(filename)s:%(lineno)s - %(funcName)20s()] %(message)s",
151 level=logging.DEBUG,
152 )
153 else:
154 logging.basicConfig(format="%(message)s", level=logging.INFO)
155 logging.getLogger("transformers").setLevel(logging.WARNING)
156
157
158def prepare_environment(cache_dir, output_dir, use_gpu, provider=None):
159 if cache_dir and not os.path.exists(cache_dir):
160 os.makedirs(cache_dir)
161
162 if output_dir and not os.path.exists(output_dir):
163 os.makedirs(output_dir)
164
165 if use_gpu:
166 if provider == "dml":
167 assert "DmlExecutionProvider" in onnxruntime.get_available_providers(), (
168 "Please install onnxruntime-directml package to test GPU inference."
169 )
170
171 else:
172 assert not set(onnxruntime.get_available_providers()).isdisjoint(
173 ["CUDAExecutionProvider", "MIGraphXExecutionProvider"]
174 ), "Please install onnxruntime-gpu package, or install migraphx, to test GPU inference."
175
176 logger.info(f"PyTorch Version:{torch.__version__}")
177 logger.info(f"Transformers Version:{transformers.__version__}")
178 logger.info(f"OnnxRuntime Version:{onnxruntime.__version__}")
179
180 # Support three major versions of PyTorch and OnnxRuntime, and up to 9 months of transformers.
181 assert version.parse(torch.__version__) >= version.parse("1.10.0")
182 assert version.parse(transformers.__version__) >= version.parse("4.12.0")
183 assert version.parse(onnxruntime.__version__) >= version.parse("1.10.0")
184
185
186def get_latency_result(latency_list, batch_size):
187 latency_ms = sum(latency_list) / float(len(latency_list)) * 1000.0
188 latency_variance = numpy.var(latency_list, dtype=numpy.float64) * 1000.0
189 throughput = batch_size * (1000.0 / latency_ms)
190
191 return {
192 "test_times": len(latency_list),
193 "latency_variance": f"{latency_variance:.2f}",
194 "latency_90_percentile": f"{numpy.percentile(latency_list, 90) * 1000.0:.2f}",
195 "latency_95_percentile": f"{numpy.percentile(latency_list, 95) * 1000.0:.2f}",
196 "latency_99_percentile": f"{numpy.percentile(latency_list, 99) * 1000.0:.2f}",
197 "average_latency_ms": f"{latency_ms:.2f}",
198 "QPS": f"{throughput:.2f}",
199 }
200
201
202def output_details(results, csv_filename):
203 with open(csv_filename, mode="a", newline="", encoding="ascii") as csv_file:
204 column_names = [
205 "engine",
206 "version",
207 "providers",
208 "device",
209 "precision",
210 "optimizer",
211 "io_binding",
212 "model_name",
213 "inputs",
214 "threads",
215 "batch_size",
216 "sequence_length",
217 "custom_layer_num",
218 "datetime",
219 "test_times",
220 "QPS",
221 "average_latency_ms",
222 "latency_variance",
223 "latency_90_percentile",
224 "latency_95_percentile",
225 "latency_99_percentile",
226 ]
227
228 csv_writer = csv.DictWriter(csv_file, fieldnames=column_names)
229 csv_writer.writeheader()
230 for result in results:
231 csv_writer.writerow(result)
232
233 logger.info(f"Detail results are saved to csv file: {csv_filename}")
234
235
236def output_summary(results, csv_filename, args):
237 with open(csv_filename, mode="a", newline="", encoding="ascii") as csv_file:
238 header_names = [
239 "model_name",
240 "inputs",
241 "custom_layer_num",
242 "engine",
243 "version",
244 "providers",
245 "device",
246 "precision",
247 "optimizer",
248 "io_binding",
249 "threads",
250 ]
251 data_names = []
252 for batch_size in args.batch_sizes:
253 if args.sequence_lengths == [""]:
254 data_names.append(f"b{batch_size}")
255 else:
256 for sequence_length in args.sequence_lengths:
257 data_names.append(f"b{batch_size}_s{sequence_length}")
258
259 csv_writer = csv.DictWriter(csv_file, fieldnames=header_names + data_names)
260 csv_writer.writeheader()
261 for model_name in args.models:
262 for input_count in [1, 2, 3]:
263 for engine_name in args.engines:
264 for io_binding in [True, False, ""]:
265 for threads in args.num_threads:
266 row = {}
267 for result in results:
268 if (
269 result["model_name"] == model_name
270 and result["inputs"] == input_count
271 and result["engine"] == engine_name
272 and result["io_binding"] == io_binding
273 and result["threads"] == threads
274 ):
275 headers = {k: v for k, v in result.items() if k in header_names}
276 if not row:
277 row.update(headers)
278 row.update(dict.fromkeys(data_names, ""))
279 else:
280 for k in header_names:
281 assert row[k] == headers[k]
282 b = result["batch_size"]
283 s = result["sequence_length"]
284 if s:
285 row[f"b{b}_s{s}"] = result["average_latency_ms"]
286 else:
287 row[f"b{b}"] = result["average_latency_ms"]
288 if row:
289 csv_writer.writerow(row)
290
291 logger.info(f"Summary results are saved to csv file: {csv_filename}")
292
293
294def output_fusion_statistics(model_fusion_statistics, csv_filename):
295 with open(csv_filename, mode="a", newline="", encoding="ascii") as csv_file:
296 column_names = [
297 "model_filename",
298 "datetime",
299 "transformers",
300 "torch",
301 *list(next(iter(model_fusion_statistics.values())).keys()),
302 ]
303 csv_writer = csv.DictWriter(csv_file, fieldnames=column_names)
304 csv_writer.writeheader()
305 for key in model_fusion_statistics:
306 model_fusion_statistics[key]["datetime"] = str(datetime.now())
307 model_fusion_statistics[key]["transformers"] = transformers.__version__
308 model_fusion_statistics[key]["torch"] = torch.__version__
309 model_fusion_statistics[key]["model_filename"] = key
310 csv_writer.writerow(model_fusion_statistics[key])
311 logger.info(f"Fusion statistics is saved to csv file: {csv_filename}")
312
313
314def inference_ort(ort_session, ort_inputs, result_template, repeat_times, batch_size, warm_up_repeat=0):
315 result = {}
316 timeit.repeat(lambda: ort_session.run(None, ort_inputs), number=1, repeat=warm_up_repeat) # Dry run
317 latency_list = timeit.repeat(lambda: ort_session.run(None, ort_inputs), number=1, repeat=repeat_times)
318 result.update(result_template)
319 result.update({"io_binding": False})
320 result.update(get_latency_result(latency_list, batch_size))
321 return result
322
323
324def inference_ort_with_io_binding(
325 ort_session,
326 ort_inputs,
327 result_template,
328 repeat_times,
329 ort_output_names,
330 ort_outputs,
331 output_buffers,
332 output_buffer_max_sizes,
333 batch_size,
334 device,
335 data_type=numpy.longlong,
336 warm_up_repeat=0,
337):
338 result = {}
339
340 # Bind inputs and outputs to onnxruntime session
341 io_binding = ort_session.io_binding()
342 # Bind inputs to device
343 for name in ort_inputs:
344 np_input = torch.from_numpy(ort_inputs[name]).to(device)
345 input_type = IO_BINDING_DATA_TYPE_MAP.get(str(ort_inputs[name].dtype), data_type)
346 io_binding.bind_input(
347 name,
348 np_input.device.type,
349 0,
350 input_type,
351 np_input.shape,
352 np_input.data_ptr(),
353 )
354 # Bind outputs buffers with the sizes needed if not allocated already
355 if len(output_buffers) == 0:
356 allocateOutputBuffers(output_buffers, output_buffer_max_sizes, device)
357
358 for i, ort_output_name in enumerate(ort_output_names):
359 io_binding.bind_output(
360 ort_output_name,
361 output_buffers[i].device.type,
362 0,
363 numpy.float32,
364 ort_outputs[i].shape,
365 output_buffers[i].data_ptr(),
366 )
367
368 timeit.repeat(
369 lambda: ort_session.run_with_iobinding(io_binding),
370 number=1,
371 repeat=warm_up_repeat,
372 ) # Dry run
373
374 latency_list = timeit.repeat(
375 lambda: ort_session.run_with_iobinding(io_binding),
376 number=1,
377 repeat=repeat_times,
378 )
379 result.update(result_template)
380 result.update({"io_binding": True})
381 result.update(get_latency_result(latency_list, batch_size))
382 return result
383
384
385def allocateOutputBuffers(output_buffers, output_buffer_max_sizes, device): # noqa: N802
386 # Allocate output tensors with the largest test size needed. So the allocated memory can be reused
387 # for each test run.
388
389 for i in output_buffer_max_sizes:
390 output_buffers.append(torch.empty(i, dtype=torch.float32, device=device))
391
392
393def set_random_seed(seed=123):
394 """Set random seed manually to get deterministic results"""
395 random.seed(seed)
396 numpy.random.seed(seed)
397 torch.manual_seed(seed)
398 torch.cuda.manual_seed(seed)
399 torch.cuda.manual_seed_all(seed)
400 # torch.backends.cudnn.enabled = False
401 # torch.backends.cudnn.benchmark = False
402 # torch.backends.cudnn.deterministic = True
403
404
405def get_gpu_info() -> list[dict[str, Any]] | None:
406 from py3nvml.py3nvml import ( # noqa: PLC0415
407 NVMLError,
408 nvmlDeviceGetCount,
409 nvmlDeviceGetHandleByIndex,
410 nvmlDeviceGetMemoryInfo,
411 nvmlDeviceGetName,
412 nvmlInit,
413 nvmlShutdown,
414 )
415
416 try:
417 nvmlInit()
418 result = []
419 device_count = nvmlDeviceGetCount()
420 if not isinstance(device_count, int):
421 return None
422
423 for i in range(device_count):
424 info = nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(i))
425 if isinstance(info, str):
426 return None
427 result.append(
428 {
429 "id": i,
430 "name": nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(i)),
431 "total": info.total,
432 "free": info.free,
433 "used": info.used,
434 }
435 )
436 nvmlShutdown()
437 return result
438 except NVMLError as error:
439 print("Error fetching GPU information using nvml: %s", error)
440 return None
441
442
443class MemoryMonitor(ABC):
444 def __init__(self, keep_measuring=True):
445 self.keep_measuring = keep_measuring
446
447 def measure_cpu_usage(self):
448 import psutil # noqa: PLC0415
449
450 max_usage = 0
451 while True:
452 max_usage = max(max_usage, psutil.Process(os.getpid()).memory_info().rss / 1024**2)
453 sleep(0.005) # 5ms
454 if not self.keep_measuring:
455 break
456 return max_usage
457
458 @abstractmethod
459 def measure_gpu_usage(self) -> list[dict[str, Any]] | None:
460 raise NotImplementedError()
461
462
463class CudaMemoryMonitor(MemoryMonitor):
464 def __init__(self, keep_measuring=True):
465 super().__init__(keep_measuring)
466
467 def measure_gpu_usage(self) -> list[dict[str, Any]] | None:
468 from py3nvml.py3nvml import ( # noqa: PLC0415
469 NVMLError,
470 nvmlDeviceGetCount,
471 nvmlDeviceGetHandleByIndex,
472 nvmlDeviceGetMemoryInfo,
473 nvmlDeviceGetName,
474 nvmlInit,
475 nvmlShutdown,
476 )
477
478 max_gpu_usage = []
479 gpu_name = []
480 try:
481 nvmlInit()
482 device_count = nvmlDeviceGetCount()
483 if not isinstance(device_count, int):
484 logger.error(f"nvmlDeviceGetCount result is not integer: {device_count}")
485 return None
486
487 max_gpu_usage = [0 for i in range(device_count)]
488 gpu_name = [nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(i)) for i in range(device_count)]
489 while True:
490 for i in range(device_count):
491 info = nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(i))
492 if isinstance(info, str):
493 logger.error(f"nvmlDeviceGetMemoryInfo returns str: {info}")
494 return None
495 max_gpu_usage[i] = max(max_gpu_usage[i], info.used / 1024**2)
496 sleep(0.005) # 5ms
497 if not self.keep_measuring:
498 break
499 nvmlShutdown()
500 return [
501 {
502 "device_id": i,
503 "name": gpu_name[i],
504 "max_used_MB": max_gpu_usage[i],
505 }
506 for i in range(device_count)
507 ]
508 except NVMLError as error:
509 logger.error("Error fetching GPU information using nvml: %s", error)
510 return None
511
512
513class RocmMemoryMonitor(MemoryMonitor):
514 def __init__(self, keep_measuring=True):
515 super().__init__(keep_measuring)
516 rocm_smi_path = "/opt/rocm/libexec/rocm_smi"
517 if os.path.exists(rocm_smi_path):
518 if rocm_smi_path not in sys.path:
519 sys.path.append(rocm_smi_path)
520 try:
521 import rocm_smi # noqa: PLC0415
522
523 self.rocm_smi = rocm_smi
524 self.rocm_smi.initializeRsmi()
525 except ImportError:
526 self.rocm_smi = None
527
528 def get_used_memory(self, dev):
529 if self.rocm_smi is None:
530 return -1
531 return self.rocm_smi.getMemInfo(dev, "VRAM")[0] / 1024 / 1024
532
533 def measure_gpu_usage(self):
534 if self.rocm_smi is None:
535 return None
536
537 device_count = len(self.rocm_smi.listDevices()) if self.rocm_smi is not None else 0
538 max_gpu_usage = [0 for i in range(device_count)]
539 gpu_name = [f"GPU{i}" for i in range(device_count)]
540 while True:
541 for i in range(device_count):
542 max_gpu_usage[i] = max(max_gpu_usage[i], self.get_used_memory(i))
543 time.sleep(0.005) # 5ms
544 if not self.keep_measuring:
545 break
546 return [
547 {
548 "device_id": i,
549 "name": gpu_name[i],
550 "max_used_MB": max_gpu_usage[i],
551 }
552 for i in range(device_count)
553 ]
554
555
556def measure_memory(is_gpu, func, monitor_type="cuda", start_memory=None):
557 memory_monitor_type = None
558 if monitor_type == "rocm":
559 memory_monitor_type = RocmMemoryMonitor
560 else:
561 memory_monitor_type = CudaMemoryMonitor
562
563 monitor = memory_monitor_type(False)
564
565 if is_gpu:
566 if start_memory is not None:
567 memory_before_test = start_memory
568 else:
569 memory_before_test = monitor.measure_gpu_usage()
570 if memory_before_test is None:
571 return None
572
573 if func is None:
574 return memory_before_test
575
576 with ThreadPoolExecutor() as executor:
577 monitor = memory_monitor_type()
578 mem_thread = executor.submit(monitor.measure_gpu_usage)
579 try:
580 fn_thread = executor.submit(func)
581 _ = fn_thread.result()
582 finally:
583 monitor.keep_measuring = False
584 max_usage = mem_thread.result()
585
586 if max_usage is None:
587 return None
588
589 logger.info(f"GPU memory usage: before={memory_before_test} peak={max_usage}")
590 if len(memory_before_test) >= 1 and len(max_usage) >= 1 and len(memory_before_test) == len(max_usage):
591 # When there are multiple GPUs, we will check the one with maximum usage.
592 max_used = 0
593 for i, memory_before in enumerate(memory_before_test):
594 before = memory_before["max_used_MB"]
595 after = max_usage[i]["max_used_MB"]
596 used = after - before
597 max_used = max(max_used, used)
598 return max_used
599 return None
600
601 # CPU memory
602 if start_memory is not None:
603 memory_before_test = start_memory
604 else:
605 memory_before_test = monitor.measure_cpu_usage()
606
607 if func is None:
608 return memory_before_test
609
610 with ThreadPoolExecutor() as executor:
611 monitor = memory_monitor_type()
612 mem_thread = executor.submit(monitor.measure_cpu_usage)
613 try:
614 fn_thread = executor.submit(func)
615 _ = fn_thread.result()
616 finally:
617 monitor.keep_measuring = False
618 max_usage = mem_thread.result()
619
620 logger.info(f"CPU memory usage: before={memory_before_test:.1f} MB, peak={max_usage:.1f} MB")
621 return max_usage - memory_before_test
622
623
624def get_ort_environment_variables():
625 # Environment variables might impact ORT performance on transformer models. Note that they are for testing only.
626 env_names = [
627 "ORT_DISABLE_FUSED_ATTENTION",
628 "ORT_ENABLE_FUSED_CAUSAL_ATTENTION",
629 "ORT_DISABLE_FUSED_CROSS_ATTENTION",
630 "ORT_DISABLE_TRT_FLASH_ATTENTION",
631 "ORT_DISABLE_MEMORY_EFFICIENT_ATTENTION",
632 "ORT_TRANSFORMER_OPTIONS",
633 "ORT_CUDA_GEMM_OPTIONS",
634 ]
635 env = ""
636 for name in env_names:
637 value = os.getenv(name)
638 if value is None:
639 continue
640 if env:
641 env += ","
642 env += f"{name}={value}"
643 return env
644