mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-11-12 02:06:41 +08:00
Izzy Putterman
2802302bf9
## What does this PR do? **Type of change:** ? <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> **Overview:** Addition of SpecBench Dataset Addition of NVIDID SPEED-Bench dataset, preproc scripts, and custom metrics aggregator Addition of example of converting SpecBench Medusa to this FW Addition of Initial TRTLLM AutoDeploy Specdec support Updates to all frameworks for better perf (overlap/async scheduling etc) ## Usage <!-- You can potentially add a usage example below. --> ```python # Add a code snippet demonstrating how to use this ``` ## Testing <!-- Mention how have you tested your change if applicable. --> ## Before your PR is "*Ready for review*" <!-- If you haven't finished some of the above items you can still open `Draft` PR. --> - **Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: Yes/No <!--- If No, explain why. --> - **Did you write any new necessary tests?**: Yes/No - **Did you add or update any necessary documentation?**: Yes/No - **Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**: Yes/No <!--- Only for new features, API changes, critical bug fixes or bw breaking changes. --> ## Additional Information <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Release Notes * **New Features** * Added SPEED-Bench dataset support with configurable throughput and qualitative configurations * Introduced SpecBench metrics with acceptance rate analysis and visualizations * Added progress bar during benchmark execution * New model implementations for auto-deployment and Medusa-style speculative decoding * Data preparation utility for benchmark datasets * Enhanced metrics with per-category analysis and performance charts * **Documentation** * Updated README with SPEED-Bench workflow and examples * New porting guide for integrating custom benchmark runners * **Refactor** * Streamlined model and runner interfaces for improved flexibility * Consolidated dataset implementations and removed deprecated base classes * **Chores** * Added required dependencies for data handling and visualizations <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Izzy Putterman <iputterman@nvidia.com>
82 lines
2.6 KiB
Python
82 lines
2.6 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
from typing import get_args
|
|
|
|
from specdec_bench import datasets
|
|
from specdec_bench.datasets.speed import config_type
|
|
|
|
datasets_available = {
|
|
"speed": datasets.SPEEDBench,
|
|
}
|
|
|
|
|
|
def prepare_data(args: argparse.Namespace) -> None:
|
|
"""Prepare and save benchmark data to disk.
|
|
|
|
Calls the dataset's ``prepare_data`` classmethod which downloads and
|
|
resolves all external data references, then saves the fully-resolved
|
|
result as a parquet file so that subsequent benchmark runs can load
|
|
directly from disk without re-downloading.
|
|
|
|
Args:
|
|
args: Parsed CLI arguments containing dataset type, config,
|
|
output directory, and optional filtering parameters.
|
|
"""
|
|
configs = get_args(config_type) if args.config == "all" else [args.config]
|
|
|
|
dataset_cls = datasets_available[args.dataset]
|
|
|
|
for config in configs:
|
|
print(f"Preparing config '{config}' ...")
|
|
|
|
output_path = dataset_cls.prepare_data(
|
|
output_dir=args.output_dir / args.dataset / config,
|
|
config_name=config,
|
|
)
|
|
|
|
print(f" -> Saved to {output_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(
|
|
description="Download and prepare benchmark datasets for specdec_bench.",
|
|
)
|
|
parser.add_argument(
|
|
"--dataset",
|
|
type=str,
|
|
default="speed",
|
|
choices=list(datasets_available.keys()),
|
|
help="Dataset to prepare (default: %(default)s)",
|
|
)
|
|
parser.add_argument(
|
|
"--config",
|
|
type=str,
|
|
default="all",
|
|
choices=[*list(get_args(config_type)), "all"],
|
|
help='SPEED-Bench configuration to prepare. Use "all" to prepare all configs. (default: %(default)s)',
|
|
)
|
|
parser.add_argument(
|
|
"--output_dir",
|
|
type=Path,
|
|
default=Path("data/"),
|
|
help="Directory to save the prepared dataset files (default: %(default)s)",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
prepare_data(args)
|