[WIP][benchmarks] overhaul benchmarks #11565

sayakpaul · 2025-05-16T08:28:23Z

What does this PR do?

This PR considerably simplifies how we do benchmarks. Instead of using entire pipeline-level benchmarks across different tasks, we will now ONLY benchmark the diffusion network that is the most compute-intensive part in a standard diffusion workflow.

To make the estimates more realistic, we will make use of pre-trained checkpoints and dummy inputs with reasonable dimensionalities.

I ran benchmarking_flux.py on an 80GB A100 on a batch size of 1 and got the following results:

By default, all benchmarks will use a batch size of 1, eliminating CFG.

How to add your benchmark?

Adding benchmarks for a new model class (SanaTransformer2DModel, for example) boils down to the following:

Define the dummy inputs of the model.
Define the benchmarking scenarios we should run the benchmark on.

This is what benchmarking_flux.py does. More modularization can be shipped afterward.

Idea would be to merge this PR with pre-configured benchmarks for a few popular models and open others to the community.

TODOs

Utilities:

To fire the execution of the individual model-level benchmarks sequentially.
To combine CSVs from multiple different model classes.
Central dataset update and Slack notification.

@DN6 could you give the approach a quick look? I can then work on resolving the TODOs.

sayakpaul · 2025-05-16T08:33:25Z

benchmarks/benchmarking_utils.py

+logger = logging.get_logger(__name__)
+
+
+def benchmark_fn(f, *args, **kwargs):


This automatically warms up the model. No need to do it explicitly.

sayakpaul · 2025-05-16T08:34:30Z

benchmarks/benchmarking_flux.py

+
+
+if __name__ == "__main__":
+    scenarios = [


Covered the following scenarios:

Regular BF16 with compilation

NF4

Layerwise upcasting

Group offloading

sayakpaul added 8 commits May 15, 2025 18:05

start overhauling the benchmarking suite.

24a46cc

fixes

ab7f381

fixes

cc0a38a

checking.

169f831

checking

ad18983

fixes.

31e34d5

error handling and logging.

36afdea

Merge branch 'main' into benchmarking-overhaul

0d3af90

sayakpaul commented May 16, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[WIP][benchmarks] overhaul benchmarks #11565

[WIP][benchmarks] overhaul benchmarks #11565

		logger = logging.get_logger(__name__)


		def benchmark_fn(f, args, *kwargs):

[WIP][benchmarks] overhaul benchmarks #11565

Are you sure you want to change the base?

[WIP][benchmarks] overhaul benchmarks #11565

Conversation

What does this PR do?

How to add your benchmark?

TODOs

Choose a reason for hiding this comment

Choose a reason for hiding this comment