[Model Suite] Add model correctness testing #179
Closed
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This was referenced Oct 2, 2025
meta-cla
bot
added
the
CLA Signed
Here we introduce model suite (model.py). The idea here to start and codify the ideas from jiannanWang/BackendBenchExamples. Specifically this PR deals with model registration and loading. ### Model Registration This PR creates a way of adding models to the suite and automatically validates them through CI. It also loads the models as well. The way these models are added is detailed in this readme. The tl;dir is we use a format similar to kernelbench and SakanaAI/robust-kbench where we pair model code with a config. Importantly the configs contain initialization code, forward pass arguments (both in a similar format to torchbench), and a list of ops in the forward and backwards passes. These ops are fairly important as they are what we want to point out to the researcher when they are optimizing a model. There is a README.md to help folks setup proper model code / configs. We also further verify these registrations are correct through CI. Specifically we run test/test_model_ops_configs.py to ensure the configs are formatted correctly. ### Added models I added 2 models here as examples. SmokeTestModel - This is simple model that uses aten.ops.mm as we can implement a correct version of this op ToyCoreOpsModel - This is a model which explicitly calls the backwards passes which are both in torchbench + core. ### Small Things - Added a --model-filter to the CLI as it will be needed to support filtering in model suite as it chooses things to test based on the model not set of ops ### Testing New tests are added so pytest resolves things here ### Future work with Model Suite #181
This PR adds another unit test to the model loading / config system in the last PR. Specifically, here we ensure that the ops specified in the config are run in the model itself. This is important as updates to torch could change how backwards passes could work. Furthermore, if we are expecting folks to write kernels for a set of ops and then run the model, we should guarentee those ops are used. ### Future work with Model Suite #181
This PR adds end to end model correctness testing testing to Model suite by comparing the outputs and gradients (after a backwards pass) with 1 iteration of the model. We also integrate it into CI. ### Testing Running `uv run python BackendBench/scripts/main.py --suite model --backend directory` with a working mm kernel and a watermarked kernel for everything else yeilds ```bash [2025-10-02 07:16:13][INFO][main.py] ============================================================ [2025-10-02 07:16:13][INFO][main.py] MODEL EVALUATION RESULTS [2025-10-02 07:16:13][INFO][main.py] ============================================================ [2025-10-02 07:16:13][INFO][model.py] Model: ToyCoreOpsModel [2025-10-02 07:16:13][INFO][model.py] Status: ✗ Failed (0/3 tests) [2025-10-02 07:16:13][INFO][model.py] ✗ small_batch [2025-10-02 07:16:13][INFO][model.py] Error: Model ToyCoreOpsModel::small_batch failed: Expected number of channels in input to be divisible by num_groups, but got input of shape [2, 3, 32, 32] and num_groups=8 [2025-10-02 07:16:13][INFO][model.py] ✗ medium_batch [2025-10-02 07:16:13][INFO][model.py] Error: Model ToyCoreOpsModel::medium_batch failed: Expected number of channels in input to be divisible by num_groups, but got input of shape [4, 3, 64, 64] and num_groups=8 [2025-10-02 07:16:13][INFO][model.py] ✗ large_input [2025-10-02 07:16:13][INFO][model.py] Error: Model ToyCoreOpsModel::large_input failed: Expected number of channels in input to be divisible by num_groups, but got input of shape [2, 3, 128, 128] and num_groups=8 [2025-10-02 07:16:13][INFO][model.py] Model: SmokeTestModel [2025-10-02 07:16:13][INFO][model.py] Status: ✓ Passed (3/3 tests) [2025-10-02 07:16:13][INFO][model.py] ✓ small_batch [2025-10-02 07:16:13][INFO][model.py] Output match: ✓ Gradients match: ✓ (4 gradients) [2025-10-02 07:16:13][INFO][model.py] ✓ medium_batch [2025-10-02 07:16:13][INFO][model.py] Output match: ✓ Gradients match: ✓ (4 gradients) [2025-10-02 07:16:13][INFO][model.py] ✓ large_batch [2025-10-02 07:16:13][INFO][model.py] Output match: ✓ Gradients match: ✓ (4 gradients) [2025-10-02 07:16:13][INFO][main.py] ============================================================ ``` ### Future work with Model Suite #181
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR adds end to end model correctness testing testing to Model suite by comparing the outputs and gradients (after a backwards pass) with 1 iteration of the model.
We also integrate it into CI.
Testing
Running
uv run python BackendBench/scripts/main.py --suite model --backend directorywith a working mm kernel and a watermarked kernel for everything elseyeilds
Future work with Model Suite
#181
Stack created with Sapling. Best reviewed with ReviewStack.