diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py b/onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py index 236923d5ff..55b143122c 100755 --- a/onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py +++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py @@ -4,7 +4,10 @@ # -------------------------------------------------------------------------- import argparse +import csv import os +import statistics +import sys import time SD_MODELS = { @@ -31,6 +34,94 @@ def example_prompts(): return prompts +def measure_gpu_memory(func, start_memory=None): + class MemoryMonitor: + def __init__(self, keep_measuring=True): + self.keep_measuring = keep_measuring + + def measure_gpu_usage(self): + from py3nvml.py3nvml import ( + NVMLError, + nvmlDeviceGetCount, + nvmlDeviceGetHandleByIndex, + nvmlDeviceGetMemoryInfo, + nvmlDeviceGetName, + nvmlInit, + nvmlShutdown, + ) + + max_gpu_usage = [] + gpu_name = [] + try: + nvmlInit() + device_count = nvmlDeviceGetCount() + if not isinstance(device_count, int): + print(f"nvmlDeviceGetCount result is not integer: {device_count}") + return None + + max_gpu_usage = [0 for i in range(device_count)] + gpu_name = [nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(i)) for i in range(device_count)] + while True: + for i in range(device_count): + info = nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(i)) + if isinstance(info, str): + print(f"nvmlDeviceGetMemoryInfo returns str: {info}") + return None + max_gpu_usage[i] = max(max_gpu_usage[i], info.used / 1024**2) + time.sleep(0.002) # 2ms + if not self.keep_measuring: + break + nvmlShutdown() + return [ + { + "device_id": i, + "name": gpu_name[i], + "max_used_MB": max_gpu_usage[i], + } + for i in range(device_count) + ] + except NVMLError as error: + print("Error fetching GPU information using nvml: %s", error) + return None + + monitor = MemoryMonitor(False) + memory_before_test = monitor.measure_gpu_usage() + + if start_memory is None: + start_memory = memory_before_test + if start_memory is None: + return None + if func is None: + return start_memory + + from concurrent.futures import ThreadPoolExecutor + + with ThreadPoolExecutor() as executor: + monitor = MemoryMonitor() + mem_thread = executor.submit(monitor.measure_gpu_usage) + try: + fn_thread = executor.submit(func) + _ = fn_thread.result() + finally: + monitor.keep_measuring = False + max_usage = mem_thread.result() + + if max_usage is None: + return None + + print(f"GPU memory usage: before={memory_before_test} peak={max_usage}") + if len(start_memory) >= 1 and len(max_usage) >= 1 and len(start_memory) == len(max_usage): + # When there are multiple GPUs, we will check the one with maximum usage. + max_used = 0 + for i, memory_before in enumerate(start_memory): + before = memory_before["max_used_MB"] + after = max_usage[i]["max_used_MB"] + used = after - before + max_used = max(max_used, used) + return max_used + return None + + def get_ort_pipeline(model_name: str, directory: str, provider: str, disable_safety_checker: bool): from diffusers import DPMSolverMultistepScheduler, OnnxStableDiffusionPipeline @@ -61,14 +152,25 @@ def get_ort_pipeline(model_name: str, directory: str, provider: str, disable_saf return pipe -def get_torch_pipeline(model_name: str, disable_safety_checker: bool): +def get_torch_pipeline(model_name: str, disable_safety_checker: bool, enable_torch_compile: bool, use_xformers: bool): from diffusers import DPMSolverMultistepScheduler, StableDiffusionPipeline from torch import channels_last, float16 - pipe = StableDiffusionPipeline.from_pretrained( - model_name, torch_dtype=float16, revision="fp16", use_auth_token=True - ).to("cuda") + pipe = StableDiffusionPipeline.from_pretrained(model_name, torch_dtype=float16).to("cuda") + pipe.unet.to(memory_format=channels_last) # in-place operation + + if use_xformers: + pipe.enable_xformers_memory_efficient_attention() + + if enable_torch_compile: + import torch + + pipe.unet = torch.compile(pipe.unet) + pipe.vae = torch.compile(pipe.vae) + pipe.text_encoder = torch.compile(pipe.text_encoder) + print("Torch compiled unet, vae and text_encoder") + pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config) pipe.set_progress_bar_config(disable=True) @@ -84,14 +186,21 @@ def get_image_filename_prefix(engine: str, model_name: str, batch_size: int, dis return f"{engine}_{short_model_name}_b{batch_size}" + ("" if disable_safety_checker else "_safe") -def run_ort_pipeline(pipe, batch_size: int, image_filename_prefix: str, height, width, steps, num_prompts, batch_count): +def run_ort_pipeline( + pipe, batch_size: int, image_filename_prefix: str, height, width, steps, num_prompts, batch_count, start_memory +): from diffusers import OnnxStableDiffusionPipeline assert isinstance(pipe, OnnxStableDiffusionPipeline) prompts = example_prompts() - pipe("warm up", height, width, num_inference_steps=steps) + def warmup(): + pipe("warm up", height, width, num_inference_steps=steps, num_images_per_prompt=batch_size) + + # Run warm up, and measure GPU memory of two runs (The first run has cuDNN algo search so it might need more memory) + first_run_memory = measure_gpu_memory(warmup, start_memory) + second_run_memory = measure_gpu_memory(warmup, start_memory) latency_list = [] for i, prompt in enumerate(prompts): @@ -111,21 +220,42 @@ def run_ort_pipeline(pipe, batch_size: int, image_filename_prefix: str, height, inference_end = time.time() latency = inference_end - inference_start latency_list.append(latency) - print(f"Inference took {latency} seconds") + print(f"Inference took {latency:.3f} seconds") for k, image in enumerate(images): image.save(f"{image_filename_prefix}_{i}_{j}_{k}.jpg") - print("Average latency in seconds:", sum(latency_list) / len(latency_list)) + from onnxruntime import __version__ as ort_version + + return { + "engine": "onnxruntime", + "version": ort_version, + "height": height, + "width": width, + "steps": steps, + "batch_size": batch_size, + "batch_count": batch_count, + "num_prompts": num_prompts, + "average_latency": sum(latency_list) / len(latency_list), + "median_latency": statistics.median(latency_list), + "first_run_memory_MB": first_run_memory, + "second_run_memory_MB": second_run_memory, + } def run_torch_pipeline( - pipe, batch_size: int, image_filename_prefix: str, height, width, steps, num_prompts, batch_count + pipe, batch_size: int, image_filename_prefix: str, height, width, steps, num_prompts, batch_count, start_memory ): import torch prompts = example_prompts() - pipe("warm up", height, width, num_inference_steps=steps) + # total 2 runs of warm up, and measure GPU memory + def warmup(): + pipe("warm up", height, width, num_inference_steps=steps, num_images_per_prompt=batch_size) + + # Run warm up, and measure GPU memory of two runs (The first run has cuDNN algo search so it might need more memory) + first_run_memory = measure_gpu_memory(warmup, start_memory) + second_run_memory = measure_gpu_memory(warmup, start_memory) torch.set_grad_enabled(False) @@ -151,11 +281,24 @@ def run_torch_pipeline( inference_end = time.time() latency = inference_end - inference_start latency_list.append(latency) - print(f"Inference took {latency} seconds") + print(f"Inference took {latency:.3f} seconds") for k, image in enumerate(images): image.save(f"{image_filename_prefix}_{i}_{j}_{k}.jpg") - print("Average latency in seconds:", sum(latency_list) / len(latency_list)) + return { + "engine": "torch", + "version": torch.__version__, + "height": height, + "width": width, + "steps": steps, + "batch_size": batch_size, + "batch_count": batch_count, + "num_prompts": num_prompts, + "average_latency": sum(latency_list) / len(latency_list), + "median_latency": statistics.median(latency_list), + "first_run_memory_MB": first_run_memory, + "second_run_memory_MB": second_run_memory, + } def run_ort( @@ -169,6 +312,7 @@ def run_ort( steps, num_prompts, batch_count, + start_memory, ): load_start = time.time() pipe = get_ort_pipeline(model_name, directory, provider, disable_safety_checker) @@ -176,18 +320,33 @@ def run_ort( print(f"Model loading took {load_end - load_start} seconds") image_filename_prefix = get_image_filename_prefix("ort", model_name, batch_size, disable_safety_checker) - run_ort_pipeline(pipe, batch_size, image_filename_prefix, height, width, steps, num_prompts, batch_count) + result = run_ort_pipeline( + pipe, batch_size, image_filename_prefix, height, width, steps, num_prompts, batch_count, start_memory + ) + + result.update( + { + "model_name": model_name, + "directory": directory, + "provider": provider, + "disable_safety_checker": disable_safety_checker, + } + ) + return result def run_torch( model_name: str, batch_size: int, disable_safety_checker: bool, + enable_torch_compile: bool, + use_xformers: bool, height, width, steps, num_prompts, batch_count, + start_memory, ): import torch @@ -198,13 +357,31 @@ def run_torch( torch.set_grad_enabled(False) load_start = time.time() - pipe = get_torch_pipeline(model_name, disable_safety_checker) + pipe = get_torch_pipeline(model_name, disable_safety_checker, enable_torch_compile, use_xformers) load_end = time.time() print(f"Model loading took {load_end - load_start} seconds") image_filename_prefix = get_image_filename_prefix("torch", model_name, batch_size, disable_safety_checker) - with torch.inference_mode(): - run_torch_pipeline(pipe, batch_size, image_filename_prefix, height, width, steps, num_prompts, batch_count) + + if not enable_torch_compile: + with torch.inference_mode(): + result = run_torch_pipeline( + pipe, batch_size, image_filename_prefix, height, width, steps, num_prompts, batch_count, start_memory + ) + else: + result = run_torch_pipeline( + pipe, batch_size, image_filename_prefix, height, width, steps, num_prompts, batch_count, start_memory + ) + + result.update( + { + "model_name": model_name, + "directory": None, + "provider": "compile" if enable_torch_compile else "xformers" if use_xformers else "default", + "disable_safety_checker": disable_safety_checker, + } + ) + return result def parse_arguments(): @@ -246,6 +423,22 @@ def parse_arguments(): ) parser.set_defaults(enable_safety_checker=False) + parser.add_argument( + "--enable_torch_compile", + required=False, + action="store_true", + help="Enable compile unet for PyTorch 2.0", + ) + parser.set_defaults(enable_torch_compile=False) + + parser.add_argument( + "--use_xformers", + required=False, + action="store_true", + help="Use xformers for PyTorch", + ) + parser.set_defaults(use_xformers=False) + parser.add_argument( "-b", "--batch_size", @@ -307,19 +500,15 @@ def main(): args = parse_arguments() print(args) + start_memory = measure_gpu_memory(None) + print("GPU memory used before loading models:", start_memory) + sd_model = SD_MODELS[args.version] if args.engine == "onnxruntime": assert args.pipeline, "--pipeline should be specified for onnxruntime engine" - if args.batch_size > 1: - # Need remove a line https://github.com/huggingface/diffusers/blob/a66f2baeb782e091dde4e1e6394e46f169e5ba58/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py#L307 - # in diffuers to run batch_size > 1. - assert ( - not args.enable_safety_checker - ), "batch_size > 1 is not compatible with safety checker due to a bug in diffuers" - provider = "CUDAExecutionProvider" # TODO: use ["CUDAExecutionProvider", "CPUExecutionProvider"] in diffuers - run_ort( + result = run_ort( sd_model, args.pipeline, provider, @@ -330,19 +519,52 @@ def main(): args.steps, args.num_prompts, args.batch_count, + start_memory, ) else: - run_torch( + result = run_torch( sd_model, args.batch_size, not args.enable_safety_checker, + args.enable_torch_compile, + args.use_xformers, args.height, args.width, args.steps, args.num_prompts, args.batch_count, + start_memory, ) + print(result) + + with open("benchmark_result.csv", mode="a", newline="") as csv_file: + column_names = [ + "model_name", + "directory", + "engine", + "version", + "provider", + "disable_safety_checker", + "height", + "width", + "steps", + "batch_size", + "batch_count", + "num_prompts", + "average_latency", + "median_latency", + "first_run_memory_MB", + "second_run_memory_MB", + ] + csv_writer = csv.DictWriter(csv_file, fieldnames=column_names) + csv_writer.writeheader() + csv_writer.writerow(result) + if __name__ == "__main__": - main() + try: + main() + except Exception as e: + tb = sys.exc_info() + print(e.with_traceback(tb[2])) diff --git a/onnxruntime/python/tools/transformers/models/stable_diffusion/requirements.txt b/onnxruntime/python/tools/transformers/models/stable_diffusion/requirements.txt index 45190f2fb9..45ce4ca172 100644 --- a/onnxruntime/python/tools/transformers/models/stable_diffusion/requirements.txt +++ b/onnxruntime/python/tools/transformers/models/stable_diffusion/requirements.txt @@ -1,5 +1,5 @@ # Install the following package in python 3.10 -diffusers==0.12.1 +diffusers==0.13.0 transformers==4.26.0 numpy==1.24.1 accelerate==0.15.0 @@ -10,6 +10,7 @@ packaging==23.0 protobuf==3.20.3 psutil==5.9.4 sympy==1.11.1 +py3nvml==0.2.7 #Tested with PyTorch 1.13.1+cu117 (see pytorch.org for more download options). #--extra-index-url https://download.pytorch.org/whl/cu117 #torch==1.13.1+cu117