Mystique: Enabling Accurate and Scalable Generation of Production AI Benchmarks

TL;DR

Mystique is a scalable framework that uses detailed execution traces to generate accurate, portable, and adaptable AI benchmarks for production environments, addressing representativeness and rapid update challenges.

Contribution

It introduces a novel use of PyTorch execution traces for scalable, accurate, and portable AI benchmark generation in production settings.

Findings

Benchmarks closely match original models in execution time and system metrics.

Generated benchmarks are portable across different platforms.

The framework enables flexible and rapid benchmark creation.

Abstract

Building large AI fleets to support the rapidly growing DL workloads is an active research topic for modern cloud providers. Generating accurate benchmarks plays an essential role in designing the fast-paced software and hardware solutions in this space. Two fundamental challenges to make this scalable are (i) workload representativeness and (ii) the ability to quickly incorporate changes to the fleet into the benchmarks. To overcome these issues, we propose Mystique, an accurate and scalable framework for production AI benchmark generation. It leverages the PyTorch execution trace (ET), a new feature that captures the runtime information of AI models at the granularity of operators, in a graph format, together with their metadata. By sourcing fleet ETs, we can build AI benchmarks that are portable and representative. Mystique is scalable, due to its lightweight data collection, in terms of runtime overhead and instrumentation effort. It is also adaptive because ET composability allows flexible control on benchmark creation. We evaluate our methodology on several production AI models, and show that benchmarks generated with Mystique closely resemble original AI models, both in execution time and system-level metrics. We also showcase the portability of the generated benchmarks across platforms, and demonstrate several use cases enabled by the fine-grained composability of the execution trace.