Signed-off-by: Fengyuan Yu <15fengyuan@gmail.com> Co-authored-by: Fengyuan Yu <15fengyuan@gmail.com> Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
co-authored by
Fengyuan Yu
ronnie_zheng
parent
e39f0f4ff3
commit
5c245d978f
@@ -29,7 +29,7 @@ import dataclasses
|
||||
import json
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Tuple
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
@@ -81,6 +81,8 @@ class BenchArgs:
|
||||
task_name: str = "unknown"
|
||||
num_prompts: int = 10
|
||||
batch_size: int = 1
|
||||
random_request_config: str = None
|
||||
random_request_seed: int = 42
|
||||
|
||||
# Benchmark Execution
|
||||
skip_warmup: bool = False
|
||||
@@ -151,6 +153,23 @@ class BenchArgs:
|
||||
help="Batch size per generation call (currently only bs=1 is supported)",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--random-request-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"JSON string defining random request profiles. "
|
||||
"Each profile may contain: width, height, num_inference_steps, etc. "
|
||||
"The 'weight' field controls sampling probability (relative weight)."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--random-request-seed",
|
||||
type=int,
|
||||
default=42,
|
||||
help="Random seed for sampling request profiles (default: 42).",
|
||||
)
|
||||
|
||||
# Benchmark Execution
|
||||
parser.add_argument(
|
||||
"--skip-warmup", action="store_true", help="Skip warmup batch"
|
||||
@@ -186,17 +205,22 @@ def generate_batch(
|
||||
engine: DiffGenerator,
|
||||
bench_args: BenchArgs,
|
||||
prompts: List[str],
|
||||
user_sampling_params: Dict[str, Any],
|
||||
user_sampling_params: List[Dict[str, Any]],
|
||||
) -> BatchOutput:
|
||||
"""Generate batch of images/videos synchronously."""
|
||||
assert len(user_sampling_params) == len(prompts), (
|
||||
f"user_sampling_params length ({len(user_sampling_params)}) must match "
|
||||
f"prompts length ({len(prompts)})"
|
||||
)
|
||||
|
||||
output = BatchOutput()
|
||||
start_time = time.perf_counter()
|
||||
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
|
||||
for prompt in prompts:
|
||||
for prompt, params in zip(prompts, user_sampling_params):
|
||||
try:
|
||||
sampling_params_kwargs = dict(user_sampling_params)
|
||||
sampling_params_kwargs = dict(params)
|
||||
sampling_params_kwargs["prompt"] = prompt
|
||||
result = engine.generate(sampling_params_kwargs=sampling_params_kwargs)
|
||||
|
||||
@@ -227,6 +251,7 @@ def calculate_metrics(
|
||||
total_duration: float,
|
||||
resolution: Tuple[int, int, int],
|
||||
num_requests: int,
|
||||
all_sampling_params: Optional[List[Dict[str, Any]]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Calculate generation-specific throughput metrics."""
|
||||
successful = [o for o in outputs if o.success]
|
||||
@@ -235,8 +260,15 @@ def calculate_metrics(
|
||||
peak_memory = max((o.peak_memory_mb for o in outputs), default=0)
|
||||
|
||||
width, height, frames = resolution
|
||||
pixels_per_sample = width * height * frames
|
||||
total_pixels = num_success * pixels_per_sample
|
||||
if all_sampling_params:
|
||||
total_pixels = sum(
|
||||
p.get("width", width)
|
||||
* p.get("height", height)
|
||||
* p.get("num_frames", frames)
|
||||
for p in all_sampling_params[:num_success]
|
||||
)
|
||||
else:
|
||||
total_pixels = num_success * width * height * frames
|
||||
|
||||
metrics = {
|
||||
"num_requests": num_requests,
|
||||
@@ -272,6 +304,11 @@ def throughput_test(
|
||||
|
||||
engine = initialize_engine(server_args)
|
||||
|
||||
if bench_args.random_request_config and bench_args.dataset != "random":
|
||||
raise ValueError(
|
||||
"--random-request-config can only be used with --dataset random"
|
||||
)
|
||||
|
||||
logger.info(f"Loading {bench_args.dataset} dataset...")
|
||||
if bench_args.dataset == "vbench":
|
||||
bench_args.task_name = engine.server_args.pipeline_config.task_type
|
||||
@@ -281,7 +318,7 @@ def throughput_test(
|
||||
else:
|
||||
raise ValueError(f"Unknown dataset: {bench_args.dataset}")
|
||||
|
||||
sampling_params = {
|
||||
_sampling_params = {
|
||||
"guidance_scale": bench_args.guidance_scale,
|
||||
"num_inference_steps": bench_args.num_inference_steps,
|
||||
"height": bench_args.height,
|
||||
@@ -290,18 +327,29 @@ def throughput_test(
|
||||
"seed": bench_args.seed,
|
||||
}
|
||||
if bench_args.disable_safety_checker:
|
||||
sampling_params["safety_checker"] = None
|
||||
_sampling_params["safety_checker"] = None
|
||||
|
||||
total_count = min(bench_args.num_prompts, len(dataset))
|
||||
all_prompts = [dataset[i].prompt for i in range(total_count)]
|
||||
|
||||
if bench_args.random_request_config:
|
||||
all_sampling_params = []
|
||||
for i in range(total_count):
|
||||
params = dict(_sampling_params)
|
||||
params.update(dataset.get_sampling_params(i))
|
||||
all_sampling_params.append(params)
|
||||
else:
|
||||
all_sampling_params = [_sampling_params] * total_count
|
||||
|
||||
if not bench_args.skip_warmup:
|
||||
logger.info("Running warmup batch...")
|
||||
warmup_count = min(bench_args.batch_size, len(dataset))
|
||||
warmup_prompts = [dataset[i].prompt for i in range(warmup_count)]
|
||||
generate_batch(engine, bench_args, warmup_prompts, sampling_params)
|
||||
warmup_count = min(bench_args.batch_size, total_count)
|
||||
warmup_prompts = all_prompts[:warmup_count]
|
||||
warmup_sampling_params = all_sampling_params[:warmup_count]
|
||||
generate_batch(engine, bench_args, warmup_prompts, warmup_sampling_params)
|
||||
|
||||
logger.info(f"Running benchmark with {bench_args.num_prompts} prompts...")
|
||||
outputs: List[BatchOutput] = []
|
||||
total_count = min(bench_args.num_prompts, len(dataset))
|
||||
all_prompts = [dataset[i].prompt for i in range(total_count)]
|
||||
|
||||
start_time = time.perf_counter()
|
||||
|
||||
@@ -315,9 +363,10 @@ def throughput_test(
|
||||
for batch_start in range(0, total_count, bench_args.batch_size):
|
||||
batch_end = min(batch_start + bench_args.batch_size, total_count)
|
||||
batch_prompts = all_prompts[batch_start:batch_end]
|
||||
batch_sampling_params = all_sampling_params[batch_start:batch_end]
|
||||
|
||||
batch_output = generate_batch(
|
||||
engine, bench_args, batch_prompts, sampling_params
|
||||
engine, bench_args, batch_prompts, batch_sampling_params
|
||||
)
|
||||
outputs.append(batch_output)
|
||||
|
||||
@@ -332,6 +381,7 @@ def throughput_test(
|
||||
total_duration,
|
||||
resolution=resolution,
|
||||
num_requests=total_count,
|
||||
all_sampling_params=all_sampling_params,
|
||||
)
|
||||
|
||||
display_results(
|
||||
|
||||
@@ -191,6 +191,8 @@ async def async_request_image_sglang(
|
||||
|
||||
if input.width and input.height:
|
||||
payload["size"] = f"{input.width}x{input.height}"
|
||||
if input.num_inference_steps:
|
||||
payload["num_inference_steps"] = input.num_inference_steps
|
||||
|
||||
# Merge extra parameters
|
||||
payload.update(input.extra_body)
|
||||
@@ -299,6 +301,8 @@ async def async_request_video_sglang(
|
||||
payload["size"] = f"{input.width}x{input.height}"
|
||||
if input.num_frames:
|
||||
payload["num_frames"] = input.num_frames
|
||||
if input.num_inference_steps:
|
||||
payload["num_inference_steps"] = input.num_inference_steps
|
||||
if input.fps:
|
||||
payload["fps"] = input.fps
|
||||
|
||||
@@ -519,6 +523,11 @@ async def benchmark(args):
|
||||
|
||||
setattr(args, "task_name", task_name)
|
||||
|
||||
if args.random_request_config and args.dataset != "random":
|
||||
raise ValueError(
|
||||
"--random-request-config can only be used with --dataset random"
|
||||
)
|
||||
|
||||
if args.dataset == "vbench":
|
||||
dataset = VBenchDataset(args, api_url, args.model)
|
||||
elif args.dataset == "random":
|
||||
@@ -720,6 +729,25 @@ if __name__ == "__main__":
|
||||
)
|
||||
parser.add_argument("--width", type=int, default=None, help="Image/Video width.")
|
||||
parser.add_argument("--height", type=int, default=None, help="Image/Video height.")
|
||||
parser.add_argument(
|
||||
"--random-request-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"JSON string defining random request profiles. "
|
||||
"Each profile may contain: width, height, num_inference_steps, etc. "
|
||||
"The 'weight' field controls sampling probability (relative weight). "
|
||||
"Example: "
|
||||
'[{"width":512,"height":512,"num_inference_steps":20,"weight":0.15},'
|
||||
'{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]'
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--random-request-seed",
|
||||
type=int,
|
||||
default=42,
|
||||
help="Random seed for sampling request profiles (default: 42).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-frames", type=int, default=None, help="Number of frames (for video)."
|
||||
)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
import subprocess
|
||||
import uuid
|
||||
@@ -286,16 +287,39 @@ class RandomDataset(BaseDataset):
|
||||
super().__init__(args, api_url, model)
|
||||
self.num_prompts = args.num_prompts or 100
|
||||
|
||||
self.random_request_config = args.random_request_config
|
||||
if self.random_request_config:
|
||||
self.random_request_config = json.loads(self.random_request_config)
|
||||
weights = [p.pop("weight") for p in self.random_request_config]
|
||||
seed = args.random_request_seed
|
||||
rng = random.Random(seed)
|
||||
self._sampled_requests = rng.choices(
|
||||
self.random_request_config, weights=weights, k=self.num_prompts
|
||||
)
|
||||
else:
|
||||
self._sampled_requests = None
|
||||
|
||||
def get_sampling_params(self, idx: int) -> dict:
|
||||
"""Return the per-request sampling profile dict, or empty dict if not mix-diffusion."""
|
||||
if self._sampled_requests:
|
||||
return self._sampled_requests[idx]
|
||||
return {}
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self.num_prompts
|
||||
|
||||
def __getitem__(self, idx: int) -> RequestFuncInput:
|
||||
profile = self._sampled_requests[idx] if self._sampled_requests else {}
|
||||
|
||||
return RequestFuncInput(
|
||||
prompt=f"Random prompt {idx} for benchmarking diffusion models",
|
||||
api_url=self.api_url,
|
||||
model=self.model,
|
||||
width=self.args.width,
|
||||
height=self.args.height,
|
||||
num_frames=self.args.num_frames,
|
||||
fps=self.args.fps,
|
||||
width=profile.get("width", self.args.width),
|
||||
height=profile.get("height", self.args.height),
|
||||
num_frames=profile.get("num_frames", self.args.num_frames),
|
||||
num_inference_steps=profile.get(
|
||||
"num_inference_steps", self.args.num_inference_steps
|
||||
),
|
||||
fps=profile.get("fps", self.args.fps),
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user