Skip to content

[benchmarks] Increase compilation cache. #6509

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Feb 9, 2024
Merged

[benchmarks] Increase compilation cache. #6509

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
40c9b1b
Increase compilation cache size for a few models.
ysiraichi Feb 9, 2024
93ca6d3
Fix attribute access.
ysiraichi Feb 9, 2024
File filter

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions benchmarks/torchbench_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@
{
"test": "eval",
"xla": "PJRT",
"dynamo": None,
}, # TIMEOUT
],
"hf_T5_generate": [
Expand All @@ -85,6 +86,7 @@
{
"test": "eval",
"xla": "PJRT",
"dynamo": None,
}, # TIMEOUT
],
"doctr_det_predictor": [{
Expand Down Expand Up @@ -135,6 +137,13 @@
],
}

# Models that had more graphs to be compiled than the actual size of
# the cache.
NEED_LARGER_CACHE = {
"cm3leon_generate",
"hf_T5_generate",
}


class TorchBenchModelLoader(ModelLoader):

Expand Down Expand Up @@ -401,6 +410,9 @@ def update_process_env(self, process_env):
if precision_flag is not None:
process_env[precision_flag] = '1'

if self.model_name in NEED_LARGER_CACHE:
process_env["XLA_COMPILATION_CACHE_SIZE"] = "2048"

def pick_grad(self):
# special case
if self.model_name in ("maml",):
Expand Down