8000 Enable sample nightly PT2 benchmark on B200 (#158011) · pytorch/benchmark@b4dd068 · GitHub
[go: up one dir, main page]
Skip to content

Commit b4dd068

Browse files
huydhnfacebook-github-bot
huydhn
authored and
facebook-github-bot
committed
Enable sample nightly PT2 benchmark on B200 (#158011)
Summary: Per the discussion with nWEIdia, this resumes the work on pytorch/pytorch#157870 to enable PT2 benchmark on B200 ### Testing https://github.com/pytorch/pytorch/actions/runs/16615101382 X-link: pytorch/pytorch#158011 Approved by: https://github.com/nWEIdia, https://github.com/atalman Reviewed By: yangw-dev Differential Revision: D79490413 fbshipit-source-id: 0202249d62bd3011a1e60925b2f8442ddfdcbff7
1 parent fdea2a1 commit b4dd068

File tree

2 files changed

+15
-1
lines changed

2 files changed

+15
-1
lines changed

‎userbenchmark/dynamo/dynamobench/huggingface.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,11 @@ def process_hf_reformer_output(out):
106106
# on A100 GPUs - 40 GB.
107107
BATCH_SIZE_KNOWN_MODELS = {}
108108

109+
# Run only this selected group of models, leave this empty to run everything
110+
TORCHBENCH_ONLY_MODELS = [
111+
m.strip() for m in os.getenv("TORCHBENCH_ONLY_MODELS", "").split(",") if m.strip()
112+
]
113+
109114

110115
# TODO(sdym): use batch-size-file parameter of common.main, like torchbench.py
111 10A75 116
# Get the list of models and their batch sizes
@@ -116,6 +121,8 @@ def process_hf_reformer_output(out):
116121
lines = [line.rstrip() for line in lines]
117122
for line in lines:
118123
model_name, batch_size = line.split(",")
124+
if TORCHBENCH_ONLY_MODELS and model_name not in TORCHBENCH_ONLY_MODELS:
125+
continue
119126
batch_size = int(batch_size)
120127
BATCH_SIZE_KNOWN_MODELS[model_name] = batch_size
121128
assert len(BATCH_SIZE_KNOWN_MODELS)

‎userbenchmark/dynamo/dynamobench/timm_models.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,13 +39,20 @@ def pip_install(package):
3939
from timm.models import create_model
4040

4141
TIMM_MODELS = {}
42-
filename = os.path.join(os.path.dirname(__file__), "timm_models_list.txt")
4342

43+
# Run only this selected group of models, leave this empty to run everything
44+
TORCHBENCH_ONLY_MODELS = [
45+
m.strip() for m in os.getenv("TORCHBENCH_ONLY_MODELS", "").split(",") if m.strip()
46+
]
47+
48+
filename = os.path.join(os.path.dirname(__file__), "timm_models_list.txt")
4449
with open(filename) as fh:
4550
lines = fh.readlines()
4651
lines = [line.rstrip() for line in lines]
4752
for line in lines:
4853
model_name, batch_size = line.split(" ")
54+
if TORCHBENCH_ONLY_MODELS and model_name not in TORCHBENCH_ONLY_MODELS:
55+
continue
4956
TIMM_MODELS[model_name] = int(batch_size)
5057

5158

0 commit comments

Comments
 (0)
0