From 0d12c41fc86f88c65aa09e575f2c22eccc522158 Mon Sep 17 00:00:00 2001 From: Yuzhou Nie <62874089+rucnyz@users.noreply.github.com> Date: Thu, 6 Feb 2025 19:30:23 -0800 Subject: [PATCH] add benchmark for ops (#176) --- fastvideo/ops/modulate/k_modulate.py | 102 ++++++++ fastvideo/ops/modulate/modulate.py | 63 +++++ fastvideo/sample/sample_t2v_cogvideo.py | 306 +++++++++++++++++++++++ tests/test_kernels/benchmark.py | 297 ++++++++++++++++++++++ tests/test_kernels/benchmark_modulate.py | 83 ++++++ 5 files changed, 851 insertions(+) create mode 100644 fastvideo/ops/modulate/k_modulate.py create mode 100644 fastvideo/ops/modulate/modulate.py create mode 100644 fastvideo/sample/sample_t2v_cogvideo.py create mode 100644 tests/test_kernels/benchmark.py create mode 100644 tests/test_kernels/benchmark_modulate.py diff --git a/fastvideo/ops/modulate/k_modulate.py b/fastvideo/ops/modulate/k_modulate.py new file mode 100644 index 000000000..5bcc890fc --- /dev/null +++ b/fastvideo/ops/modulate/k_modulate.py @@ -0,0 +1,102 @@ +import triton +import triton.language as tl + +CONFIG_LIST = [ + triton.Config({"BLOCK_M": 256, "BLOCK_N": 32}, num_stages=2, num_warps=4), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_stages=2, num_warps=4), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 32}, num_stages=2, num_warps=4), + triton.Config({"BLOCK_M": 64, "BLOCK_N": 128}, num_stages=2, num_warps=4), + triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}, num_stages=2, num_warps=4), + triton.Config({"BLOCK_M": 64, "BLOCK_N": 32}, num_stages=2, num_warps=4), + triton.Config({"BLOCK_M": 32, "BLOCK_N": 64}, num_stages=2, num_warps=4), + triton.Config({"BLOCK_M": 32, "BLOCK_N": 128}, num_stages=2, num_warps=4), + triton.Config({"BLOCK_M": 32, "BLOCK_N": 256}, num_stages=2, num_warps=4), +] + + +@triton.autotune( + configs=CONFIG_LIST, + key=["M", "N"], +) +@triton.jit +def _modulate_fwd( + x_ptr, # *Pointer* to first input vector. + output_ptr, # *Pointer* to output vector. + scale_ptr, + shift_ptr, + m_stride, + s_stride, + M, + N, + seq_len, + BLOCK_M: tl.constexpr, # Number of elements each program should process. + BLOCK_N: tl.constexpr, + # NOTE: `constexpr` so it can be used as a shape value. +): + row_id = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. + rows = row_id * BLOCK_M + tl.arange(0, BLOCK_M) + s_rows = (row_id // seq_len) * BLOCK_M + col_id = tl.program_id(axis=1) + cols = col_id * BLOCK_N + tl.arange(0, BLOCK_N) + + x_ptrs = x_ptr + rows[:, None] * m_stride + cols[None, :] + scale_ptrs = scale_ptr + s_rows * s_stride + cols[None, :] + shift_ptrs = shift_ptr + s_rows * s_stride + cols[None, :] + + col_mask = cols[None, :] < N + block_mask = (rows[:, None] < M) & col_mask + s_block_mask = col_mask + x = tl.load(x_ptrs, mask=block_mask, other=0.0) + scale = tl.load(scale_ptrs, mask=s_block_mask, other=0.0) + shift = tl.load(shift_ptrs, mask=s_block_mask, other=0.0) + + output = x * (1 + scale) + shift + # Write x + y back to DRAM. + tl.store(output_ptr + rows[:, None] * m_stride + cols[None, :], output, mask=block_mask) + + +@triton.autotune( + configs=CONFIG_LIST, + key=["M", "N"], +) +@triton.jit +def _modulate_bwd( + dx_ptr, # *Pointer* to first input vector. + x_ptr, + dy_ptr, # *Pointer* to output vector. + scale_ptr, + dscale_ptr, + m_stride, + s_stride, + M, + N, + seq_len, + BLOCK_M: tl.constexpr, # Number of elements each program should process. + BLOCK_N: tl.constexpr, + # NOTE: `constexpr` so it can be used as a shape value. +): + row_id = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. + rows = row_id * BLOCK_M + tl.arange(0, BLOCK_M) + s_rows = (row_id // seq_len) * BLOCK_M + col_id = tl.program_id(axis=1) + cols = col_id * BLOCK_N + tl.arange(0, BLOCK_N) + + x_ptrs = x_ptr + rows[:, None] * m_stride + cols[None, :] + dy_ptrs = dy_ptr + rows[:, None] * m_stride + cols[None, :] + dx_ptrs = dx_ptr + rows[:, None] * m_stride + cols[None, :] + dscale_ptrs = dscale_ptr + rows[:, None] * m_stride + cols[None, :] + + scale_ptrs = scale_ptr + s_rows * s_stride + cols[None, :] + + col_mask = cols[None, :] < N + block_mask = (rows[:, None] < M) & col_mask + s_block_mask = col_mask + x = tl.load(x_ptrs, mask=block_mask, other=0.0) + dy = tl.load(dy_ptrs, mask=block_mask, other=0.0) + scale = tl.load(scale_ptrs, mask=s_block_mask, other=0.0) + + dx = dy * (1 + scale) + dscale = dy * x + # Write x + y back to DRAM. + tl.store(dx_ptrs, dx, mask=block_mask) + tl.store(dscale_ptrs, dscale, mask=block_mask) \ No newline at end of file diff --git a/fastvideo/ops/modulate/modulate.py b/fastvideo/ops/modulate/modulate.py new file mode 100644 index 000000000..dc6314eda --- /dev/null +++ b/fastvideo/ops/modulate/modulate.py @@ -0,0 +1,63 @@ +import torch +import triton + +from fastvideo.ops.modulate.k_modulate import _modulate_fwd, _modulate_bwd + + +class _FusedModulate(torch.autograd.Function): + @staticmethod + def forward(ctx, x, scale, shift): + y = torch.empty_like(x) + batch, seq_len, dim = x.shape + M = batch * seq_len + N = dim + x = x.view(-1, dim).contiguous() + scale = scale.view(-1, dim).contiguous() + shift = shift.view(-1, dim).contiguous() + + def grid(meta): + return ( + triton.cdiv(batch * seq_len, meta["BLOCK_M"]), + triton.cdiv(dim, meta["BLOCK_N"]), + ) + + _modulate_fwd[grid](x, y, scale, shift, x.stride(0), scale.stride(0), M, N, seq_len) + + ctx.save_for_backward(x, scale) + ctx.batch = batch + ctx.seq_len = seq_len + ctx.dim = dim + return y + + @staticmethod + def backward(ctx, dy): # pragma: no cover # this is covered, but called directly from C++ + x, scale = ctx.saved_tensors + + batch, seq_len, dim = ctx.batch, ctx.seq_len, ctx.dim + M = batch * seq_len + N = dim + + # allocate output + dy = dy.contiguous() + dx = torch.empty_like(dy) + dscale = torch.empty_like(dy) + dshift = torch.sum(dy, dim=1) + + def grid(meta): + return ( + triton.cdiv(batch * seq_len, meta["BLOCK_M"]), + triton.cdiv(dim, meta["BLOCK_N"]), + ) + + _modulate_bwd[grid](dx, x, dy, scale, dscale, x.stride(0), scale.stride(0), M, N, seq_len) + + dscale = torch.sum(dscale, dim=1) + return dx, dscale, dshift + + +def fused_modulate( + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, +) -> torch.Tensor: + return _FusedModulate.apply(x, scale, shift) \ No newline at end of file diff --git a/fastvideo/sample/sample_t2v_cogvideo.py b/fastvideo/sample/sample_t2v_cogvideo.py new file mode 100644 index 000000000..b6f40df7d --- /dev/null +++ b/fastvideo/sample/sample_t2v_cogvideo.py @@ -0,0 +1,306 @@ +""" +This script demonstrates how to generate a video using the CogVideoX model with the Hugging Face `diffusers` pipeline. +The script supports different types of video generation, including text-to-video (t2v), image-to-video (i2v), +and video-to-video (v2v), depending on the input data and different weight. + +- text-to-video: THUDM/CogVideoX-5b, THUDM/CogVideoX-2b or THUDM/CogVideoX1.5-5b +- video-to-video: THUDM/CogVideoX-5b, THUDM/CogVideoX-2b or THUDM/CogVideoX1.5-5b +- image-to-video: THUDM/CogVideoX-5b-I2V or THUDM/CogVideoX1.5-5b-I2V + +Running the Script: +To run the script, use the following command with appropriate arguments: + +```bash +$ python cli_demo.py --prompt "A girl riding a bike." --model_path THUDM/CogVideoX1.5-5b --generate_type "t2v" +``` + +Additional options are available to specify the model path, guidance scale, number of inference steps, video generation type, and output paths. +""" + +import argparse +import logging +import os +from typing import Literal, Optional + +import torch +from diffusers import ( + CogVideoXDPMScheduler, + CogVideoXImageToVideoPipeline, + CogVideoXPipeline, + CogVideoXVideoToVideoPipeline, +) +from diffusers.utils import export_to_video, load_image, load_video + + +logging.basicConfig(level=logging.INFO) + +# Recommended resolution for each model (width, height) +RESOLUTION_MAP = { + # cogvideox1.5-* + "cogvideox1.5-5b-i2v": (1360, 768), + "cogvideox1.5-5b": (1360, 768), + # cogvideox-* + "cogvideox-5b-i2v": (720, 480), + "cogvideox-5b": (720, 480), + "cogvideox-2b": (720, 480), +} + + +def generate_video( + prompt: str, + model_path: str, + lora_path: str = None, + lora_rank: int = 128, + num_frames: int = 81, + width: Optional[int] = None, + height: Optional[int] = None, + output_path: str = "./output.mp4", + image_or_video_path: str = "", + num_inference_steps: int = 50, + guidance_scale: float = 6.0, + num_videos_per_prompt: int = 1, + dtype: torch.dtype = torch.bfloat16, + generate_type: str = Literal[ + "t2v", "i2v", "v2v" + ], # i2v: image to video, v2v: video to video + seed: int = 42, + fps: int = 16, +): + """ + Generates a video based on the given prompt and saves it to the specified path. + + Parameters: + - prompt (str): The description of the video to be generated. + - model_path (str): The path of the pre-trained model to be used. + - lora_path (str): The path of the LoRA weights to be used. + - lora_rank (int): The rank of the LoRA weights. + - output_path (str): The path where the generated video will be saved. + - num_inference_steps (int): Number of steps for the inference process. More steps can result in better quality. + - num_frames (int): Number of frames to generate. CogVideoX1.0 generates 49 frames for 6 seconds at 8 fps, while CogVideoX1.5 produces either 81 or 161 frames, corresponding to 5 seconds or 10 seconds at 16 fps. + - width (int): The width of the generated video, applicable only for CogVideoX1.5-5B-I2V + - height (int): The height of the generated video, applicable only for CogVideoX1.5-5B-I2V + - guidance_scale (float): The scale for classifier-free guidance. Higher values can lead to better alignment with the prompt. + - num_videos_per_prompt (int): Number of videos to generate per prompt. + - dtype (torch.dtype): The data type for computation (default is torch.bfloat16). + - generate_type (str): The type of video generation (e.g., 't2v', 'i2v', 'v2v').ยท + - seed (int): The seed for reproducibility. + - fps (int): The frames per second for the generated video. + """ + + # 1. Load the pre-trained CogVideoX pipeline with the specified precision (bfloat16). + # add device_map="balanced" in the from_pretrained function and remove the enable_model_cpu_offload() + # function to use Multi GPUs. + + image = None + video = None + + model_name = model_path.split("/")[-1].lower() + desired_resolution = RESOLUTION_MAP[model_name] + if width is None or height is None: + width, height = desired_resolution + logging.info( + f"\033[1mUsing default resolution {desired_resolution} for {model_name}\033[0m" + ) + elif (width, height) != desired_resolution: + if generate_type == "i2v": + # For i2v models, use user-defined width and height + logging.warning( + f"\033[1;31mThe width({width}) and height({height}) are not recommended for {model_name}. The best resolution is {desired_resolution}.\033[0m" + ) + else: + # Otherwise, use the recommended width and height + logging.warning( + f"\033[1;31m{model_name} is not supported for custom resolution. Setting back to default resolution {desired_resolution}.\033[0m" + ) + width, height = desired_resolution + + if generate_type == "i2v": + pipe = CogVideoXImageToVideoPipeline.from_pretrained( + model_path, torch_dtype=dtype + ) + image = load_image(image=image_or_video_path) + elif generate_type == "t2v": + pipe = CogVideoXPipeline.from_pretrained(model_path, torch_dtype=dtype) + else: + pipe = CogVideoXVideoToVideoPipeline.from_pretrained( + model_path, torch_dtype=dtype + ) + video = load_video(image_or_video_path) + + # If you're using with lora, add this code + if lora_path: + pipe.load_lora_weights( + lora_path, + weight_name="pytorch_lora_weights.safetensors", + adapter_name="test_1", + ) + pipe.fuse_lora(lora_scale=1 / lora_rank) + + # 2. Set Scheduler. + # Can be changed to `CogVideoXDPMScheduler` or `CogVideoXDDIMScheduler`. + # We recommend using `CogVideoXDDIMScheduler` for CogVideoX-2B. + # using `CogVideoXDPMScheduler` for CogVideoX-5B / CogVideoX-5B-I2V. + + # pipe.scheduler = CogVideoXDDIMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing") + pipe.scheduler = CogVideoXDPMScheduler.from_config( + pipe.scheduler.config, timestep_spacing="trailing" + ) + + # 3. Enable CPU offload for the model. + # turn off if you have multiple GPUs or enough GPU memory(such as H100) and it will cost less time in inference + # and enable to("cuda") + + # pipe.to("cuda") + pipe.enable_sequential_cpu_offload() + pipe.vae.enable_slicing() + pipe.vae.enable_tiling() + + # 4. Generate the video frames based on the prompt. + # `num_frames` is the Number of frames to generate. + if generate_type == "i2v": + video_generate = pipe( + height=height, + width=width, + prompt=prompt, + image=image, + # The path of the image, the resolution of video will be the same as the image for CogVideoX1.5-5B-I2V, otherwise it will be 720 * 480 + num_videos_per_prompt=num_videos_per_prompt, # Number of videos to generate per prompt + num_inference_steps=num_inference_steps, # Number of inference steps + num_frames=num_frames, # Number of frames to generate + use_dynamic_cfg=True, # This id used for DPM scheduler, for DDIM scheduler, it should be False + guidance_scale=guidance_scale, + generator=torch.Generator().manual_seed( + seed + ), # Set the seed for reproducibility + ).frames[0] + elif generate_type == "t2v": + video_generate = pipe( + height=height, + width=width, + prompt=prompt, + num_videos_per_prompt=num_videos_per_prompt, + num_inference_steps=num_inference_steps, + num_frames=num_frames, + use_dynamic_cfg=True, + guidance_scale=guidance_scale, + generator=torch.Generator().manual_seed(seed), + ).frames[0] + else: + video_generate = pipe( + height=height, + width=width, + prompt=prompt, + video=video, # The path of the video to be used as the background of the video + num_videos_per_prompt=num_videos_per_prompt, + num_inference_steps=num_inference_steps, + num_frames=num_frames, + use_dynamic_cfg=True, + guidance_scale=guidance_scale, + generator=torch.Generator().manual_seed( + seed + ), # Set the seed for reproducibility + ).frames[0] + export_to_video(video_generate, output_path, fps=fps) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Generate a video from a text prompt using CogVideoX" + ) + parser.add_argument( + "--prompt", + type=str, + required=True, + help="The description of the video to be generated", + ) + parser.add_argument( + "--image_or_video_path", + type=str, + default=None, + help="The path of the image to be used as the background of the video", + ) + parser.add_argument( + "--model_path", + type=str, + default="THUDM/CogVideoX1.5-5B", + help="Path of the pre-trained model use", + ) + parser.add_argument( + "--lora_path", + type=str, + default=None, + help="The path of the LoRA weights to be used", + ) + parser.add_argument( + "--lora_rank", type=int, default=128, help="The rank of the LoRA weights" + ) + parser.add_argument( + "--output_path", + type=str, + default="./output.mp4", + help="The path save generated video", + ) + parser.add_argument( + "--guidance_scale", + type=float, + default=6.0, + help="The scale for classifier-free guidance", + ) + parser.add_argument( + "--num_inference_steps", type=int, default=50, help="Inference steps" + ) + parser.add_argument( + "--num_frames", + type=int, + default=81, + help="Number of steps for the inference process", + ) + parser.add_argument( + "--width", type=int, default=None, help="The width of the generated video" + ) + parser.add_argument( + "--height", type=int, default=None, help="The height of the generated video" + ) + parser.add_argument( + "--fps", + type=int, + default=16, + help="The frames per second for the generated video", + ) + parser.add_argument( + "--num_videos_per_prompt", + type=int, + default=1, + help="Number of videos to generate per prompt", + ) + parser.add_argument( + "--generate_type", type=str, default="t2v", help="The type of video generation" + ) + parser.add_argument( + "--dtype", type=str, default="bfloat16", help="The data type for computation" + ) + parser.add_argument( + "--seed", type=int, default=42, help="The seed for reproducibility" + ) + + args = parser.parse_args() + dtype = torch.float16 if args.dtype == "float16" else torch.bfloat16 + os.makedirs(os.path.dirname(args.output_path), exist_ok=True) + generate_video( + prompt=args.prompt, + model_path=args.model_path, + lora_path=args.lora_path, + lora_rank=args.lora_rank, + output_path=args.output_path, + num_frames=args.num_frames, + width=args.width, + height=args.height, + image_or_video_path=args.image_or_video_path, + num_inference_steps=args.num_inference_steps, + guidance_scale=args.guidance_scale, + num_videos_per_prompt=args.num_videos_per_prompt, + dtype=dtype, + generate_type=args.generate_type, + seed=args.seed, + fps=args.fps, + ) diff --git a/tests/test_kernels/benchmark.py b/tests/test_kernels/benchmark.py new file mode 100644 index 000000000..c931d5241 --- /dev/null +++ b/tests/test_kernels/benchmark.py @@ -0,0 +1,297 @@ +# Copyright (c) 2023, Tri Dao. +"""Useful functions for writing test code.""" + +import torch +import torch.utils.benchmark as benchmark + + +def get_abs_err(x, y): + return (x - y).flatten().abs().max().item() + + +def get_err_ratio(x, y): + err = (x - y).flatten().square().mean().sqrt().item() + base = x.flatten().square().mean().sqrt().item() + return err / base + + +def assert_close(prefix, ref, tri, ratio): + msg = f"{prefix} diff: {get_abs_err(ref, tri):.6f} ratio: {get_err_ratio(ref, tri):.6f}" + print(msg) + assert get_err_ratio(ref, tri) < ratio, msg + + +def benchmark_forward( + fn, + *inputs, + repeats=10, + desc="", + verbose=True, + amp=False, + amp_dtype=torch.float16, + **kwinputs, +): + """Use Pytorch Benchmark on the forward pass of an arbitrary function.""" + if verbose: + print(desc, "- Forward pass") + + def amp_wrapper(*inputs, **kwinputs): + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + fn(*inputs, **kwinputs) + + t = benchmark.Timer( + stmt="fn_amp(*inputs, **kwinputs)", + globals={"fn_amp": amp_wrapper, "inputs": inputs, "kwinputs": kwinputs}, + num_threads=torch.get_num_threads(), + ) + m = t.timeit(repeats) + if verbose: + print(m) + return t, m + + +def benchmark_backward( + fn, + *inputs, + grad=None, + repeats=10, + desc="", + verbose=True, + amp=False, + amp_dtype=torch.float16, + **kwinputs, +): + """Use Pytorch Benchmark on the backward pass of an arbitrary function.""" + if verbose: + print(desc, "- Backward pass") + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + y = fn(*inputs, **kwinputs) + if type(y) is tuple: + y = y[0] + if grad is None: + grad = torch.randn_like(y) + else: + if grad.shape != y.shape: + raise RuntimeError("Grad shape does not match output shape") + + def f(*inputs, y, grad): + # Set .grad to None to avoid extra operation of gradient accumulation + for x in inputs: + if isinstance(x, torch.Tensor): + x.grad = None + y.backward(grad, retain_graph=True) + + t = benchmark.Timer( + stmt="f(*inputs, y=y, grad=grad)", + globals={"f": f, "inputs": inputs, "y": y, "grad": grad}, + num_threads=torch.get_num_threads(), + ) + m = t.timeit(repeats) + if verbose: + print(m) + return t, m + + +def benchmark_combined( + fn, + *inputs, + grad=None, + repeats=10, + desc="", + verbose=True, + amp=False, + amp_dtype=torch.float16, + **kwinputs, +): + """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function.""" + if verbose: + print(desc, "- Forward + Backward pass") + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + y = fn(*inputs, **kwinputs) + if type(y) is tuple: + y = y[0] + if grad is None: + grad = torch.randn_like(y) + else: + if grad.shape != y.shape: + raise RuntimeError("Grad shape does not match output shape") + + def f(grad, *inputs, **kwinputs): + for x in inputs: + if isinstance(x, torch.Tensor): + x.grad = None + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + y = fn(*inputs, **kwinputs) + if type(y) is tuple: + y = y[0] + y.backward(grad, retain_graph=True) + + t = benchmark.Timer( + stmt="f(grad, *inputs, **kwinputs)", + globals={ + "f": f, + "fn": fn, + "inputs": inputs, + "grad": grad, + "kwinputs": kwinputs, + }, + num_threads=torch.get_num_threads(), + ) + m = t.timeit(repeats) + if verbose: + print(m) + return t, m + + +def benchmark_fwd_bwd( + fn, + *inputs, + grad=None, + repeats=10, + desc="", + verbose=True, + amp=False, + amp_dtype=torch.float16, + **kwinputs, +): + """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function.""" + return ( + benchmark_forward( + fn, + *inputs, + repeats=repeats, + desc=desc, + verbose=verbose, + amp=amp, + amp_dtype=amp_dtype, + **kwinputs, + ), + benchmark_backward( + fn, + *inputs, + grad=grad, + repeats=repeats, + desc=desc, + verbose=verbose, + amp=amp, + amp_dtype=amp_dtype, + **kwinputs, + ), + ) + + +def benchmark_all( + fn, + *inputs, + grad=None, + repeats=10, + desc="", + verbose=True, + amp=False, + amp_dtype=torch.float16, + **kwinputs, +): + """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function.""" + return ( + benchmark_forward( + fn, + *inputs, + repeats=repeats, + desc=desc, + verbose=verbose, + amp=amp, + amp_dtype=amp_dtype, + **kwinputs, + ), + benchmark_backward( + fn, + *inputs, + grad=grad, + repeats=repeats, + desc=desc, + verbose=verbose, + amp=amp, + amp_dtype=amp_dtype, + **kwinputs, + ), + benchmark_combined( + fn, + *inputs, + grad=grad, + repeats=repeats, + desc=desc, + verbose=verbose, + amp=amp, + amp_dtype=amp_dtype, + **kwinputs, + ), + ) + + +def pytorch_profiler( + fn, + *inputs, + trace_filename=None, + backward=False, + amp=False, + amp_dtype=torch.float16, + cpu=False, + verbose=True, + **kwinputs, +): + """Wrap benchmark functions in Pytorch profiler to see CUDA information.""" + if backward: + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + out = fn(*inputs, **kwinputs) + if type(out) is tuple: + out = out[0] + g = torch.randn_like(out) + for _ in range(30): # Warm up + if backward: + for x in inputs: + if isinstance(x, torch.Tensor): + x.grad = None + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + out = fn(*inputs, **kwinputs) + if type(out) is tuple: + out = out[0] + # Backward should be done outside autocast + if backward: + out.backward(g, retain_graph=True) + activities = ([torch.profiler.ProfilerActivity.CPU] if cpu else []) + [ + torch.profiler.ProfilerActivity.CUDA + ] + with torch.profiler.profile( + activities=activities, + record_shapes=True, + # profile_memory=True, + with_stack=True, + ) as prof: + if backward: + for x in inputs: + if isinstance(x, torch.Tensor): + x.grad = None + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + out = fn(*inputs, **kwinputs) + if type(out) is tuple: + out = out[0] + if backward: + out.backward(g, retain_graph=True) + if verbose: + # print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=50)) + print(prof.key_averages().table(row_limit=50)) + if trace_filename is not None: + prof.export_chrome_trace(trace_filename) + + +def benchmark_memory(fn, *inputs, desc="", verbose=True, **kwinputs): + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + torch.cuda.synchronize() + fn(*inputs, **kwinputs) + torch.cuda.synchronize() + mem = torch.cuda.max_memory_allocated() / ((2**20) * 1000) + if verbose: + print(f"{desc} max memory: {mem}GB") + torch.cuda.empty_cache() + return mem diff --git a/tests/test_kernels/benchmark_modulate.py b/tests/test_kernels/benchmark_modulate.py new file mode 100644 index 000000000..27b0930c1 --- /dev/null +++ b/tests/test_kernels/benchmark_modulate.py @@ -0,0 +1,83 @@ +import torch +from torch import nn + +from benchmark import ( + benchmark_combined, + benchmark_forward, + benchmark_backward, + assert_close, +) +from fastvideo.ops.modulate.modulate import fused_modulate + + +def torch_modulate(x, scale, shift): + return x * (scale.unsqueeze(1) + 1) + shift.unsqueeze(1) + + +# from flash_attn import flash_attn_func + + +def time_fwd(func, *args, **kwargs): + time_fb = benchmark_forward(func, *args, **kwargs) + return time_fb[1].mean + + +def time_fwd_bwd(func, *args, **kwargs): + time_fb = benchmark_combined(func, *args, **kwargs) + return time_fb[1].mean + + +def time_bwd(func, *args, **kwargs): + time_fb = benchmark_backward(func, *args, **kwargs) + return time_fb[1].mean + + +device = "cuda" +dtype = torch.bfloat16 + + +batch_sizes = [1, 8, 32] +seq_lengths = [128, 512, 1024] +hidden_dims = [768, 1024, 2048] + +# methods = (["torch", "triton", "thunderkitten"]) +methods = ["torch", "triton"] +time_f = {} +time_b = {} +time_f_b = {} +speed_f = {} +speed_b = {} +speed_f_b = {} +for B in batch_sizes: + for T in seq_lengths: + for D in hidden_dims: + config = (B, T, D) + + norm_func = nn.LayerNorm(D, elementwise_affine=False, eps=1e-6) + x = torch.randn(B, T, D, device="cuda", requires_grad=True, dtype=dtype) + x = norm_func(x.to(torch.float32)).to(dtype) + shift = torch.randn(B, D, device="cuda", requires_grad=True, dtype=dtype) + scale = torch.randn(B, D, device="cuda", requires_grad=True, dtype=dtype) + # test torch + o_ref = torch_modulate(x, scale, shift) + o_ref.sum().backward(retain_graph=True) + f_b = time_fwd_bwd(torch_modulate, x, scale, shift, verbose=False) + time_f_b[config, "torch"] = f_b + # test triton + o2 = fused_modulate(x, scale, shift) + o2.sum().backward(retain_graph=True) + f_b = time_fwd_bwd(fused_modulate, x, scale, shift, verbose=False) + time_f_b[config, "triton"] = f_b + # test if the results are close + assert_close(" o", o_ref, o2, 0.005) + # time_f_b[config, "thunderkitten"] = f_b + + print(f"### batch size={B}, seq length={T}, B={B}, hidden dim={D} ###") + for method in methods: + # time_f_b[config, method] = time_f[config, method] + time_b[config, method] + print( + f"{method:>50} fwd + bwd:\t {time_f_b[config, method]*1000:>6.4f} ms " + ) + +# with open('flash2_attn_time.plk', 'wb') as fp: +# pickle.dump((speed_f, speed_b, speed_f_b), fp, protocol=pickle.HIGHEST_PROTOCOL)