Dynamo TorchBench Integration Plan

[gs-olive]

Dynamo TorchBench Integration

TL;DR

With the novel PyTorch compile and export paths in 2.0 and our corresponding TensorRT implementations thereof (torch_tensorrt.dynamo.compile + torch_tensorrt.dynamo.fx_ts_compat.compile(..., is_aten=True, ...)), using TorchBench to efficiently test our performance on these models will help in improving Torch-TensorRT.

Goal(s)

The objective of adding our torch.compile and torch._dynamo.export implementations to TorchBench for performance testing is to evaluate our model coverage on a large set of key models, and to measure our performance on critical metrics.

Usecases

Proposed APIs / UX

Our compile and export paths will be added in a single folder called "torch_tensorrt" in the benchmark/userbenchmark directory in torchbenchmark. This directory is intended for custom benchmarking scripts and would be the location for testing our custom backends.

Users could then benchmark their choice of models on our Torch-TensorRT compile and export paths using the convenient TorchBench APIs. To run all of the models on the benchmark, for example, they might use:

# Compile Path
python run_benchmark.py torch_tensorrt --path compile

# Export Path
python run_benchmark.py torch_tensorrt --path export

Internal Implementation

Design

Following the instructions provided in the documentation here, we would create a directory in the benchmark/userbenchmark folder. Then, we implement an __init__.py file which has a run method which can take in arbitrary arguments and which outputs a formatted TorchBench JSON file containing key metrics.

Additionally, taking inspiration from Torch distributed, we can implement an augmented CLI interface which can take in the --path argument (or similar) and select between benchmarking compile or export. This augmented CLI should also allow the specification of compile-time arguments for our backend implementations. For instance, min_block_size or enabled_precisions could be supported.

Question: Should a user have to edit the TorchBench source to add compile-time arguments, or should this CLI handle it?

We would need to ensure that failures are caught and reported as is expected by TorchBench so as not to crash the program early if one model fails to compile.

Extensions Required to Core API implementations

No changes are needed in the Torch-TensorRT codebase.

Implementation Phases

Prototype - Small/Medium

• Implement custom TorchBench directory for torch_tensorrt.dynamo.compile, without any custom CI
• Validate generation of JSON files in the correct logging format
• Upstream directory as a PR to torchbench

MVP - Medium

• Implement torch_tensorrt.dynamo.fx_ts_compat.compile for export path in TorchBench. This is possible in one of two ways:
  • A new folder in torchbench entirely dedicated to export
  • A CLI interface to be shared between compile and export
• Allow selection of a subset of models and specification of compile-time options for the backend being profiled/benchmarked