add_benchmarks-pytorch-AtomGit

已合并

add_benchmarks #30535

Lu_G创建于 2月5日从Lu_G/pytorch_6045： add_benchmarks_v2.7 合入到 Ascend/pytorch： v2.7.1

add_benchmarks #30535

已合并

Lu_G创建于 2月5日

从Lu_G/pytorch_6045 add_benchmarks_v2.7合入到Ascend/pytorch v2.7.1

共 8 个文件变更+3086-0

Abenchmarks/.gitignore+49-0

@@ -0,0 +1,49 @@
	1	+# Benchmark output files (CSV results)
	2	+*.csv
	3	+
	4	+# Profiler outputs
	5	+*.json
	6	+profile/
	7	+prof_output/
	8	+profiler_info.json
	9	+
	10	+# Python cache
	11	+__pycache__/
	12	+*.py[cod]
	13	+*$py.class
	14	+*.so
	15	+
	16	+# Torch/Inductor cache
	17	+traced_graph_cache/
	18	+
	19	+# NPU/CUDA Profiler outputs
	20	+npu_profile/
	21	+cuda_profile/
	22	+profiling/
	23	+
	24	+# fusion config
	25	+fusion_switch.cfg
	26	+
	27	+# Temporary files
	28	+*.tmp
	29	+*.log
	30	+*.out
	31	+
	32	+# IDE specific
	33	+.vscode/
	34	+.idea/
	35	+*.swp
	36	+*.swo
	37	+*~
	38	+
	39	+# OS specific
	40	+.DS_Store
	41	+Thumbs.db
	42	+
	43	+# External torchbenchmark dependency (cloned from GitHub)
	44	+# These directories are created by following the instructions in README.md:
	45	+# git clone https://github.com/pytorch/benchmark.git --depth=1
	46	+# The leading slash ensures we only ignore these directories in benchmarks/
	47	+# and not in other locations.
	48	+/benchmark/
	49	+/torchbenchmark/

Abenchmarks/README.md+123-0

@@ -0,0 +1,123 @@
	1	+# Torchbenchmark
	2	+
	3	+## 简介
	4	+
	5	+为了评测图模式能力，pytorch社区在CI HUD中提供了图模式应用在torchbench仓库的多个模型的加速效果。参考社区的实现，benchmarks提供了对部分指定模型的支持。用户可按照此README进行NPU图模式测试。
	6	+
	7	+## Torchbench
	8	+
	9	+## 准备环境
	10	+1. 安装requirements.txt中的依赖包
	11	+ ```shell
	12	+ pip install -r requirements.txt
	13	+ ```
	14	+
	15	+2. 下载pytorch/benchmark源码，并切换至指定commit id
	16	+ ```shell
	17	+ git clone https://github.com/pytorch/benchmark.git --depth=1
	zichun_ye2月18日下面的文件路径跟benchmark clone的路径相关，这里需要指定下载在哪里，为了方便用户指定，可以采用变量的方式，例如 `BENCHMARK_DIR=./benchmarks cd $BENCHMARK_DIR git clone https://github.com/pytorch/benchmark.git --depth=1` 下面所有的路径都可以通过`$BENCHMARK_DIR`方便替换
	18	+ cd benchmark
	19	+ git remote set-branches origin '9910b31cc17d175a781412fd9ca6f18a4ee04610'
	20	+ git fetch --depth 1 origin 9910b31cc17d175a781412fd9ca6f18a4ee04610
	21	+ git checkout 9910b31cc17d175a781412fd9ca6f18a4ee04610
	22	+ cd ..
	23	+ ```
	24	+
	25	+## 准备数据集
	26	+1. 部分模型需要少量数据集，名单如下
	27	+ ```
	28	+ INPUT_TARBALLS:
	29	+ # index file for S3 storage of the input data
	30	+ - pytorch_stargan_inputs.tar.gz
	31	+ - LearningToPaint_inputs.tar.gz
	32	+ - speech_transformer_inputs.tar.gz
	33	+ ```
	34	+
	35	+2. 下载数据集可以使用下方URL
	36	+ ```
	37	+ https://ossci-datasets.s3.amazonaws.com/torchbench/data/<TARBALL_NAME>
	38	+ ```
	39	+ 例如，下载pytorch_stargan_inputs数据集，可以使用
	40	+ ```
	41	+ https://ossci-datasets.s3.amazonaws.com/torchbench/data/pytorch_stargan_inputs.tar.gz
	42	+ ```
	43	+
	44	+3. 下载后的数据集放到指定目录下，以pytorch_stargan_inputs.tar.gz为例
	45	+ ```shell
	46	+ # 创建.data/目录
	47	+ cd ./benchmark/torchbenchmark/data/
	48	+ mkdir .data/
	49	+
	50	+ # 将数据集移动到./benchmark/torchbenchmark/data/.data目录下
	51	+ cd ../../../
	52	+ cp your/path/to/dataset/pytorch_stargan_inputs.tar.gz ./benchmark/torchbenchmark/data/.data
	53	+
	54	+ # 解压数据集
	55	+ cd ./benchmark/torchbenchmark/data/
	56	+ tar -xvzf ./benchmark/torchbenchmark/data/.data/pytorch_stargan_inputs.tar.gz
	57	+ ```
	58	+
	59	+## 运行测试
	60	+1. 精度和端到端总时间验证
	61	+
	62	+ ```shell
	63	+ # 不使能--only，默认运行torchbench_models_list.txt目录下的所有模型
	64	+ python3 torchbench.py --accuracy --cold-start-latency --train --float32 --backend inductor --iterations 50
	65	+
	66	+ # 使能--only，运行指定模型
	67	+ python3 torchbench.py --accuracy --cold-start-latency --train --float32 --backend inductor --only BERT_pytorch --iterations 50
	68	+ ```
	69	+ 执行上述命令后，会在终端界面分别打印出模型执行（eager模式和图模式）的单步"端到端时间"、单步loss、"端到端时间"平均值以及eager模式和图模式精度比较的结果（pass_accuracy or fail_accuracy）
	70	+
	71	+2. 编译总时间验证
	72	+ ```shell
	73	+ # 添加参数--dump-compile-time，会在模型测试结束后输出编译时间的测量结果
	74	+ python3 torchbench.py --accuracy --cold-start-latency --train --float32 --backend inductor --iterations 50 --dump-compile-time
	75	+ ```
	76	+ 执行上述命令后，会在终端界面打印出图模式下的算子编译时间
	77	+
	78	+3. 算子总时间验证
	79	+
	80	+ 算子时间验证需要开启profile工具，添加参数`--enable-profiler`，profile结果默认输出路径为`./profile`，用户指定输出路径可使用参数`--prof-output-path`。
	81	+ ```shell
	82	+ # 不指定profile输出路径
	83	+ python3 torchbench.py --accuracy --cold-start-latency --train --float32 --backend inductor --iterations 50 --enable-profiler
	84	+
	85	+ # 指定profile输出路径
	86	+ python3 torchbench.py --accuracy --cold-start-latency --train --float32 --backend inductor --iterations 50 --enable-profiler --prof-output-path 'your/path/for/profile/output/'
	87	+ ```
	88	+
	89	+ `./profile`下的目录结构示例如下，`step_trace_time.csv`文件中记录了模型执行（eager模式和图模式）的单步算子时间
	90	+ ```
	91	+ ./profile
	92	+ ├── BERT_pytorch
	93	+ │ ├── compile
	94	+ │ │ └── localhost.localdomain_xxxx_ascend_pt
	95	+ │ │ └── ASCEND_PROFILER_OUTPUT
	96	+ \| \| └── step_trace_time.csv
	97	+ │ └── eager
	98	+ │ └── localhost.localdomain_xxxx_ascend_pt
	99	+ │ └── ASCEND_PROFILER_OUTPUT
	100	+ \| └── step_trace_time.csv
	101	+ \| ······
	102	+ ```
	103	+
	104	+## 结果展示
	105	+1. 执行下述命令，测试精度、编译时间、端到端总时间，并保存日志
	106	+ ```shell
	107	+ python3 torchbench.py --accuracy --cold-start-latency --train --float32 --backend inductor --iterations 50 --dump-compile-time 2>&1 \| tee models.log
	108	+ ```
	109	+
	110	+2. 执行下述命令，打开profile，获取算子总时间
	111	+ ```shell
	112	+ python3 torchbench.py --accuracy --cold-start-latency --train --float32 --backend inductor --iterations 50 --enable-profiler
	113	+ ```
	114	+
	115	+3. 执行下述命令，获得模型的测试结果，默认输出到`./analysis.xlsx`文件中
	116	+ ```shell
	117	+ python3 extract_log.py --log_file models.log
	118	+ ```
	119	+4. `./analysis.xlsx`文件的基本内容如下，展示了模型名称，图模式与eager模式的精度比较结果，算子编译时间（ms），eager模式的端到端时间（ms），图模式的端到端时间（ms），端到端时间加速比，eager模式的算子时间（ms），图模式的算子时间（ms），算子时间加速比
	120	+
	121	+\| model_name \| accuracy \| op_compile_time \| eager_E2E_avg_time \| compile_E2E_avg_time \| E2E_speed_up_rate \| eager_OP_avg_time \| compile_OP_avg_time \| OP_speed_up_rate \|
	122	+\| --- \| --- \| --- \| --- \| --- \| --- \| --- \| --- \| --- \|
	123	+\| BERT_pytorch \| pass_accuracy \| 41055.906 \| 48.49 \| 20.26 \| 2.393385982 \| 20.399543 \| 17.4955115 \| 1.16598723 \|

Abenchmarks/common.py+1573-0

@@ -0,0 +1,1573 @@
	1	+#!/usr/bin/env python3
	2	+# BSD 3-Clause License
	3	+
	4	+# Copyright (c) 2019, pytorch
	5	+# All rights reserved.
	6	+
	7	+# Redistribution and use in source and binary forms, with or without
	8	+# modification, are permitted provided that the following conditions are met:
	9	+
	10	+# 1. Redistributions of source code must retain the above copyright notice, this
	11	+# list of conditions and the following disclaimer.
	12	+
	13	+# 2. Redistributions in binary form must reproduce the above copyright notice,
	14	+# this list of conditions and the following disclaimer in the documentation
	15	+# and/or other materials provided with the distribution.
	16	+
	17	+# 3. Neither the name of the copyright holder nor the names of its
	18	+# contributors may be used to endorse or promote products derived from
	19	+# this software without specific prior written permission.
	20	+
	21	+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
	22	+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
	23	+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
	24	+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
	25	+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
	26	+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
	27	+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
	28	+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
	29	+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
	30	+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
	31	+
	32	+from __future__ import annotations
	33	+
	34	+import argparse
	35	+import collections
	36	+import contextlib
	37	+import copy
	38	+import csv
	39	+import dataclasses
	40	+import functools
	41	+import importlib
	42	+import itertools
	43	+import logging
	44	+import os
	45	+import pathlib
	46	+import random
	47	+import shutil
	48	+import signal
	49	+import subprocess
	50	+import sys
	51	+import time
	52	+import warnings
	53	+from contextlib import contextmanager
	54	+
	55	+from typing import Any, Callable, Mapping, NamedTuple, Optional, Tuple, Type
	56	+from unittest.mock import MagicMock
	57	+
	58	+import numpy as np
	59	+import pandas as pd
	60	+import psutil
	61	+import torch
	62	+import torch._dynamo
	63	+import torch._dynamo.utils
	64	+import torch.distributed
	65	+import torch.fx._pytree as fx_pytree
	66	+import torch.multiprocessing as mp
	67	+from scipy.stats import gmean, ttest_ind
	68	+from torch._dynamo.profiler import fx_insert_profiling, Profiler
	69	+from torch._dynamo.testing import dummy_fx_compile, format_speedup, same
	70	+from torch._dynamo.utils import clone_inputs, graph_break_reasons
	71	+from torch._functorch.aot_autograd import set_model_name
	72	+from torch._inductor import config as inductor_config
	73	+from torch._inductor.utils import fresh_inductor_cache
	74	+from torch._subclasses.fake_tensor import FakeTensorMode
	75	+
	76	+from torch.utils import _pytree as pytree
	77	+from torch.utils._pytree import tree_map, tree_map_only
	78	+
	79	+from tqdm.auto import tqdm, trange
	80	+
	81	+try:
	82	+ import torch_npu
	83	+ is_npu_available = torch_npu.npu.is_available()
	84	+ from npu_support import patch_model
	85	+ from profiler import NPUProfiler
	86	+except ImportError:
	87	+ # ignore the error if torch_npu is not installed
	88	+ is_npu_available = False
	89	+ from profiler import CUDAProfiler
	90	+
	91	+from benchmark.userbenchmark.dynamo.dynamobench.common import (
	92	+ load_model_from_path, Stats, randomize_input, speedup_experiment_ds,
	93	+ baselines, null_experiment, DummyGradScaler, cast_to_bf16, cast_to_fp16,
	94	+ cast_to_fp64, cast_to_fp32, maybe_fresh_cache, maybe_init_distributed,
	95	+)
	96	+log = logging.getLogger(__name__)
	97	+
	98	+# We are primarily interested in TF32
	99	+torch.backends.cuda.matmul.allow_tf32 = True
	100	+
	101	+# Suppress torch.profiler spam
	102	+os.environ["KINETO_LOG_LEVEL"] = "5"
	103	+
	104	+current_name = ""
	105	+current_device = ""
	106	+current_batch_size = None
	107	+output_filename = None
	108	+
	109	+MAX_DOWNLOAD_ATTEMPTS = 5
	110	+
	111	+
	112	+class PathManager:
	113	+ MAX_PATH_LENGTH = 4096
	114	+ MAX_FILE_NAME_LENGTH = 255
	115	+ DATA_FILE_AUTHORITY = 0o640
	116	+ DATA_DIR_AUTHORITY = 0o750
	117	+
	118	+ @classmethod
	119	+ def check_path_owner_consistent(cls, path: str):
	120	+ if not os.path.exists(path):
	121	+ msg = f"The path does not exist: {path}"
	122	+ raise RuntimeError(msg)
	123	+ if os.stat(path).st_uid != os.getuid():
	124	+ warnings.warn(f"Warning: The {path} owner does not match the current user.")
	125	+
	126	+ @classmethod
	127	+ def create_file_safety(cls, path: str):
	128	+ msg = f"Failed to create file: {path}"
	129	+ if os.path.islink(path):
	130	+ raise RuntimeError(msg)
	131	+ if os.path.exists(path):
	132	+ return
	133	+ try:
	134	+ path = os.path.realpath(path)
	135	+ os.close(os.open(path, os.O_WRONLY \| os.O_CREAT, cls.DATA_FILE_AUTHORITY))
	136	+ except Exception as err:
	137	+ raise RuntimeError(msg) from err
	138	+
	139	+ @classmethod
	140	+ def check_directory_path_readable(cls, path):
	141	+ cls.check_path_owner_consistent(path)
	142	+ if os.path.islink(path):
	143	+ msg = f"Invalid path is a soft chain: {path}"
	144	+ raise RuntimeError(msg)
	145	+ if not os.access(path, os.R_OK):
	146	+ msg = f"The path permission check failed: {path}"
	147	+ raise RuntimeError(msg)
	148	+
	149	+ @classmethod
	150	+ def check_directory_path_writeable(cls, path):
	151	+ cls.check_path_owner_consistent(path)
	152	+ if os.path.islink(path):
	153	+ msg = f"Invalid path is a soft chain: {path}"
	154	+ raise RuntimeError(msg)
	155	+ if not os.access(path, os.W_OK):
	156	+ msg = f"The path permission check failed: {path}"
	157	+ raise RuntimeError(msg)
	158	+
	159	+callbacks = []
	160	+
	161	+
	162	+def register_callback(callback):
	163	+ callbacks.append(callback)
	164	+
	165	+
	166	+def model_specified_by_path(path_and_class_str):
	167	+ return ":" in path_and_class_str
	168	+
	169	+
	170	+
	171	+def output_csv(filename, headers, row):
	172	+ abspath = os.path.abspath(filename)
	173	+ if os.path.exists(filename):
	174	+ PathManager.check_directory_path_readable(abspath)
	175	+ with open(filename) as fd:
	176	+ lines = list(csv.reader(fd)) or [[]]
	177	+ if headers and len(headers) > len(lines[0]):
	178	+ # if prior results failed the header might not be filled in yet
	179	+ lines[0] = headers
	180	+ else:
	181	+ headers = lines[0]
	182	+ else:
	183	+ lines = [headers]
	184	+ lines.append([(f"{x:.6f}" if isinstance(x, float) else x) for x in row])
	185	+ PathManager.create_file_safety(abspath)
	186	+ PathManager.check_directory_path_writeable(abspath)
	187	+ with open(filename, "w") as fd:
	188	+ writer = csv.writer(fd, lineterminator="\n")
	189	+ for line in lines:
	190	+ writer.writerow(list(line) + ["0"] * (len(headers) - len(line)))
	191	+
	192	+
	193	+def nothing(f):
	194	+ return f
	195	+
	196	+
	197	+@functools.lru_cache(None)
	198	+def patch_torch_manual_seed():
	199	+ """Make torch manual seed deterministic. Helps with accuracy testing."""
	200	+
	201	+ def deterministic_torch_manual_seed(args, *kwargs):
	202	+ from torch._C import default_generator
	203	+
	204	+ seed = 1337
	205	+
	206	+ if not torch.cuda._is_in_bad_fork():
	207	+ torch.cuda.manual_seed_all(seed)
	208	+ if is_npu_available:
	209	+ torch_npu.npu.manual_seed_all(seed)
	210	+ return default_generator.manual_seed(seed)
	211	+
	212	+ torch.manual_seed = deterministic_torch_manual_seed
	213	+
	214	+
	215	+def synchronize():
	216	+ pass
	217	+
	218	+
	219	+
	220	+def timed(
	221	+ model,
	222	+ model_iter_fn,
	223	+ example_inputs,
	224	+ times=1,
	225	+ return_result=False,
	226	+ collect_outputs=False,
	227	+):
	228	+ synchronize()
	229	+ time_total = 0
	230	+ # Dont collect outputs to correctly measure timing
	231	+ for _ in range(times):
	232	+ # Put this call inside the loop to reset the seed for each iteration.
	233	+ # Don't include reset_rng_state() to correctly measure timing
	234	+ reset_rng_state()
	235	+ t_iter_begin = time.perf_counter()
	236	+ result = model_iter_fn(model, example_inputs, collect_outputs=collect_outputs)
	237	+ t_iter_end = time.perf_counter()
	238	+ time_total += t_iter_end - t_iter_begin
	239	+
	240	+ t_0 = time.perf_counter()
	241	+ synchronize()
	242	+ t_1 = time.perf_counter()
	243	+ time_total += t_1 - t_0
	244	+ return (time_total, result) if return_result else time_total
	245	+
	246	+
	247	+
	248	+
	249	+
	250	+def speedup_experiment(args, model_iter_fn, model, example_inputs, **kwargs):
	251	+ """
	252	+ Measure speedups over eager.
	253	+
	254	+ Writes to ./speedups.csv
	255	+ """
	256	+
	257	+ timings = np.zeros((args.repeat, 2), np.float64)
	258	+ # if we randomize the input, we should also check the result is correct
	259	+ should_randomize_input = args.randomize_input
	260	+
	261	+ from torch._inductor.utils import maybe_profile
	262	+
	263	+ @contextlib.contextmanager
	264	+ def maybe_mark_profile(args, *kwargs):
	265	+ prof: torch.profiler.profile = kwargs.pop("p", None)
	266	+ mark = kwargs.pop("mark", None)
	267	+ if prof:
	268	+ with torch.profiler.record_function(mark):
	269	+ yield
	270	+ else:
	271	+ yield
	272	+
	273	+ times = args.iterations_per_run
	274	+
	275	+ tolerance = 1e-4
	276	+ torch._dynamo.config.repro_tolerance = tolerance
	277	+
	278	+ with maybe_profile(args.export_profiler_trace) as p:
	279	+ frozen_model_iter_fn = torch._dynamo.run(model_iter_fn)
	280	+
	281	+ for rep in trange(args.repeat, desc="running benchmark"):
	282	+ inputs = (
	283	+ randomize_input(copy.deepcopy(example_inputs))
	284	+ if should_randomize_input
	285	+ else example_inputs
	286	+ )
	287	+
	288	+ # interleave the runs to handle frequency scaling and load changes
	289	+ with maybe_mark_profile(p=p, mark="expected"):
	290	+ timings[rep, 0], _ = timed(
	291	+ model,
	292	+ model_iter_fn,
	293	+ inputs,
	294	+ return_result=True,
	295	+ times=times,
	296	+ collect_outputs=args.collect_outputs,
	297	+ )
	298	+
	299	+ with maybe_mark_profile(p=p, mark="actual"):
	300	+ timings[rep, 1], _ = timed(
	301	+ model,
	302	+ frozen_model_iter_fn,
	303	+ inputs,
	304	+ return_result=True,
	305	+ times=times,
	306	+ collect_outputs=args.collect_outputs,
	307	+ )
	308	+
	309	+ if args.export_profiler_trace:
	310	+ name = args.profiler_trace_name + "_" + model.name + ".json"
	311	+ name = os.path.join(torch._dynamo.config.base_dir, name)
	312	+ p.export_chrome_trace(name)
	313	+ median = np.median(timings, axis=0)
	314	+ speedup = median[0] / median[1]
	315	+
	316	+ first_headers = ["dev", "name", "batch_size"]
	317	+ first_fields = [current_device, current_name, current_batch_size]
	318	+ if "tag" in kwargs:
	319	+ first_headers.append("tag")
	320	+ first_fields.append(kwargs["tag"])
	321	+ headers = first_headers + ["speedup", "abs_latency"]
	322	+ row = first_fields + [float(speedup), median[1] * 1000]
	323	+ msg = f"{speedup:.3f}x"
	324	+ if args.baseline:
	325	+ headers.extend(
	326	+ [
	327	+ "baseline",
	328	+ "speedup_vs_baseline",
	329	+ ]
	330	+ )
	331	+ df = pd.read_csv(args.baseline)
	332	+ try:
	333	+ baseline_speedup = df[df["name"] == current_name]["speedup"].item()
	334	+ row.extend([baseline_speedup, speedup / baseline_speedup])
	335	+ msg = f"{baseline_speedup:.3f}x -> {speedup:.3f}x [{speedup / baseline_speedup:.3f}x]"
	336	+ except (KeyError, ZeroDivisionError):
	337	+ row.extend(
	338	+ [
	339	+ 0.0,
	340	+ 0.0,
	341	+ ]
	342	+ )
	343	+ if "compilation_latency" in kwargs:
	344	+ headers += [
	345	+ "compilation_latency",
	346	+ "compression_ratio",
	347	+ "eager_peak_mem",
	348	+ "dynamo_peak_mem",
	349	+ ]
	350	+ row.append(kwargs["compilation_latency"])
	351	+ row.append(kwargs["compression_ratio"])
	352	+ row.append(kwargs["eager_peak_mem"])
	353	+ row.append(kwargs["dynamo_peak_mem"])
	354	+ if "dynamo_stats" in kwargs:
	355	+ for k, v in kwargs["dynamo_stats"].items():
	356	+ headers.append(k)
	357	+ row.append(v)
	358	+ output_csv(
	359	+ output_filename,
	360	+ headers,
	361	+ row,
	362	+ )
	363	+ headers, data = torch._dynamo.utils.compile_times(repr="csv", aggregate=True)
	364	+ if output_filename.find(".csv") <= 0:
	365	+ raise AssertionError(f"expected output_filename to be a .csv, but got {output_filename}")
	366	+ output_csv(
	367	+ output_filename[:-4] + "_compilation_metrics.csv",
	368	+ first_headers + headers,
	369	+ first_fields + data,
	370	+ )
	371	+ return msg
	372	+
	373	+
	374	+
	375	+
	376	+
	377	+
	378	+def read_batch_size_from_file(args, filename, model_name):
	379	+ batch_size = None
	380	+ if os.path.exists("benchmarks"):
	381	+ filename = os.path.join("benchmarks", filename)
	382	+ if not os.path.exists(filename):
	383	+ raise AssertionError(filename)
	384	+ abspath = os.path.abspath(filename)
	385	+ PathManager.check_directory_path_readable(abspath)
	386	+ with open(filename) as f:
	387	+ lines = f.readlines()
	388	+ lines = [i.split(",") for i in lines if len(i.strip()) > 0]
	389	+ for val in lines:
	390	+ cur_name, b = val
	391	+ if model_name == cur_name:
	392	+ batch_size = int(b)
	393	+ if batch_size is None:
	394	+ log.warning("Could not find batch size for %s", model_name)
	395	+ elif batch_size == -1:
	396	+ raise RuntimeError(
	397	+ f"Batch size is unset for {model_name} in {args.batch_size_file}"
	398	+ )
	399	+ print(f"batch size: {batch_size}")
	400	+ return batch_size
	401	+
	402	+
	403	+def get_peak_memory():
	404	+ return torch.cuda.max_memory_allocated() / 10**9
	405	+
	406	+
	407	+def get_peak_memory_npu():
	408	+ return torch_npu.npu.max_memory_allocated() / 10**9
	409	+
	410	+
	411	+
	412	+def reset_rng_state():
	413	+ torch.manual_seed(1337)
	414	+ random.seed(1337)
	415	+ np.random.seed(1337)
	416	+
	417	+
	418	+def get_dynamo_stats():
	419	+ # adding a helper to do subtraction on it
	420	+ return collections.Counter(
	421	+ {
	422	+ "calls_captured": torch._dynamo.utils.counters["stats"]["calls_captured"],
	423	+ "unique_graphs": torch._dynamo.utils.counters["stats"]["unique_graphs"],
	424	+ "graph_breaks": sum(torch._dynamo.utils.counters["graph_break"].values()),
	425	+ # NB: The plus removes zero counts
	426	+ "unique_graph_breaks": len(+torch._dynamo.utils.counters["graph_break"]),
	427	+ }
	428	+ )
	429	+
	430	+
	431	+class BenchmarkRunner:
	432	+ def __init__(self):
	433	+ self.model_iter_fn = None
	434	+ self.grad_scaler = DummyGradScaler()
	435	+ self.autocast = contextlib.nullcontext
	436	+ self.optimizer = None
	437	+ self._args = None
	438	+
	439	+ def setup_amp(self):
	440	+ if self.args.only in self.fp32_only_models:
	441	+ return
	442	+
	443	+ if self.args.amp and self.args.devices == ["cuda"]:
	444	+ # AMP training can lead to small loss values which can undeflow
	445	+ # gradient values returning in zero gradients. To solve this
	446	+ # problem, PyTorch introduces GradScaler. GradScaler is a stateful
	447	+ # structure, that scales the loss values to prevent underflow. Loss
	448	+ # values are big at the beginning of training (therefore not
	449	+ # requiring scaling), while loss value tends to be small as network
	450	+ # starts getting better (requiring scaling). GradScaler manages all
	451	+ # of this fine tuning, checking the gradients are turning to inf,
	452	+ # discarding such batches.
	453	+
	454	+ # Since we are not running a long iteration, default value of
	455	+ # init_scale 65536 is going to turn all gradients to inf. Therefore,
	456	+ # we just use a init_scale of 2.0 for benchmarking purpose.
	457	+
	458	+ # Disabling Gradscaler because
	459	+ # 1) Benchmark setup runs 2 iterations of fwd-bwd. So, not useful.
	460	+ # 2) Current setup shares grad_scaler for eager and dynamo model,
	461	+ # which is bad as Gradscaler has state and can adjust the scaling
	462	+ # factor between eager and dynamo run, making accuracy check
	463	+ # harder.
	464	+ self.autocast = torch.cuda.amp.autocast
	465	+ elif (self.args.bfloat16 or self.args.amp) and self.args.devices == ["cpu"]:
	466	+ self.autocast = torch.cpu.amp.autocast
	467	+ elif self.args.amp and self.args.devices == ["npu"]:
	468	+ self.autocast = torch_npu.npu.amp.autocast
	469	+
	470	+ def init_optimizer(self, name, device, params, learning_rate=0.01):
	471	+ if device == "cuda" and self.args.training:
	472	+ self.optimizer = torch.optim.SGD(params, lr=learning_rate, foreach=True)
	473	+ elif device == "npu" and self.args.training:
	474	+ # Currently, npu dynamo doesn't support foreach=True.
	475	+ # There may be changes here in the future.
	476	+ self.optimizer = torch.optim.SGD(params, lr=learning_rate, foreach=False)
	477	+ else:
	478	+ self.optimizer = None
	479	+
	480	+ @property
	481	+ def args(self):
	482	+ return self._args
	483	+
	484	+ @args.setter
	485	+ def args(self, args):
	486	+ self._args = args
	487	+
	488	+ @property
	489	+ def skip_models(self):
	490	+ return set()
	491	+
	492	+ @property
	493	+ def skip_models_for_cuda(self):
	494	+ return set()
	495	+
	496	+ @property
	497	+ def slow_models(self):
	498	+ return set()
	499	+
	500	+ @property
	501	+ def very_slow_models(self):
	502	+ return set()
	503	+
	504	+ @property
	505	+ def non_deterministic_models(self):
	506	+ return set()
	507	+
	508	+ @property
	509	+ def fp32_only_models(self):
	510	+ return set()
	511	+
	512	+ @property
	513	+ def force_amp_for_fp16_bf16_models(self):
	514	+ return set()
	515	+
	516	+ @property
	517	+ def failing_torchinductor_models(self):
	518	+ return set()
	519	+
	520	+ def get_tolerance_and_cosine_flag(self, is_training, curr_device, name):
	521	+ raise NotImplementedError()
	522	+
	523	+ def get_learning_rate(self, is_training, curr_device, name):
	524	+ raise NotImplementedError()
	525	+
	526	+ @property
	527	+ def equal_nan(self):
	528	+ equal_nan = True
	529	+ if self.args.float32:
	530	+ equal_nan = False
	531	+ return equal_nan
	532	+
	533	+ def iter_models(self, args):
	534	+ for model_name in self.iter_model_names(args):
	535	+ for device in args.devices:
	536	+ try:
	537	+ yield self.load_model(
	538	+ device,
	539	+ model_name,
	540	+ batch_size=args.batch_size,
	541	+ )
	542	+ except NotImplementedError:
	543	+ continue # bad benchmark implementation
	544	+
	545	+ def deepcopy_model(self, model):
	546	+ return copy.deepcopy(model)
	547	+
	548	+ def cast_based_on_args(self, model, example_inputs):
	549	+ if self.args.float32 or self.args.only in self.fp32_only_models:
	550	+ if not self.args.float32:
	551	+ log.warning("Model %s supports float32 only", self.args.only)
	552	+ model, example_inputs = cast_to_fp32(model, example_inputs)
	553	+ elif self.args.float16:
	554	+ if self.args.only in self.force_amp_for_fp16_bf16_models:
	555	+ log.warning(
	556	+ "Model %s does not support float16, running with amp instead",
	557	+ self.args.only,
	558	+ )
	559	+ self.args.amp = True
	560	+ self.setup_amp()
	561	+ else:
	562	+ model, example_inputs = cast_to_fp16(model, example_inputs)
	563	+ elif self.args.bfloat16:
	564	+ if self.args.only in self.force_amp_for_fp16_bf16_models:
	565	+ log.warning(
	566	+ "Model %s does not support bfloat16, running with amp instead",
	567	+ self.args.only,
	568	+ )
	569	+ self.args.amp = True
	570	+ self.setup_amp()
	571	+ else:
	572	+ model, example_inputs = cast_to_bf16(model, example_inputs)
	573	+
	574	+ return model, example_inputs
	575	+
	576	+ def validate_model(self, model, example_inputs):
	577	+ """
	578	+ Runs the eager model with example inputs to ensure that eager passes.
	579	+ """
	580	+ model = self.deepcopy_model(model)
	581	+ example_inputs = clone_inputs(example_inputs)
	582	+ model, example_inputs = self.cast_based_on_args(model, example_inputs)
	583	+ try:
	584	+ self.model_iter_fn(model, example_inputs)
	585	+ except Exception as e:
	586	+ raise NotImplementedError("Eager model failed to run") from e
	587	+
	588	+ def maybe_cast(self, model, example_inputs):
	589	+ model = self.deepcopy_model(model)
	590	+ example_inputs = clone_inputs(example_inputs)
	591	+ model, example_inputs = self.cast_based_on_args(model, example_inputs)
	592	+ return model, example_inputs
	593	+
	594	+ def run_n_iterations(self, mod, inputs, run_mode=None):
	595	+ n = self.args.iterations
	596	+ if run_mode is None:
	597	+ for _ in range(n - 1):
	598	+ self.model_iter_fn(mod, inputs, collect_outputs=False)
	599	+ return self.model_iter_fn(mod, inputs, collect_outputs=True)
	600	+
	601	+ start_step = int(n * 0.3)
	602	+ end_step = min(int(n * 0.8) - 1, n - 1)
	603	+
	604	+ step_times = []
	605	+ prof_output_dir = os.path.join(self.args.prof_output_path, self.args.only, run_mode)
	606	+ if is_npu_available:
	607	+ prof = NPUProfiler(enable=self.args.enable_profiler, warmup=10, active=n, save_path=prof_output_dir)
	608	+ else:
	609	+ prof = CUDAProfiler(enable=self.args.enable_profiler, warmup=10, active=n, save_path=prof_output_dir)
	610	+
	611	+ prof.start()
	zichun_ye2月18日建立开关，prof可以不默认使能或者关闭
	612	+ for i in range(n):
	613	+ start = time.perf_counter()
	614	+ output = self.model_iter_fn(mod, inputs, collect_outputs=(i == n - 1))
	615	+ synchronize()
	616	+ end = time.perf_counter()
	617	+ elapsed_ms = (end - start) * 1000
	618	+ prof.step()
	619	+ step_times.append(elapsed_ms)
	620	+ if i != n - 1:
	621	+ print(f"[{run_mode}] step: {i+1} step_time: {elapsed_ms} ms loss: {output}")
	622	+ else:
	623	+ print(f"[{run_mode}] step: {i+1} step_time: {elapsed_ms} ms")
	624	+ prof.stop()
	625	+
	626	+ steps = step_times[start_step:end_step + 1]
	627	+ if steps:
	628	+ total_ms = sum(steps)
	629	+ avg_ms = total_ms / len(steps)
	630	+ print(f"[{run_mode}] summary [{start_step+1}-{end_step+1}] "
	631	+ f"total steps time: {total_ms:.2f} ms, "
	632	+ f"avg step time: {avg_ms:.2f} ms")
	633	+
	634	+ return output
	635	+
	636	+ def optimizer_zero_grad(self, mod):
	637	+ if self.optimizer is not None:
	638	+ self.optimizer.zero_grad(True)
	639	+ else:
	640	+ mod.zero_grad(True)
	641	+
	642	+ def optimizer_step(self):
	643	+ if self.optimizer is not None:
	644	+ self.optimizer.step()
	645	+
	646	+ def deepcopy_and_maybe_ddp(self, model):
	647	+ model = self.deepcopy_model(model)
	648	+ if self.args.ddp:
	649	+ if not torch.distributed.is_available():
	650	+ raise AssertionError("Can't use DDP without a distributed enabled build")
	651	+ from torch.nn.parallel import DistributedDataParallel as DDP
	652	+
	653	+ model = DDP(model, find_unused_parameters=True)
	654	+ return model
	655	+
	656	+ def check_accuracy(
	657	+ self, name, model, example_inputs, optimize_ctx, experiment, tag
	658	+ ):
	659	+ """
	660	+ Checks accuracy.
	661	+ 1) Collect the outputs with fp64 datatype. This is useful for error checking.
	662	+ 2) Checks if eager itself has variations.
	663	+ """
	664	+ start_stats = get_dynamo_stats()
	665	+ lr = self.get_learning_rate(self.args.training, current_device, name)
	666	+ print(f"learning rate: {lr}")
	667	+
	668	+ def record_status(accuracy_status, dynamo_start_stats):
	669	+ """
	670	+ Records the status in the csv file
	671	+ """
	672	+ if current_name in self.non_deterministic_models:
	673	+ if accuracy_status in (
	674	+ "pass_accuracy",
	675	+ "eager_two_runs_differ",
	676	+ "fail_accuracy",
	677	+ ):
	678	+ accuracy_status = "pass_accuracy"
	679	+
	680	+ headers = ["dev", "name", "batch_size", "accuracy"]
	681	+ fields = [current_device, current_name, current_batch_size, accuracy_status]
	682	+
	683	+ if tag is not None:
	684	+ headers.insert(3, "tag")
	685	+ fields.insert(3, tag)
	686	+
	687	+ dynamo_stats = get_dynamo_stats()
	688	+ dynamo_stats.subtract(dynamo_start_stats)
	689	+ for k, v in dynamo_stats.items():
	690	+ headers.append(k)
	691	+ fields.append(v)
	692	+
	693	+ output_csv(output_filename, headers, fields)
	694	+ return accuracy_status
	695	+
	696	+ # Collect the fp64 reference outputs to be used later for accuracy checking.
	697	+ fp64_outputs = None
	698	+ try:
	699	+ model_fp64, inputs_fp64 = cast_to_fp64(
	700	+ self.deepcopy_and_maybe_ddp(model),
	701	+ clone_inputs(example_inputs),
	702	+ )
	703	+ self.init_optimizer(name, current_device, model_fp64.parameters(), lr)
	704	+ fp64_outputs = self.run_n_iterations(model_fp64, inputs_fp64)
	705	+ fp64_outputs = tree_map(
	706	+ lambda x: x.to(torch.float64)
	707	+ if isinstance(x, torch.Tensor) and x.is_floating_point()
	708	+ else x,
	709	+ fp64_outputs,
	710	+ )
	711	+ except Exception:
	712	+ log.warning(
	713	+ "fp64 golden ref were not generated for %s. Setting accuracy check to cosine",
	714	+ name,
	715	+ )
	716	+ self.args.cosine = True
	717	+ fp64_outputs = None
	718	+
	719	+ tolerance, cos_similarity = self.get_tolerance_and_cosine_flag(
	720	+ self.args.training, current_device, name
	721	+ )
	722	+ print(f"tolerance: {tolerance}")
	723	+
	724	+ # Cast the model to float16/float32 as necessary
	725	+ model, example_inputs = self.maybe_cast(model, example_inputs)
	726	+ accuracy_status = "pass_accuracy"
	727	+
	728	+
	729	+ with self.pick_grad(name, self.args.training):
	730	+ # Get results of native pytorch
	731	+ reset_rng_state()
	732	+ try:
	733	+ model_copy = self.deepcopy_and_maybe_ddp(model)
	734	+ self.init_optimizer(name, current_device, model_copy.parameters(), lr)
	735	+ correct_result = self.run_n_iterations(
	736	+ model_copy, clone_inputs(example_inputs), "eager"
	737	+ )
	738	+ except Exception as e:
	739	+ accuracy_status = (
	740	+ "eager_1st_run_OOM"
	741	+ if isinstance(e, torch.cuda.OutOfMemoryError)
	742	+ else "eager_1st_run_fail"
	743	+ )
	744	+ log.exception(e)
	745	+ return record_status(accuracy_status, dynamo_start_stats=start_stats)
	746	+
	747	+ # Rerun native pytorch
	748	+ reset_rng_state()
	749	+ try:
	750	+ model_copy = self.deepcopy_and_maybe_ddp(model)
	751	+ self.init_optimizer(name, current_device, model_copy.parameters(), lr)
	752	+ correct_rerun_result = self.run_n_iterations(
	753	+ model_copy, clone_inputs(example_inputs)
	754	+ )
	755	+ except Exception as e:
	756	+ accuracy_status = (
	757	+ "eager_2nd_run_OOM"
	758	+ if isinstance(e, torch.cuda.OutOfMemoryError)
	759	+ else "eager_2nd_run_fail"
	760	+ )
	761	+ return record_status(accuracy_status, dynamo_start_stats=start_stats)
	762	+
	763	+ # Two eager runs should have exactly same result
	764	+ is_same = True
	765	+ try:
	766	+ if (
	767	+ not same(
	768	+ correct_result,
	769	+ correct_rerun_result,
	770	+ fp64_ref=None,
	771	+ cos_similarity=False,
	772	+ tol=0,
	773	+ equal_nan=self.equal_nan,
	774	+ )
	775	+ ):
	776	+ is_same = False
	777	+ except Exception as e:
	778	+ # Sometimes torch.allclose may throw RuntimeError
	779	+ is_same = False
	780	+
	781	+ if not is_same:
	782	+ accuracy_status = "eager_two_runs_differ"
	783	+ return record_status(accuracy_status, dynamo_start_stats=start_stats)
	784	+
	785	+ correct_rerun_result = None
	786	+
	787	+ # Run with Dynamo
	788	+ reset_rng_state()
	789	+ torch._dynamo.reset()
	790	+ try:
	791	+ model_copy = self.deepcopy_and_maybe_ddp(model)
	792	+ self.init_optimizer(name, current_device, model_copy.parameters(), lr)
	793	+ optimized_model = optimize_ctx(model_copy)
	794	+ new_result = self.run_n_iterations(optimized_model, example_inputs, "compile")
	795	+ except Exception as e:
	796	+ log.exception(e)
	797	+ print(
	798	+ "TorchDynamo optimized model failed to run because of following error"
	799	+ )
	800	+ accuracy_status = (
	801	+ "OOM"
	802	+ if isinstance(e, torch.cuda.OutOfMemoryError)
	803	+ else "fail_to_run"
	804	+ )
	805	+ return record_status(accuracy_status, dynamo_start_stats=start_stats)
	806	+
	807	+ try:
	808	+ if not same(
	809	+ correct_result,
	810	+ new_result,
	811	+ fp64_outputs,
	812	+ equal_nan=self.equal_nan,
	813	+ cos_similarity=cos_similarity,
	814	+ tol=tolerance,
	815	+ ):
	816	+ is_same = False
	817	+ except Exception as e:
	818	+ # Sometimes torch.allclose may throw RuntimeError
	819	+ is_same = False
	820	+
	821	+ if not is_same:
	822	+ accuracy_status = "fail_accuracy"
	823	+ return record_status(accuracy_status, dynamo_start_stats=start_stats)
	824	+
	825	+ return record_status(accuracy_status, dynamo_start_stats=start_stats)
	826	+
	827	+ def run_performance_test(
	828	+ self, name, model, example_inputs, optimize_ctx, experiment, tag=None
	829	+ ):
	830	+ def warmup(fn, model, example_inputs, mode, niters=5):
	831	+ peak_mem = 0
	832	+ start_stats = get_dynamo_stats()
	833	+ try:
	834	+ if current_device == "cuda":
	835	+ torch.cuda.reset_peak_memory_stats()
	836	+ torch.cuda.empty_cache()
	837	+ elif current_device == "npu":
	838	+ torch_npu.npu.reset_peak_memory_stats()
	839	+ torch_npu.npu.empty_cache()
	840	+ t0 = time.perf_counter()
	841	+ for _ in range(niters):
	842	+ fn(model, example_inputs)
	843	+ t1 = time.perf_counter()
	844	+ latency = t1 - t0
	845	+ if current_device == "cuda":
	846	+ peak_mem = get_peak_memory()
	847	+ elif current_device == "npu":
	848	+ peak_mem = get_peak_memory_npu()
	849	+ elif current_device == "cpu":
	850	+ total = psutil.virtual_memory().total
	851	+ percentage = psutil.Process(os.getpid()).memory_percent()
	852	+ peak_mem = percentage * total / 10**9
	853	+ except Exception as e:
	854	+ log.exception("Backend %s failed in warmup()", mode)
	855	+ raise RuntimeError(f"Backend {mode} failed in warmup()") from e
	856	+ dynamo_stats = get_dynamo_stats()
	857	+ dynamo_stats.subtract(start_stats)
	858	+ return latency, peak_mem, dynamo_stats
	859	+
	860	+ # Cast the model to float16/float32 as necessary
	861	+ model, example_inputs = self.maybe_cast(model, example_inputs)
	862	+
	863	+ model = self.deepcopy_and_maybe_ddp(model)
	864	+
	865	+ self.init_optimizer(name, current_device, model.parameters())
	866	+ with self.pick_grad(name, self.args.training):
	867	+ ok, total = Stats.reset_counters()
	868	+ experiment_kwargs = {}
	869	+ if tag is not None:
	870	+ experiment_kwargs["tag"] = tag
	871	+ results = []
	872	+ eager_latency, eager_peak_mem, _ = warmup(
	873	+ self.model_iter_fn, model, example_inputs, "eager"
	874	+ )
	875	+ optimized_model_iter_fn = optimize_ctx(self.model_iter_fn)
	876	+ dynamo_latency, dynamo_peak_mem, dynamo_stats = warmup(
	877	+ optimized_model_iter_fn, model, example_inputs, "dynamo"
	878	+ )
	879	+
	880	+ compilation_time = dynamo_latency - eager_latency
	881	+ compression_ratio = (
	882	+ eager_peak_mem / dynamo_peak_mem if dynamo_peak_mem else 0.0
	883	+ )
	884	+
	885	+ if experiment.func is speedup_experiment:
	886	+ experiment_kwargs["compilation_latency"] = compilation_time
	887	+ experiment_kwargs["compression_ratio"] = compression_ratio
	888	+ experiment_kwargs["eager_peak_mem"] = eager_peak_mem
	889	+ experiment_kwargs["dynamo_peak_mem"] = dynamo_peak_mem
	890	+ experiment_kwargs["dynamo_stats"] = dynamo_stats
	891	+
	892	+ if not hasattr(model, name):
	893	+ model.name = name
	894	+ results.append(experiment(model, example_inputs, **experiment_kwargs))
	895	+ return " ".join(map(str, results))
	896	+
	897	+
	898	+ def run_one_model(
	899	+ self,
	900	+ name,
	901	+ model,
	902	+ example_inputs,
	903	+ optimize_ctx,
	904	+ experiment,
	905	+ tag=None,
	906	+ ):
	907	+ mode = "train" if self.args.training else "eval"
	908	+ msg = f"{current_device:4} {mode:5} {current_name:34} "
	909	+ if tag:
	910	+ msg += f" {tag:26}"
	911	+ print(msg, flush=True)
	912	+
	913	+ start_stats = get_dynamo_stats()
	914	+
	915	+ if self.args.accuracy:
	916	+ status = self.check_accuracy(
	917	+ name, model, example_inputs, optimize_ctx, experiment, tag
	918	+ )
	919	+ print(status)
	920	+ if self.args.dump_compile_time:
	921	+ headers, values = torch._dynamo.utils.compile_times("csv")
	922	+ for header, value in zip(headers, values):
	923	+ if header == "async_compile.wait":
	924	+ numbers = [float(num.strip()) for num in value.split(',') if num.strip()]
	925	+ op_compile_time = sum(numbers)
	926	+ print(f"op_compile_time:{op_compile_time * 1e3} ms", )
	927	+
	928	+ elif self.args.performance:
	929	+ status = self.run_performance_test(
	930	+ name, model, example_inputs, optimize_ctx, experiment, tag
	931	+ )
	932	+ print(status)
	933	+ stats = get_dynamo_stats()
	934	+ stats.subtract(start_stats)
	935	+
	936	+ if self.args.log_graph_breaks or self.args.print_graph_breaks:
	937	+ filename = f"{output_filename.rstrip('.csv')}_graph_breaks.csv"
	938	+
	939	+ def add_double_quotes(x):
	940	+ # Delimiter because reason could have comma
	941	+ return f'"{x}"'
	942	+
	943	+ for graph_break in graph_break_reasons:
	944	+ reason = add_double_quotes(graph_break.reason)
	945	+ user_stack = add_double_quotes(
	946	+ ", ".join([str(x) for x in graph_break.user_stack])
	947	+ )
	948	+ output_csv(
	949	+ filename,
	950	+ ["model", "reason", "user_stack"],
	951	+ [current_name, reason, user_stack],
	952	+ )
	953	+
	954	+ if self.args.stats:
	955	+ Stats.print_summary()
	956	+
	957	+
	958	+def parse_args(args=None):
	959	+ parser = argparse.ArgumentParser()
	960	+ parser.add_argument(
	961	+ "--devices", "--device", "-d", action="append", help="cpu or cuda"
	962	+ )
	963	+ parser.add_argument("--device-index", help="CUDA device index")
	964	+ parser.add_argument(
	965	+ "--repeat", "-n", type=int, default=30, help="number of timing runs"
	966	+ )
	967	+ iterations_per_run_help = """
	968	+ Run this may iterations for each time measurement. This is mainly used for
	969	+ XLA training. We want to run multiple iterations per measurement so the
	970	+ tracing and computation for different iteartions can overlap with each
	971	+ other. This makes sure we have an accurate xla baseline.
	972	+ """
	973	+ parser.add_argument(
	974	+ "--iterations-per-run", type=int, default=1, help=iterations_per_run_help
	975	+ )
	976	+ parser.add_argument(
	977	+ "--randomize-input",
	978	+ action="store_true",
	979	+ help="Whether to randomize the input values. Dimensions will be kept the same.",
	980	+ )
	981	+ parser.add_argument(
	982	+ "--nopython", action="store_true", help="Turn graph breaks into errors"
	983	+ )
	984	+ parser.add_argument(
	985	+ "--no-skip",
	986	+ action="store_true",
	987	+ help="run models that are in the global SKIP list",
	988	+ )
	989	+ parser.add_argument(
	990	+ "--batch-size", "--batch_size", type=int, help="batch size for benchmarking"
	991	+ )
	992	+ parser.add_argument(
	993	+ "--iterations", type=int, default=50, help="how many iterations to run"
	994	+ )
	995	+ parser.add_argument(
	996	+ "--batch-size-file", type=str, help="String to load batch size from"
	997	+ )
	998	+ parser.add_argument("--cosine", action="store_true", help="use cosine similarity")
	999	+ parser.add_argument(
	1000	+ "--cpp-wrapper", action="store_true", help="turn on cpp/cuda wrapper codegen"
	1001	+ )
	1002	+ parser.add_argument(
	1003	+ "--freezing", action="store_true", help="turn on freezing", default=False
	1004	+ )
	1005	+ parser.add_argument(
	1006	+ "--only",
	1007	+ help="""Run just one model from torchbench. Or
	1008	+ specify the path and class name of the model in format like:
	1009	+ --only=path:<MODEL_FILE_PATH>,class:<CLASS_NAME>
	1010	+
	1011	+ Due to the fact that dynamo changes current working directory,
	1012	+ the path should be an absolute path.
	1013	+
	1014	+ The class should have a method get_example_inputs to return the inputs
	1015	+ for the model. An example looks like
	1016	+ ```
	1017	+ class LinearModel(nn.Module):
	1018	+ def __init__(self):
	1019	+ super().__init__()
	1020	+ self.linear = nn.Linear(10, 10)
	1021	+
	1022	+ def forward(self, x):
	1023	+ return self.linear(x)
	1024	+
	1025	+ def get_example_inputs(self):
	1026	+ return (torch.randn(2, 10),)
	1027	+ ```
	1028	+ """,
	1029	+ )
	1030	+ parser.add_argument(
	1031	+ "--multiprocess",
	1032	+ action="store_true",
	1033	+ help="Create n processes based on the number of devices (distributed use case).",
	1034	+ )
	1035	+ parser.add_argument(
	1036	+ "--ddp",
	1037	+ action="store_true",
	1038	+ help="Wraps model in DDP before running it, and uses dynamo DDPOptmizer (graph breaks) by default.",
	1039	+ )
	1040	+ parser.add_argument(
	1041	+ "--distributed-master-port",
	1042	+ default="6789",
	1043	+ help="Port to bind for for torch.distributed. Use the default unless it's conflicting with another user",
	1044	+ )
	1045	+ parser.add_argument(
	1046	+ "--dynamic-shapes",
	1047	+ action="store_true",
	1048	+ help="Runs a dynamic shapes version of the benchmark, if available.",
	1049	+ )
	1050	+ parser.add_argument(
	1051	+ "--dynamic-batch-only",
	1052	+ action="store_true",
	1053	+ help="Only assume batch dimension is dynamic. Implies --dynamic-shapes",
	1054	+ )
	1055	+ parser.add_argument(
	1056	+ "--output",
	1057	+ help="Overrides the output filename",
	1058	+ )
	1059	+ parser.add_argument(
	1060	+ "--output-directory",
	1061	+ help="Overrides the directory to place output files.",
	1062	+ )
	1063	+ parser.add_argument(
	1064	+ "--baseline",
	1065	+ help="Compare with a prior --output",
	1066	+ )
	1067	+ parser.add_argument(
	1068	+ "--part",
	1069	+ default=None,
	1070	+ help="Specify the part of the model to run.",
	1071	+ )
	1072	+ parser.add_argument(
	1073	+ "--export-profiler-trace",
	1074	+ action="store_true",
	1075	+ help="exports trace of kineto profiler",
	1076	+ )
	1077	+ parser.add_argument(
	1078	+ "--profiler-trace-name",
	1079	+ "--profiler_trace_name",
	1080	+ help="Overwrites exported trace name",
	1081	+ )
	1082	+ parser.add_argument(
	1083	+ "--tag", default=None, help="Specify a tag to be included in csv files."
	1084	+ )
	1085	+ parser.add_argument(
	1086	+ "--stats",
	1087	+ action="store_true",
	1088	+ help="print graph counter stats",
	1089	+ )
	1090	+ parser.add_argument(
	1091	+ "--cold-start-latency",
	1092	+ "--cold_start_latency",
	1093	+ action="store_true",
	1094	+ help="Use a fresh triton cachedir when running each model, to force cold-start compile.",
	1095	+ )
	1096	+ parser.add_argument(
	1097	+ "--disable-cudagraphs",
	1098	+ action="store_true",
	1099	+ help="Disables cudagraphs for Inductor",
	1100	+ )
	1101	+ parser.add_argument(
	1102	+ "--disable-split-reductions",
	1103	+ action="store_true",
	1104	+ help="Disables split reductions for Inductor",
	1105	+ )
	1106	+ parser.add_argument(
	1107	+ "--disable-persistent-reductions",
	1108	+ action="store_true",
	1109	+ help="Disables split reductions for Inductor",
	1110	+ )
	1111	+ parser.add_argument(
	1112	+ "--inductor-compile-mode",
	1113	+ default=None,
	1114	+ help="torch.compile mode argument for inductor runs.",
	1115	+ )
	1116	+ parser.add_argument(
	1117	+ "--print-graph-breaks",
	1118	+ action="store_true",
	1119	+ help="Show a warning whenever graph break",
	1120	+ )
	1121	+ parser.add_argument(
	1122	+ "--log-graph-breaks",
	1123	+ action="store_true",
	1124	+ help="log graph breaks in a file",
	1125	+ )
	1126	+ parser.add_argument(
	1127	+ "--collect-outputs",
	1128	+ action="store_true",
	1129	+ help="""Whether to collect outputs for training. Set this to true if we
	1130	+ want to verify the numerical correctness of graidents. But that may
	1131	+ cause time measurement not accurate""",
	1132	+ )
	1133	+
	1134	+ parser.add_argument(
	1135	+ "--timeout",
	1136	+ type=int,
	1137	+ default=2000,
	1138	+ help="timeout (second) for benchmarking.",
	1139	+ )
	1140	+
	1141	+ parser.add_argument(
	1142	+ "--enable-profiler",
	1143	+ action="store_true",
	1144	+ help="Enable profile for NPU and GPU."
	1145	+ )
	1146	+
	1147	+ parser.add_argument(
	1148	+ "--prof-output-path",
	1149	+ help="Overrides the profile output path",
	1150	+ default="./profile"
	1151	+ )
	1152	+
	1153	+ parser.add_argument(
	1154	+ "--dump-compile-time",
	1155	+ action="store_true",
	1156	+ help="dump compile time",
	1157	+ )
	1158	+
	1159	+ group_prec = parser.add_mutually_exclusive_group()
	1160	+ group_prec.add_argument("--float16", action="store_true", help="cast model to fp16")
	1161	+ group_prec.add_argument(
	1162	+ "--bfloat16", action="store_true", help="cast model to bf16"
	1163	+ )
	1164	+ group_prec.add_argument("--float32", action="store_true", help="cast model to fp32")
	1165	+ group_prec.add_argument(
	1166	+ "--amp", action="store_true", help="use automatic mixed precision"
	1167	+ )
	1168	+
	1169	+ group = parser.add_mutually_exclusive_group()
	1170	+ group.add_argument(
	1171	+ "--inductor",
	1172	+ action="store_true",
	1173	+ help="Measure speedup with TorchInductor",
	1174	+ )
	1175	+ group.add_argument(
	1176	+ "--backend",
	1177	+ choices=torch._dynamo.list_backends(exclude_tags=None),
	1178	+ help="measure speedup with a given backend",
	1179	+ )
	1180	+ group.add_argument("--nothing", action="store_true", help="A no-op experiment useful for making sure TorchBenchark alone works properly")
	1181	+
	1182	+ mode_group = parser.add_mutually_exclusive_group(required=True)
	1183	+ mode_group.add_argument(
	1184	+ "--accuracy",
	1185	+ action="store_true",
	1186	+ help="Checks accuracy with small batch size and eval mode",
	1187	+ )
	1188	+ mode_group.add_argument(
	1189	+ "--performance", action="store_true", help="Measures performance speedup"
	1190	+ )
	1191	+ run_mode_group = parser.add_mutually_exclusive_group(required=True)
	1192	+ run_mode_group.add_argument(
	1193	+ "--training",
	1194	+ action="store_true",
	1195	+ help="Performs training",
	1196	+ )
	1197	+ run_mode_group.add_argument(
	1198	+ "--inference", action="store_true", help="Performs inference"
	1199	+ )
	1200	+ return parser.parse_args(args)
	1201	+
	1202	+
	1203	+def process_entry(rank, runner, original_dir, args):
	1204	+ args.rank = rank
	1205	+ with maybe_init_distributed(
	1206	+ args.use_distributed,
	1207	+ rank=rank,
	1208	+ world_size=args.world_size,
	1209	+ port=args.distributed_master_port,
	1210	+ ):
	1211	+ return maybe_fresh_cache(
	1212	+ run, (args.cold_start_latency and args.only)
	1213	+ )(runner, args, original_dir)
	1214	+
	1215	+
	1216	+def main(runner, original_dir=None):
	1217	+ if original_dir:
	1218	+ os.chdir(original_dir)
	1219	+ args = parse_args()
	1220	+ if args.baseline:
	1221	+ args.baseline = os.path.abspath(args.baseline)
	1222	+
	1223	+ if is_npu_available and args.only:
	1224	+ patch_model(args.only)
	1225	+
	1226	+ args.use_distributed = (args.ddp) and args.only
	1227	+ if args.multiprocess:
	1228	+ # NB: Do NOT query device count before CUDA initialization; we're
	1229	+ # going to overwrite CUDA_VISIBLE_DEVICES and this will result in
	1230	+ device_count = torch.cuda.device_count()
	1231	+ if device_count <= 1:
	1232	+ log.warning(
	1233	+ "The use multiprocess flag is set but there are <= 1 devices available."
	1234	+ )
	1235	+ # multiprocess path
	1236	+ args.world_size = device_count
	1237	+ mp.spawn(process_entry, args=(runner, original_dir, args), nprocs=device_count)
	1238	+ else:
	1239	+ # single process path just uses the main process
	1240	+ args.world_size = 1
	1241	+ process_entry(0, runner, original_dir, args)
	1242	+
	1243	+
	1244	+def run(runner, args, original_dir=None):
	1245	+ # Pass the parsed args object to benchmark runner object
	1246	+ runner.args = args
	1247	+
	1248	+ if args.inductor:
	1249	+ if args.backend is not None:
	1250	+ raise AssertionError
	1251	+ args.backend = "inductor"
	1252	+ if args.dynamic_batch_only:
	1253	+ args.dynamic_shapes = True
	1254	+ torch._dynamo.config.assume_static_by_default = True
	1255	+ if args.dynamic_shapes:
	1256	+ if not args.dynamic_batch_only:
	1257	+ torch._dynamo.config.assume_static_by_default = False
	1258	+ if args.ddp:
	1259	+ # but just to measure impact on singlenode of performing graph-breaks.
	1260	+ # Left it as a follow up to keep this PR isolated.
	1261	+ if not args.accuracy:
	1262	+ raise AssertionError("DDP benchmark is currently only hooked up to --accuracy bench")
	1263	+ if not args.training:
	1264	+ raise AssertionError("DDP benchmark requires --training mode")
	1265	+ if args.accuracy:
	1266	+ # Use small batch size. We use >1 batch size to ensure we test
	1267	+ # batch_norm type of operators that work on batch dims.
	1268	+ if args.batch_size is None:
	1269	+ if runner.suite_name == "huggingface":
	1270	+ args.batch_size = 1
	1271	+ elif runner.suite_name == "torchbench":
	1272	+ args.batch_size = 4
	1273	+ else:
	1274	+ # Larger batch size of TIMM models to have stable batch_norm
	1275	+ if runner.suite_name != "timm_models":
	1276	+ raise AssertionError
	1277	+ args.batch_size = 8
	1278	+
	1279	+ # Remove sources of randomness
	1280	+ inductor_config.fallback_random = True
	1281	+ if args.only is not None and args.only not in {
	1282	+ "alexnet",
	1283	+ "Background_Matting",
	1284	+ "pytorch_CycleGAN_and_pix2pix",
	1285	+ "pytorch_unet",
	1286	+ "Super_SloMo",
	1287	+ "vgg16",
	1288	+ "Wav2Vec2ForCTC",
	1289	+ "Wav2Vec2ForPreTraining",
	1290	+ "sam",
	1291	+ }:
	1292	+ # some of the models do not support use_deterministic_algorithms
	1293	+ torch.use_deterministic_algorithms(True)
	1294	+ else:
	1295	+ log.warning("Currently, all models keep deterministic open on npu. "
	1296	+ "But on gpu, this model does not support use_deterministic_algorithms. "
	1297	+ "Please check it to prevent bugs.")
	1298	+ torch.use_deterministic_algorithms(True)
	1299	+
	1300	+ os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
	1301	+ torch.backends.cudnn.deterministic = True
	1302	+ torch.backends.cudnn.allow_tf32 = False
	1303	+ torch.backends.cudnn.benchmark = False
	1304	+ torch.backends.cuda.matmul.allow_tf32 = False
	1305	+
	1306	+ # Remove randomeness when torch manual seed is called
	1307	+ patch_torch_manual_seed()
	1308	+
	1309	+ # Some models e.g. yolov3 assert batch size on n_gpus
	1310	+ if "CUDA_VISIBLE_DEVICES" not in os.environ:
	1311	+ args.device_index = "0"
	1312	+
	1313	+ if args.device_index is not None:
	1314	+ os.environ["CUDA_VISIBLE_DEVICES"] = args.device_index
	1315	+
	1316	+ def __check_if_transfer_to_npu():
	1317	+ return torch.cuda.is_available == is_npu_available
	1318	+
	1319	+ if not args.devices:
	1320	+ if torch.cuda.is_available():
	1321	+ if __check_if_transfer_to_npu():
	1322	+ args.devices = ["npu"]
	1323	+ else:
	1324	+ args.devices = ["cuda"]
	1325	+ elif is_npu_available:
	1326	+ args.devices = ["npu"]
	1327	+ else:
	1328	+ log.warning("torch.cuda.is_available() == False, using CPU")
	1329	+ args.devices = ["cpu"]
	1330	+
	1331	+ if args.devices != ["cpu"] and torch.cuda.is_available():
	1332	+ global synchronize
	1333	+ synchronize = torch.cuda.synchronize
	1334	+ elif is_npu_available:
	1335	+ synchronize = torch_npu.npu.synchronize
	1336	+
	1337	+ if (
	1338	+ args.devices == ["cuda"]
	1339	+ and torch.cuda.get_device_properties(0).total_memory < 25 * 2**30
	1340	+ ):
	1341	+ # OOM errors on an RTX 3090 with 24gb RAM
	1342	+ runner.skip_models.update(
	1343	+ {
	1344	+ # torchbench
	1345	+ "hf_Longformer",
	1346	+ "timm_nfnet",
	1347	+ "timm_efficientdet",
	1348	+ }
	1349	+ )
	1350	+ if args.training:
	1351	+ runner.skip_models.add("hf_T5")
	1352	+
	1353	+ if args.print_graph_breaks:
	1354	+ torch._dynamo.config.print_graph_breaks = True
	1355	+
	1356	+ if args.training:
	1357	+ runner.model_iter_fn = runner.forward_and_backward_pass
	1358	+ else:
	1359	+ runner.model_iter_fn = runner.forward_pass
	1360	+
	1361	+ if args.devices == ["cpu"]:
	1362	+ runner.skip_models.update(runner.very_slow_models)
	1363	+
	1364	+ if args.no_skip:
	1365	+ runner.skip_models.clear()
	1366	+
	1367	+ experiment = null_experiment
	1368	+ global current_name, current_device, current_batch_size, output_filename
	1369	+ optimize_ctx = contextlib.nullcontext()
	1370	+
	1371	+ if args.inductor:
	1372	+ optimize_ctx = functools.partial(
	1373	+ torch.compile,
	1374	+ backend="inductor",
	1375	+ fullgraph=args.nopython,
	1376	+ mode=args.inductor_compile_mode,
	1377	+ )
	1378	+ experiment = speedup_experiment
	1379	+ output_filename = "inductor.csv"
	1380	+ elif args.nothing:
	1381	+ optimize_ctx = nothing
	1382	+ experiment = speedup_experiment
	1383	+ output_filename = "nothing.csv"
	1384	+ elif args.backend:
	1385	+ optimize_ctx = compile_with_backend(args)
	1386	+ experiment = speedup_experiment
	1387	+ if args.accuracy:
	1388	+ output_filename = f"accuracy_{args.backend}.csv"
	1389	+ else:
	1390	+ output_filename = f"speedup_{args.backend}.csv"
	1391	+
	1392	+ if args.only is not None and args.only not in {"hf_Bart", "torch_multimodal_clip", "timm_vision_transformer"}:
	1393	+ if args.inductor or args.backend == "inductor" or args.export_aot_inductor:
	1394	+ inductor_config.triton.cudagraphs = not args.disable_cudagraphs
	1395	+
	1396	+ runner.setup_amp()
	1397	+
	1398	+ if args.output:
	1399	+ output_filename = args.output
	1400	+
	1401	+ if output_filename:
	1402	+ if args.output_directory:
	1403	+ output_filename = os.path.join(args.output_directory, output_filename)
	1404	+ else:
	1405	+ output_filename = os.path.join(
	1406	+ # pytorch use torch._dynamo.config.base_dir originally,
	1407	+ # but the generated file will be saved under directory where pytorch was installed,
	1408	+ # change the default output saved directory.
	1409	+ os.path.dirname(os.path.abspath(__file__)), output_filename
	1410	+ )
	1411	+
	1412	+ if args.export_profiler_trace:
	1413	+ if args.profiler_trace_name is None:
	1414	+ if args.backend:
	1415	+ args.profiler_trace_name = args.backend
	1416	+ elif args.inductor:
	1417	+ args.profiler_trace_name = "inductor"
	1418	+ else:
	1419	+ args.profiler_trace_name = "profile"
	1420	+ else:
	1421	+ args.profiler_trace_name = args.profiler_trace_name
	1422	+
	1423	+ experiment = functools.partial(experiment, args, runner.model_iter_fn)
	1424	+
	1425	+ if args.only:
	1426	+ # use aclnn by default, otherwise compared with aclop
	1427	+ if os.environ.get("USE_ACLOP", "0").upper() in ["1", "ON"]:
	1428	+ torch_npu.npu.set_compile_mode(jit_compile=True)
	1429	+
	1430	+ model_name = args.only
	1431	+ for device in args.devices:
	1432	+ batch_size = args.batch_size
	1433	+ if args.batch_size_file:
	1434	+ batch_size = read_batch_size_from_file(
	1435	+ args, args.batch_size_file, model_name
	1436	+ )
	1437	+ if model_specified_by_path(args.only):
	1438	+ model, example_inputs = load_model_from_path(args.only)
	1439	+ name = model.__class__.__name__
	1440	+ model = model.to(device=device)
	1441	+ example_inputs = tree_map_only(
	1442	+ torch.Tensor, lambda x: x.to(device=device), example_inputs
	1443	+ )
	1444	+ else:
	1445	+ try:
	1446	+ with tqdm(desc="loading model"):
	1447	+ if args.part:
	1448	+ (
	1449	+ device,
	1450	+ name,
	1451	+ model,
	1452	+ example_inputs,
	1453	+ batch_size,
	1454	+ ) = runner.load_model(
	1455	+ device,
	1456	+ model_name,
	1457	+ batch_size=batch_size,
	1458	+ part=args.part,
	1459	+ )
	1460	+ else:
	1461	+ (
	1462	+ device,
	1463	+ name,
	1464	+ model,
	1465	+ example_inputs,
	1466	+ batch_size,
	1467	+ ) = runner.load_model(
	1468	+ device, model_name, batch_size=batch_size
	1469	+ )
	1470	+ except NotImplementedError as e:
	1471	+ print(e)
	1472	+ import traceback
	1473	+
	1474	+ print(traceback.format_exc())
	1475	+ logging.warning("%s failed to load", args.only)
	1476	+ continue # bad benchmark implementation
	1477	+
	1478	+ current_name = name
	1479	+ current_device = device
	1480	+ current_batch_size = batch_size
	1481	+ set_model_name(name)
	1482	+
	1483	+ # Look for stuff that looks like batch size, and mark it dynamic.
	1484	+ # Better integration would integrate directly with benchmark suite
	1485	+ # but cannot conveniently do this
	1486	+ # NB: This must be done late enough so that we don't do more
	1487	+ # conversions on the inputs
	1488	+ # NB: Assumes only the first batch-y like dimension is the batch
	1489	+ marked = False
	1490	+
	1491	+ def detect_and_mark_batch(t, target_size=batch_size):
	1492	+ nonlocal marked
	1493	+ for i, s in enumerate(t.size()):
	1494	+ if s == target_size:
	1495	+ torch._dynamo.mark_dynamic(t, i)
	1496	+ marked = True
	1497	+ break
	1498	+
	1499	+ if (
	1500	+ args.dynamic_batch_only
	1501	+ and batch_size > 1
	1502	+ ):
	1503	+ tree_map_only(torch.Tensor, detect_and_mark_batch, example_inputs)
	1504	+ if not marked:
	1505	+ raise AssertionError(f"nothing in example_inputs had a dim with {batch_size}")
	1506	+
	1507	+ model, example_inputs = runner.cast_based_on_args(model, example_inputs)
	1508	+ runner.run_one_model(
	1509	+ name,
	1510	+ model,
	1511	+ example_inputs,
	1512	+ optimize_ctx,
	1513	+ experiment,
	1514	+ tag=args.tag,
	1515	+ )
	1516	+ # exec callback functions registered in npu_support.py
	1517	+ for fn in callbacks:
	1518	+ fn()
	1519	+ else:
	1520	+ if output_filename and os.path.exists(output_filename):
	1521	+ os.unlink(output_filename)
	1522	+ if original_dir:
	1523	+ os.chdir(original_dir)
	1524	+ model_names = list(runner.iter_model_names(args))
	1525	+ nmodels = len(model_names)
	1526	+ for i, name in enumerate(model_names):
	1527	+ current_name = name
	1528	+ placeholder_batch_size = 0
	1529	+ print(f"Running model {i+1}/{nmodels}", flush=True)
	1530	+
	1531	+ def write_csv(status, name=name, placeholder_batch_size=placeholder_batch_size):
	1532	+ if args.accuracy:
	1533	+ headers = ["dev", "name", "batch_size", "accuracy"]
	1534	+ rows = [
	1535	+ [device, name, placeholder_batch_size, status]
	1536	+ for device in args.devices
	1537	+ ]
	1538	+ elif args.performance:
	1539	+ headers = ["dev", "name", "batch_size", "speedup", "abs_latency"]
	1540	+ rows = [
	1541	+ [device, name, placeholder_batch_size, 0.0, 0.0]
	1542	+ for device in args.devices
	1543	+ ]
	1544	+ else:
	1545	+ headers = []
	1546	+ rows = [
	1547	+ [device, name, placeholder_batch_size, 0.0]
	1548	+ for device in args.devices
	1549	+ ]
	1550	+
	1551	+ for row in rows:
	1552	+ output_csv(output_filename, headers, row)
	1553	+
	1554	+ try:
	1555	+ subprocess.check_call(
	1556	+ [sys.executable] + sys.argv + [f"--only={name}"], timeout=args.timeout
	1557	+ )
	1558	+ except subprocess.TimeoutExpired:
	1559	+ print("TIMEOUT", file=sys.stderr)
	1560	+ write_csv("timeout")
	1561	+ except subprocess.SubprocessError:
	1562	+ print("ERROR", file=sys.stderr)
	1563	+ write_csv("infra_error")
	1564	+
	1565	+
	1566	+def compile_with_backend(args):
	1567	+ return torch._dynamo.optimize(args.backend, nopython=args.nopython)
	1568	+
	1569	+
	1570	+if __name__ == "__main__":
	1571	+ raise RuntimeError(
	1572	+ f"You shouldn't run {sys.argv[0]} directly, instead try torchbench.py"
	1573	+ )

Abenchmarks/extract_log.py+196-0

@@ -0,0 +1,196 @@
	1	+import argparse
	2	+from collections import defaultdict
	3	+from pathlib import Path
	4	+import re
	5	+import pandas as pd
	6	+
	7	+
	8	+def extract_log_info(log_file, profile_dir='./profile', output_file='log_analysis.xlsx'):
	9	+ """
	10	+ 提取日志文件中各模型的训练时间信息和profile数据
	11	+
	12	+ Args:
	13	+ log_file: 日志文件路径
	14	+ profile_dir: profile数据目录路径
	15	+ output_file: 输出Excel文件路径
	16	+ """
	17	+ # 模式匹配
	18	+ # 修改1: 同时支持npu和cuda
	19	+ model_pattern = r'(npu\|cuda)\s+train\s+(\S+)'
	20	+ eager_pattern = r'eager.avg step time:\s([\d\.]+)\s*ms'
	21	+ compile_pattern = r'compile.avg step time:\s([\d\.]+)\s*ms'
	22	+
	23	+ # 模式匹配：算子编译时间
	24	+ op_compile_time_pattern = r'op_compile_time:\s([\d\.]+)\sms'
	25	+
	26	+ # 存储结果
	27	+ data = defaultdict(lambda: {
	28	+ 'accuracy': None, # 修改2: 存储完整的精度校验日志
	29	+ 'eager_E2E_avg_time': None,
	30	+ 'compile_E2E_avg_time': None,
	31	+ 'op_compile_time': None,
	32	+ 'eager_OP_avg_time': None,
	33	+ 'compile_OP_avg_time': None
	34	+ })
	35	+ current_model = None
	36	+ in_compile_block = False
	37	+ compile_block_lines = []
	38	+
	39	+ with open(log_file, 'r', encoding='utf-8') as f:
	40	+ lines = f.readlines()
	41	+
	42	+ for _, line in enumerate(lines):
	43	+ # 匹配模型名
	44	+ model_match = re.search(model_pattern, line)
	45	+ if model_match:
	46	+ # 处理上一个模型的compile块日志
	47	+ if current_model and in_compile_block and compile_block_lines:
	48	+ # 提取pass_accuracy日志
	49	+ for log_line in compile_block_lines:
	50	+ if 'pass_accuracy' in log_line:
	51	+ data[current_model]['accuracy'] = log_line.strip()
	52	+ break
	53	+
	54	+ # 提取op_compile_time
	55	+ for log_line in compile_block_lines:
	56	+ op_compile_match = re.search(op_compile_time_pattern, log_line)
	57	+ if op_compile_match:
	58	+ data[current_model]['op_compile_time'] = float(op_compile_match.group(1))
	59	+ break
	60	+
	61	+ # 重置状态
	62	+ current_model = model_match.group(2) # 第二个分组是模型名
	63	+ in_compile_block = False
	64	+ compile_block_lines = []
	65	+ continue
	66	+
	67	+ # 匹配eager模式时间
	68	+ if current_model:
	69	+ eager_match = re.search(eager_pattern, line)
	70	+ if eager_match:
	71	+ data[current_model]['eager_E2E_avg_time'] = float(eager_match.group(1))
	72	+
	73	+ # 匹配compile模式时间
	74	+ compile_match = re.search(compile_pattern, line)
	75	+ if compile_match:
	76	+ data[current_model]['compile_E2E_avg_time'] = float(compile_match.group(1))
	77	+ in_compile_block = True
	78	+ compile_block_lines = [] # 开始收集compile块日志
	79	+
	80	+ # 收集compile块的日志
	81	+ if current_model and in_compile_block:
	82	+ compile_block_lines.append(line)
	83	+
	84	+ # 处理最后一个模型的compile块日志
	85	+ if current_model and in_compile_block and compile_block_lines:
	86	+ for log_line in compile_block_lines:
	87	+ if 'pass_accuracy' in log_line:
	88	+ data[current_model]['accuracy'] = log_line.strip()
	89	+ break
	90	+
	91	+ for log_line in compile_block_lines:
	92	+ op_compile_match = re.search(op_compile_time_pattern, log_line)
	93	+ if op_compile_match:
	94	+ data[current_model]['op_compile_time'] = float(op_compile_match.group(1))
	95	+ break
	96	+
	97	+ # 需求3: 读取profile目录中的step_trace_time.csv文件
	98	+ profile_path = Path(profile_dir)
	99	+ if profile_path.exists():
	100	+ for model_dir in profile_path.iterdir():
	101	+ if model_dir.is_dir():
	102	+ model_name = model_dir.name
	103	+
	104	+ # 读取eager模式下的step_trace_time.csv
	105	+ # 修改3: 自动获取下一级目录
	106	+ eager_dir = model_dir / 'eager'
	107	+ if eager_dir.exists() and eager_dir.is_dir():
	108	+ # 获取eager目录下的第一个子目录
	109	+ eager_subdirs = list(eager_dir.iterdir())
	110	+ if eager_subdirs:
	111	+ eager_subdir = eager_subdirs[0] # 假设只有一个子目录
	112	+ eager_csv_path = eager_subdir / 'ASCEND_PROFILER_OUTPUT' / 'step_trace_time.csv'
	113	+ if eager_csv_path.exists():
	114	+ try:
	115	+ eager_df = pd.read_csv(eager_csv_path, sep=',')
	116	+ if 'Computing' in eager_df.columns:
	117	+ data[model_name]['eager_OP_avg_time'] = eager_df['Computing'].mean() / 1e3
	118	+ else:
	119	+ print(f"警告: {eager_csv_path} 中没有Computing列")
	120	+ except Exception as e:
	121	+ print(f"读取{eager_csv_path}时出错: {e}")
	122	+
	123	+ # 读取compile模式下的step_trace_time.csv
	124	+ compile_dir = model_dir / 'compile'
	125	+ if compile_dir.exists() and compile_dir.is_dir():
	126	+ # 获取compile目录下的第一个子目录
	127	+ compile_subdirs = list(compile_dir.iterdir())
	128	+ if compile_subdirs:
	129	+ compile_subdir = compile_subdirs[0] # 假设只有一个子目录
	130	+ compile_csv_path = compile_subdir / 'ASCEND_PROFILER_OUTPUT' / 'step_trace_time.csv'
	131	+ if compile_csv_path.exists():
	132	+ try:
	133	+ compile_df = pd.read_csv(compile_csv_path, sep=',')
	134	+ if 'Computing' in compile_df.columns:
	135	+ data[model_name]['compile_OP_avg_time'] = compile_df['Computing'].mean() / 1e3
	136	+ else:
	137	+ print(f"警告: {compile_csv_path} 中没有Computing列")
	138	+ except Exception as e:
	139	+ print(f"读取{compile_csv_path}时出错: {e}")
	140	+ else:
	141	+ print(f"警告: profile目录不存在: {profile_dir}")
	142	+
	143	+ # 转换为DataFrame
	144	+ df = pd.DataFrame.from_dict(data, orient='index')
	145	+ df.index.name = 'model_name'
	146	+ df.reset_index(inplace=True)
	147	+
	148	+ # 计算速度提升率
	149	+ # 1. 计算E2E_speed_up_rate = eager_E2E_avg_time / compile_E2E_avg_time
	150	+ # 2. 计算OP_speed_up_rate = eager_OP_avg_time / compile_OP_avg_time
	151	+ df['E2E_speed_up_rate'] = df.apply(
	152	+ lambda row: row['eager_E2E_avg_time'] / row['compile_E2E_avg_time']
	153	+ if row['compile_E2E_avg_time'] and row['compile_E2E_avg_time'] != 0 else None,
	154	+ axis=1
	155	+ )
	156	+
	157	+ df['OP_speed_up_rate'] = df.apply(
	158	+ lambda row: row['eager_OP_avg_time'] / row['compile_OP_avg_time']
	159	+ if row['compile_OP_avg_time'] and row['compile_OP_avg_time'] != 0 else None,
	160	+ axis=1
	161	+ )
	162	+
	163	+ # 重排列顺序，使相关列更清晰
	164	+ column_order = [
	165	+ 'model_name', 'accuracy', 'op_compile_time',
	166	+ 'eager_E2E_avg_time', 'compile_E2E_avg_time', 'E2E_speed_up_rate',
	167	+ 'eager_OP_avg_time', 'compile_OP_avg_time', 'OP_speed_up_rate'
	168	+ ]
	169	+ existing_columns = [col for col in column_order if col in df.columns]
	170	+ df = df[existing_columns + [col for col in df.columns if col not in existing_columns]]
	171	+
	172	+ # 保存到Excel
	173	+ df.to_excel(output_file, index=False)
	174	+ return df
	175	+
	176	+
	177	+def main():
	178	+ parser = argparse.ArgumentParser(description='提取日志文件中的训练时间信息和profile数据')
	179	+ parser.add_argument('--log_file', required=True, help='日志文件路径')
	180	+ parser.add_argument('--profile_dir', default='./profile', help='profile数据目录路径，默认为./profile')
	181	+ parser.add_argument('--output_file', default='analysis.xlsx', help='输出Excel文件路径，默认为log_analysis.xlsx')
	182	+
	183	+ args = parser.parse_args()
	184	+
	185	+ result = extract_log_info(
	186	+ log_file=args.log_file,
	187	+ profile_dir=args.profile_dir,
	188	+ output_file=args.output_file
	189	+ )
	190	+ print(f"提取完成，共找到{len(result)}个模型")
	191	+ print("\n提取结果:")
	192	+ print(result.head())
	193	+
	194	+# 使用示例
	195	+if __name__ == "__main__":
	196	+ main()

Abenchmarks/npu_support.py+592-0

@@ -0,0 +1,592 @@
	1	+import importlib
	2	+import json
	3	+import logging
	4	+import os
	5	+import sys
	6	+from typing import Optional, Tuple, Type, Union
	7	+
	8	+import torch
	9	+import torch.nn as nn
	10	+import torchair
	11	+import common
	12	+import torch_npu
	13	+from torch_npu.dynamo.torchair._utils.path_manager import PathManager
	14	+
	15	+log = logging.getLogger(__name__)
	16	+_patch_table = {}
	17	+
	18	+
	19	+def register_patch(*model_names):
	20	+ def meta_decorator(fn):
	21	+ for model_name in model_names:
	22	+ _patch_table[model_name] = fn
	23	+ return fn
	24	+
	25	+ return meta_decorator
	26	+
	27	+
	28	+def check_transformers_version(required_version):
	29	+ import transformers
	30	+ if transformers.__version__ != required_version:
	31	+ log.warning(f"transformers.__version__ is not equal to {required_version}, which may cause error patch.")
	32	+
	33	+
	34	+def use_aclnn():
	35	+ os.environ["USE_ACLOP"] = "0"
	36	+
	37	+
	38	+def _hf_t5_mt5_conditionalgeneration_forward_new(
	39	+ self,
	40	+ hidden_states,
	41	+ mask=None,
	42	+ key_value_states=None,
	43	+ position_bias=None,
	44	+ past_key_value=None,
	45	+ layer_head_mask=None,
	46	+ query_length=None,
	47	+ use_cache=False,
	48	+ output_attentions=False,
	49	+):
	50	+ batch_size, seq_length = hidden_states.shape[:2]
	51	+
	52	+ real_seq_length = seq_length
	53	+
	54	+ if past_key_value is not None:
	55	+ if len(past_key_value) != 2:
	56	+ raise ValueError(f"past_key_value should have 2 past states. Got {len(past_key_value)} past states")
	57	+ real_seq_length += past_key_value[0].shape[2] if query_length is None else query_length
	58	+
	59	+ key_length = real_seq_length if key_value_states is None else key_value_states.shape[1]
	60	+
	61	+ def shape(states):
	62	+ return states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)
	63	+
	64	+ def unshape(states):
	65	+ return states.transpose(1, 2).contiguous().view(batch_size, -1, self.inner_dim)
	66	+
	67	+ def project(hidden_states, proj_layer, key_value_states, past_key_value):
	68	+ if key_value_states is None:
	69	+ hidden_states = shape(proj_layer(hidden_states))
	70	+ elif past_key_value is None:
	71	+ hidden_states = shape(proj_layer(key_value_states))
	72	+
	73	+ if past_key_value is not None:
	74	+ if key_value_states is None:
	75	+ hidden_states = torch.cat([past_key_value, hidden_states], dim=2)
	76	+ elif past_key_value.shape[2] != key_value_states.shape[1]:
	77	+ hidden_states = shape(proj_layer(key_value_states))
	78	+ else:
	79	+ hidden_states = past_key_value
	80	+ return hidden_states
	81	+
	82	+ query_states = shape(self.q(hidden_states))
	83	+
	84	+ key_states = project(
	85	+ hidden_states, self.k, key_value_states, past_key_value[0] if past_key_value is not None else None
	86	+ )
	87	+ value_states = project(
	88	+ hidden_states, self.v, key_value_states, past_key_value[1] if past_key_value is not None else None
	89	+ )
	90	+
	91	+ scores = torch.matmul(query_states, key_states.transpose(3, 2))
	92	+
	93	+ def process_position_bias():
	94	+ if not self.has_relative_attention_bias:
	95	+ position_bias = torch.zeros(
	96	+ (1, self.n_heads, real_seq_length, key_length), device=scores.device, dtype=scores.dtype
	97	+ )
	98	+ if self.gradient_checkpointing and self.training:
	99	+ position_bias.requires_grad = True
	100	+ else:
	101	+ position_bias = self.compute_bias(real_seq_length, key_length, device=scores.device)
	102	+
	103	+ if past_key_value is not None:
	104	+ position_bias = position_bias[:, :, -hidden_states.size(1):, :]
	105	+
	106	+ if mask is not None:
	107	+ position_bias = position_bias + mask
	108	+ return position_bias
	109	+
	110	+ if position_bias is None:
	111	+ position_bias = process_position_bias()
	112	+
	113	+ if self.pruned_heads:
	114	+ mask = torch.ones(position_bias.shape[1])
	115	+ mask[list(self.pruned_heads)] = 0
	116	+ position_bias_masked = position_bias[:, mask.bool()]
	117	+ else:
	118	+ position_bias_masked = position_bias
	119	+
	120	+ # Only patch here, src code: [scores += position_bias_masked]
	121	+ # Prevent from two continuous _to_copy.
	122	+ scores = scores.float() + position_bias_masked
	123	+
	124	+ attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(scores)
	125	+
	126	+ attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
	127	+
	128	+ if layer_head_mask is not None:
	129	+ attn_weights = attn_weights * layer_head_mask
	130	+
	131	+ attn_output = unshape(torch.matmul(attn_weights, value_states))
	132	+ attn_output = self.o(attn_output)
	133	+
	134	+ present_key_value_state = (key_states, value_states) if (self.is_decoder and use_cache) else None
	135	+ outputs = (attn_output,) + (present_key_value_state,) + (position_bias,)
	136	+
	137	+ if output_attentions:
	138	+ outputs = outputs + (attn_weights,)
	139	+ return outputs
	140	+
	141	+
	142	+@register_patch("LearningToPaint")
	143	+def _patch_model_1():
	144	+ # For model LearningToPaint.
	145	+ from torchbenchmark.models import LearningToPaint
	146	+ USE_DEVICE = torch.cuda.is_available() or torch_npu.npu.is_available()
	147	+ LearningToPaint.baseline.utils.util.USE_CUDA = USE_DEVICE
	148	+
	149	+
	150	+@register_patch("hf_T5", "hf_T5_base")
	151	+def _patch_model_3():
	152	+ # For model hf_T5 and hf_T5_base.
	153	+ # In these models, accuracy check will fail because in the model's block [T5Attention],
	154	+ # two continuous _to_copy are invoked: the first _to_copy converts Tensor to half
	155	+ # and the second converts it to float. In eager, there will be a loss of precision.
	156	+ # But in graph, there will be a fusion pass to prevent it happens, causing acc check fail.
	157	+ try:
	158	+ from transformers.models.t5.modeling_t5 import T5Attention
	159	+ except ImportError:
	160	+ log.warning("Import transformers failed or could not get T5Attention "
	161	+ "from module transformers.models.t5.modeling_t5")
	162	+ return
	163	+ check_transformers_version("4.36.0")
	164	+
	165	+ T5Attention.forward = _hf_t5_mt5_conditionalgeneration_forward_new
	166	+
	167	+
	168	+@register_patch("fastNLP_Bert")
	169	+def _patch_model_5():
	170	+ os.environ['BREAK_GRAPH_OP_LIST'] = 'NN.LINEAR'
	171	+ # None-public interface, just for test.
	172	+ # This env is added after torchair's init,
	173	+ # so need to call break_graph patch again,
	174	+ torchair._utils.npu_patch_break_graph()
	175	+
	176	+
	177	+@register_patch("hf_Longformer")
	178	+def _patch_model_6():
	179	+ """
	180	+ Hf_Longformer failed accurazy test because of discontiguous memory.
	181	+ Solving the problem by adding .contiguous() after .view() and .as_strided in LongformerSelfAttention._chunk.
	182	+ This patch would be removed in the near future.
	183	+ """
	184	+ # close AddLayerNormFusionPass
	185	+ close_view_optimise()
	186	+ module_spec = importlib.util.find_spec("transformers")
	187	+ if module_spec is None:
	188	+ return
	189	+ from transformers.models.longformer import LongformerSelfAttention
	190	+ src_chunk = LongformerSelfAttention._chunk
	191	+
	192	+ def _chunk(cls, hidden_states, window_overlap, onnx_export: bool = False):
	193	+ if not onnx_export:
	194	+ hidden_states = hidden_states.view(
	195	+ hidden_states.size(0),
	196	+ torch.div(hidden_states.size(1), (window_overlap * 2), rounding_mode="trunc"),
	197	+ window_overlap * 2,
	198	+ hidden_states.size(2),
	199	+ ).contiguous()
	200	+ chunk_size = list(hidden_states.size())
	201	+ chunk_size[1] = chunk_size[1] * 2 - 1
	202	+
	203	+ chunk_stride = list(hidden_states.stride())
	204	+ chunk_stride[1] = chunk_stride[1] // 2
	205	+ return hidden_states.as_strided(size=chunk_size, stride=chunk_stride).contiguous()
	206	+ return src_chunk(hidden_states, window_overlap, True)
	207	+
	208	+ LongformerSelfAttention._chunk = _chunk
	209	+
	210	+
	211	+@register_patch("soft_actor_critic")
	212	+def _patch_model_7():
	213	+ """
	214	+ soft_actor_critic failed accurazy test because of discontiguous memory.
	215	+ Solving the problem by adding .contiguous() in soft_actor_critic/net.py line:242 SquashedNormal.__init__
	216	+ This patch would be removed in the near future.
	217	+ """
	218	+ from torchbenchmark.models.soft_actor_critic.nets import StochasticActor, SquashedNormal, BetaDist
	219	+ import torch.nn.functional as F
	220	+
	221	+ def new_forward(self, state):
	222	+ x = F.relu(self.fc1(state))
	223	+ x = F.relu(self.fc2(x))
	224	+ out = self.fc3(x)
	225	+ mu, log_std = out.chunk(2, dim=1)
	226	+ if self.dist_impl == "pyd":
	227	+ log_std = torch.tanh(log_std)
	228	+ log_std = self.log_std_low + 0.5 * (
	229	+ self.log_std_high - self.log_std_low
	230	+ ) * (log_std + 1)
	231	+ std = log_std.exp()
	232	+ dist = SquashedNormal(mu.contiguous(), std.contiguous())
	233	+ elif self.dist_impl == "beta":
	234	+ out = 1.0 + F.softplus(out)
	235	+ alpha, beta = out.chunk(2, dim=1)
	236	+ dist = BetaDist(alpha, beta)
	237	+ return dist
	238	+
	239	+ StochasticActor.forward = new_forward
	240	+
	241	+
	242	+@register_patch("dcgan", "mobilenet_v2", "phlippe_resnet", "shufflenet_v2_x1_0", "squeezenet1_1", "vgg16",
	243	+ "alexnet", "densenet121", "maml_omniglot")
	244	+def _patch_model_8():
	245	+ """
	246	+ close conv amp for some model only in accuracy mode.
	247	+ This patch would be removed in the near future.
	248	+ """
	249	+ if {"--only", "--amp", "--accuracy"} <= set(sys.argv):
	250	+ from torch.nn.modules.conv import Conv2d
	251	+
	252	+ def conv2d_amp_disabled(self, x):
	253	+ with torch.npu.amp.autocast(enabled=False):
	254	+ return self._conv_forward(x, self.weight, self.bias)
	255	+
	256	+ Conv2d.forward = conv2d_amp_disabled
	257	+
	258	+
	259	+@register_patch("timm_nfnet")
	260	+def _patch_model_9():
	261	+ # close conv amp for timm_nfnet only in accuracy mode.
	262	+ # Increase the batch_size to a larger size 16,
	263	+ # to mitigate the impact of BatchNorm's tolerance on convolution
	264	+ if {"--only", "--amp", "--accuracy"} <= set(sys.argv):
	265	+ try:
	266	+ import timm
	267	+ import torch.nn.functional as F
	268	+ from timm.layers.std_conv import ScaledStdConv2dSame
	269	+ from timm.layers.padding import pad_same
	270	+ except ImportError:
	271	+ log.warning("Import timm failed or could not get ScaledStdConv2dSame"
	272	+ "from module timm.layers.std_conv.ScaledStdConv2dSame")
	273	+ return
	274	+ if timm.__version__ != '0.9.16':
	275	+ log.warning("timm.__version__ is not equal to 0.9.16, which may cause error patch.")
	276	+
	277	+ def new_forward(self, x):
	278	+ if self.same_pad:
	279	+ x = pad_same(x, self.kernel_size, self.stride, self.dilation)
	280	+ weight = F.batch_norm(
	281	+ self.weight.reshape(1, self.out_channels, -1), None, None,
	282	+ weight=(self.gain * self.scale).view(-1),
	283	+ training=True, momentum=0., eps=self.eps).reshape_as(self.weight)
	284	+ with torch.npu.amp.autocast(enabled=False):
	285	+ return F.conv2d(x, weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
	286	+
	287	+ ScaledStdConv2dSame.forward = new_forward
	288	+
	289	+ try:
	290	+ from torchbenchmark.models.timm_nfnet import Model
	291	+ except ImportError:
	292	+ log.warning("Import Model failed or could not find timm_nfnet"
	293	+ "from module torchbenchmark.models.timm_nfnet.Model")
	294	+ return
	295	+
	296	+ def new__init(self, test, device, jit=False, batch_size=None, extra_args=None):
	297	+ super(Model, self).__init__(test=test, model_name='dm_nfnet_f0',
	298	+ device=device, batch_size=16, extra_args=extra_args)
	299	+
	300	+ Model.__init__ = new__init
	301	+
	302	+
	303	+@register_patch("nvidia_deeprecommender")
	304	+def _patch_model_10():
	305	+ try:
	306	+ from torch_npu.contrib import transfer_to_npu
	307	+ except ImportError:
	308	+ log.warning("NPU_FlAG is False!")
	309	+ return
	310	+
	311	+ try:
	312	+ from torchbenchmark.models.nvidia_deeprecommender.nvtrain import DeepRecommenderTrainBenchmark
	313	+ except ImportError:
	314	+ log.warning("Import nvidia_deeprecommender failed or could not get DeepRecommenderTrainBenchmark"
	315	+ "from module torchbenchmark.models.nvidia_deeprecommender.nvtrain.DeepRecommenderTrainBenchmark")
	316	+ return
	317	+
	318	+ def new_init(self, device="cpu", jit=False, batch_size=256, process_command_line=False):
	319	+ self.TrainInit("cuda", jit, batch_size, process_command_line)
	320	+
	321	+ DeepRecommenderTrainBenchmark.__init__ = new_init
	322	+
	323	+
	324	+@register_patch("functorch_dp_cifar10")
	325	+def _patch_model_11():
	326	+ if {"--only", "--amp", "--accuracy"} <= set(sys.argv):
	327	+ try:
	328	+ from torchbenchmark.models.functorch_dp_cifar10 import Model
	329	+ import torchvision.models as models
	330	+ except ImportError:
	331	+ log.warning("import torchvision fail or could not get Model from module "
	332	+ "torchbenchmark.models.functorch_dp_cifar10")
	333	+ return
	334	+
	335	+ def new_init(self, test, device, batch_size=None, extra_args=None):
	336	+ if extra_args is None:
	337	+ extra_args = []
	338	+ super(Model, self).__init__(test=test, device=device, batch_size=32, extra_args=extra_args)
	339	+ self.model = models.__dict__['resnet18'](
	340	+ pretrained=False, norm_layer=(lambda c: nn.GroupNorm(min(c, 32), c)))
	341	+ self.model = self.model.to(device)
	342	+ self.example_inputs = (
	343	+ torch.randn((self.batch_size, 3, 32, 32), device=self.device),
	344	+ )
	345	+ self.example_target = torch.randint(0, 10, (self.batch_size,), device=self.device)
	346	+ self.optimizer = torch.optim.Adam(self.model.parameters(), lr=0.001)
	347	+ self.criterion = nn.CrossEntropyLoss()
	348	+
	349	+ Model.__init__ = new_init
	350	+
	351	+
	352	+def create_fusion_switch_file():
	353	+ fusion_config = {}
	354	+ fusion_config.setdefault("Switch", {}).setdefault("GraphFusion", {})["AddLayerNormFusionPass"] = "off"
	355	+ fusion_config_file = os.path.join(os.getcwd(), "fusion_switch.cfg")
	356	+ PathManager.check_path_writeable_and_safety(fusion_config_file)
	357	+ with os.fdopen(os.open(fusion_config_file, os.O_WRONLY \| os.O_CREAT, mode=600), 'w') as f:
	358	+ json.dump(fusion_config, f)
	359	+ config = torchair.CompilerConfig()
	360	+ config.fusion_config.fusion_switch_file = fusion_config_file
	361	+
	362	+ def clean_fusion_config_file():
	363	+ PathManager.remove_file_safety(fusion_config_file)
	364	+
	365	+ from common import register_callback
	366	+ register_callback(clean_fusion_config_file)
	367	+ return config
	368	+
	369	+
	370	+def close_view_optimise():
	371	+ config = create_fusion_switch_file()
	372	+ config.experimental_config.enable_view_optimize = False
	373	+ npu_backend = torchair.get_npu_backend(compiler_config=config)
	374	+
	375	+ def compile_with_view_switch(args):
	376	+ return torch._dynamo.optimize(npu_backend, nopython=args.nopython)
	377	+ common.compile_with_backend = compile_with_view_switch
	378	+
	379	+
	380	+def close_add_layer_norm_fusion_pass():
	381	+ npu_backend = torchair.get_npu_backend(compiler_config=create_fusion_switch_file())
	382	+
	383	+ def compile_with_fusion_switch(args):
	384	+ return torch._dynamo.optimize(npu_backend, nopython=args.nopython)
	385	+ common.compile_with_backend = compile_with_fusion_switch
	386	+
	387	+
	388	+@register_patch("moco")
	389	+def _patch_model_13():
	390	+ from argparse import Namespace
	391	+ import torch.distributed as dist
	392	+
	393	+ try:
	394	+ from torch_npu.contrib import transfer_to_npu
	395	+ except ImportError:
	396	+ log.warning("NPU_FlAG is False!")
	397	+ return
	398	+
	399	+ try:
	400	+ import torchvision.models as models
	401	+ from torchbenchmark.models.moco import Model
	402	+ from torchbenchmark.models.moco.moco.builder import MoCo
	403	+ except ImportError:
	404	+ log.warning("import torchvision fail or could not get Model,MoCo from module torchbenchmark.models.moco ")
	405	+ return
	406	+
	407	+ def new_init(self, test, device, batch_size=None, extra_args=None):
	408	+ if extra_args is None:
	409	+ extra_args = []
	410	+ super(Model, self).__init__(test=test, device=device, batch_size=batch_size, extra_args=extra_args)
	411	+ self.opt = Namespace(**{
	412	+ "arch": "resnet50", "epochs": 2, "start_epoch": 0, "lr": 0.03, "schedule": [120, 160], "momentum": 0.9,
	413	+ "weight_decay": 1e-4, "gpu": None, "moco_dim": 128, "moco_k": 32000, "moco_m": 0.999, "moco_t": 0.07,
	414	+ "mlp": False, "aug_plus": False, "cos": False, "fake_data": True, "distributed": True,
	415	+ })
	416	+ try:
	417	+ dist.init_process_group(backend="nccl", init_method="tcp://localhost:10001", world_size=1, rank=0)
	418	+ except RuntimeError:
	419	+ pass # already initialized?
	420	+
	421	+ if device == "cpu":
	422	+ raise NotImplementedError("DistributedDataParallel/allgather requires npu")
	423	+
	424	+ self.model = MoCo(
	425	+ models.__dict__[self.opt.arch],
	426	+ self.opt.moco_dim,
	427	+ self.opt.moco_k,
	428	+ self.opt.moco_m,
	429	+ self.opt.moco_t,
	430	+ self.opt.mlp,
	431	+ )
	432	+ self.model.to(self.device)
	433	+
	434	+ # Define loss function (criterion) and optimizer
	435	+ self.criterion = nn.CrossEntropyLoss().to(self.device)
	436	+
	437	+ self.optimizer = torch.optim.SGD(
	438	+ self.model.parameters(),
	439	+ self.opt.lr,
	440	+ momentum=self.opt.momentum,
	441	+ weight_decay=self.opt.weight_decay,
	442	+ )
	443	+
	444	+ def collate_train_fn(data):
	445	+ ind = data[0]
	446	+ return [batches[2 * ind], batches[2 * ind + 1]], 0
	447	+
	448	+ batches = []
	449	+ for _ in range(4):
	450	+ batches.append(torch.randn(self.batch_size, 3, 224, 224).to(self.device))
	451	+ self.example_inputs = torch.utils.data.DataLoader(range(2), collate_fn=collate_train_fn)
	452	+ if torch.cuda.is_available():
	453	+ for _, (images, _) in enumerate(self.example_inputs):
	454	+ images[0] = images[0].cuda(device=0, non_blocking=True)
	455	+ images[1] = images[1].cuda(device=0, non_blocking=True)
	456	+ else:
	457	+ for _, (images, _) in enumerate(self.example_inputs):
	458	+ images[0] = images[0].npu(device=0, non_blocking=True)
	459	+ images[1] = images[1].npu(device=0, non_blocking=True)
	460	+
	461	+ Model.__init__ = new_init
	462	+
	463	+ if {"--only", "--amp", "--accuracy"} <= set(sys.argv):
	464	+ try:
	465	+ import torchvision
	466	+ from torchvision.models.resnet import ResNet
	467	+ except ImportError:
	468	+ log.warning("Import torchvision failed or could not get ResNet "
	469	+ "from module torchvision.models.resnet")
	470	+ return
	471	+
	472	+ def _new_forward_impl(self, x):
	473	+ # See note [TorchScript super()]
	474	+ x = self.conv1(x)
	475	+ x = self.bn1(x)
	476	+ x = self.relu(x)
	477	+ x = self.maxpool(x)
	478	+
	479	+ x = self.layer1(x)
	480	+ x = self.layer2(x)
	481	+ x = self.layer3(x)
	482	+ x = self.layer4(x)
	483	+
	484	+ @torch.compiler.disable(recursive=False)
	485	+ def avgpool(x):
	486	+ x = self.avgpool(x)
	487	+ return x
	488	+
	489	+ x = self.avgpool(x)
	490	+ x = torch.flatten(x, 1)
	491	+ x = self.fc(x)
	492	+ return x
	493	+
	494	+ ResNet._forward_impl = _new_forward_impl
	495	+
	496	+
	497	+@register_patch("timm_vovnet")
	498	+def _patch_model_14():
	499	+ if {"--only", "--amp", "--accuracy"} <= set(sys.argv):
	500	+ import torch.nn.functional as F
	501	+ from torch.nn.modules.conv import Conv2d
	502	+ from torch.nn.modules.pooling import AdaptiveAvgPool2d
	503	+ from torch.nn.modules.linear import Linear
	504	+
	505	+ def conv2d_amp_disabled(self, x):
	506	+ with torch.npu.amp.autocast(enabled=False):
	507	+ return self._conv_forward(x, self.weight, self.bias)
	508	+ Conv2d.forward = conv2d_amp_disabled
	509	+
	510	+ def adaptive_avgpool_amp_disabled(self, x):
	511	+ with torch.npu.amp.autocast(enabled=False):
	512	+ return F.adaptive_avg_pool2d(x.float(), self.output_size)
	513	+ AdaptiveAvgPool2d.forward = adaptive_avgpool_amp_disabled
	514	+
	515	+ def linear_amp_disabled(self, x):
	516	+ with torch.npu.amp.autocast(enabled=False):
	517	+ return F.linear(x, self.weight, self.bias)
	518	+ Linear.forward = linear_amp_disabled
	519	+
	520	+
	521	+@register_patch("resnet50", "resnet152", "resnext50_32x4d")
	522	+def _patch_model_18():
	523	+ if {"--only", "--amp", "--accuracy"} <= set(sys.argv):
	524	+ try:
	525	+ import torchvision.models as models
	526	+ except ImportError:
	527	+ log.warning("Import torchvision failed or could not get models "
	528	+ "from module torchvision.models")
	529	+ return
	530	+
	531	+ if 'resnet50' in sys.argv:
	532	+ from torchbenchmark.models.resnet50 import Model
	533	+ model = 'resnet50'
	534	+ weight = models.ResNet50_Weights.IMAGENET1K_V1
	535	+ elif 'resnet152' in sys.argv:
	536	+ from torchbenchmark.models.resnet152 import Model
	537	+ model = 'resnet152'
	538	+ weight = models.ResNet152_Weights.IMAGENET1K_V1
	539	+ elif 'resnext50_32x4d' in sys.argv:
	540	+ from torchbenchmark.models.resnext50_32x4d import Model
	541	+ model = 'resnext50_32x4d'
	542	+ weight = models.ResNeXt50_32X4D_Weights.IMAGENET1K_V1
	543	+ else:
	544	+ raise RuntimeError("args.only expect model resnet50, resnet152 or resnext50_32x4d")
	545	+
	546	+ def new_init(self, test, device, batch_size=None, extra_args=None):
	547	+ if extra_args is None:
	548	+ extra_args = []
	549	+ super(Model, self).__init__(model_name=model, test=test, device=device,
	550	+ batch_size=32, weights=weight,
	551	+ extra_args=extra_args)
	552	+ Model.__init__ = new_init
	553	+
	554	+
	555	+@register_patch("torch_multimodal_clip")
	556	+def _patch_model_19():
	557	+ try:
	558	+ from torchmultimodal.models.clip.text_encoder import CLIPTextEncoder
	559	+ except ImportError:
	560	+ log.warning("from torchmultimodal.models.clip.text_encoder import CLIPTextEncoder failed")
	561	+ return
	562	+
	563	+ def new_forward(self, text, return_hidden_state: bool = False):
	564	+ if text.size(1) != self.context_length:
	565	+ raise ValueError(
	566	+ f"length of input should be {self.context_length} but found {text.size(1)}"
	567	+ )
	568	+ embeddings = self.token_embedding(text)
	569	+ embeddings = embeddings + self.positional_embedding
	570	+ embeddings = embeddings.permute(1, 0, 2)
	571	+ embeddings = self.encoder(embeddings, mask=self.mask, is_causal=True)
	572	+
	573	+ # [n_ctx, bs, transformer.width] -> [bs, n_ctx, transformer.width]
	574	+ embeddings = torch.permute(embeddings, (1, 0, 2))
	575	+ hidden_state = self.ln_final(embeddings)
	576	+ hidden_state = hidden_state * 1 # pass
	577	+ if return_hidden_state:
	578	+ return hidden_state
	579	+
	580	+ projected_embeddings = self.projection(
	581	+ hidden_state[torch.arange(hidden_state.shape[0]), text.argmax(dim=-1)]
	582	+ )
	583	+ return projected_embeddings
	584	+
	585	+ CLIPTextEncoder.forward = new_forward
	586	+
	587	+
	588	+def patch_model(model_name):
	589	+ if model_name not in _patch_table.keys():
	590	+ return
	591	+ # do patch
	592	+ _patch_table[model_name]()

Abenchmarks/profiler.py+102-0

@@ -0,0 +1,102 @@
	1	+import torch
	2	+
	3	+
	4	+class CUDAProfiler:
	5	+ def __init__(self, enable=False, warmup=10, active=20, with_stack=False, with_memory=False, record_shapes=False, save_path="./profile"):
	zichun_ye2月18日这个给每个类加上注释，说清楚每个类是什么的在torch_npu里面加上CUDAProfiler比较奇怪，可以叫做DefaultProfiler，表示支持Torch原生的后端支持中的Profiling工作
	6	+ self.enable = enable
	7	+
	8	+ self.sp_with_stack = with_stack
	9	+ self.sp_with_memory = with_memory
	10	+ self.sp_record_shapes = record_shapes
	11	+ self.warmup = warmup
	12	+ self.active = active
	13	+ self.sp_save_path = save_path
	14	+
	15	+ activites = [torch.profiler.ProfilerActivity.CUDA,
	16	+ torch.profiler.ProfilerActivity.CPU]
	17	+
	18	+ self.prof = torch.profiler.profile(
	19	+ with_stack=self.sp_with_stack,
	20	+ record_shapes=self.sp_record_shapes,
	21	+ profile_memory=self.sp_with_memory,
	22	+ activities=activites,
	23	+ schedule=torch.profiler.schedule(wait=0, warmup=self.warmup, active=self.active, repeat=1),
	24	+ on_trace_ready=torch.profiler.tensorboard_trace_handler(self.sp_save_path))
	25	+
	26	+ def start(self):
	27	+ if self.enable:
	28	+ self.prof.start()
	29	+
	30	+ def step(self):
	31	+ if self.enable:
	32	+ self.prof.step()
	33	+
	34	+ def stop(self):
	35	+ if self.enable:
	36	+ self.prof.stop()
	37	+try:
	zichun_ye2月18日在PTA内部，就不要Try Except了，直接默认用户会安装torch_npu zichun_ye 2月18日评论： 2月18日评论：
	38	+ import torch_npu
	39	+
	40	+ class NPUProfiler:
	41	+ def __init__(self, enable=False, warmup=10, active=20, level="level0", with_stack=False, with_memory=False, record_shapes=False, save_path="./profile"):
	42	+ self.enable = enable
	43	+
	44	+ self.sp_level = level
	45	+ self.sp_with_stack = with_stack
	46	+ self.sp_with_memory = with_memory
	47	+ self.sp_record_shapes = record_shapes
	48	+ self.warmup = warmup
	49	+ self.active = active
	50	+ self.sp_save_path = save_path
	51	+
	52	+ if self.sp_level == 'level0':
	53	+ profiler_level = torch_npu.profiler.ProfilerLevel.Level0
	54	+ elif self.sp_level == 'level1':
	55	+ profiler_level = torch_npu.profiler.ProfilerLevel.Level1
	56	+ elif self.sp_level == 'level2':
	57	+ profiler_level = torch_npu.profiler.ProfilerLevel.Level2
	58	+ else:
	59	+ raise ValueError(f"profiler_level only supports level0,"
	60	+ f" 1, and 2, but gets {self.sp_level}")
	61	+
	62	+ experimental_config = torch_npu.profiler._ExperimentalConfig(
	63	+ export_type=[
	64	+ torch_npu.profiler.ExportType.Text,
	65	+ torch_npu.profiler.ExportType.Db
	66	+ ],
	67	+ profiler_level=profiler_level,
	68	+ msprof_tx=False,
	69	+ aic_metrics=torch_npu.profiler.AiCMetrics.AiCoreNone,
	70	+ l2_cache=False,
	71	+ op_attr=False,
	72	+ data_simplification=False,
	73	+ record_op_args=False,
	74	+ gc_detect_threshold=None
	75	+ )
	76	+
	77	+ activites = [torch_npu.profiler.ProfilerActivity.NPU,
	78	+ torch_npu.profiler.ProfilerActivity.CPU]
	79	+
	80	+ self.prof = torch_npu.profiler.profile(
	81	+ with_stack=self.sp_with_stack,
	82	+ record_shapes=self.sp_record_shapes,
	83	+ profile_memory=self.sp_with_memory,
	84	+ activities=activites,
	85	+ schedule=torch_npu.profiler.schedule(wait=0, warmup=self.warmup, active=self.active, repeat=1),
	86	+ on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(self.sp_save_path),
	87	+ experimental_config=experimental_config)
	88	+
	89	+ def start(self):
	90	+ if self.enable:
	91	+ self.prof.start()
	92	+
	93	+ def step(self):
	94	+ if self.enable:
	95	+ self.prof.step()
	96	+
	97	+ def stop(self):
	98	+ if self.enable:
	99	+ self.prof.stop()
	100	+except ImportError:
	101	+ # ignore the error if torch_npu is not installed
	102	+ pass

Abenchmarks/torchbench.py+420-0

Abenchmarks/torchbench_models_list.txt+31-0