Skip to content
New issue

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

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

Already on GitHub? Sign in to your account

Add --profile-details and --export-perfdoctor option #2565

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 32 additions & 5 deletions userbenchmark/dynamo/dynamobench/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,12 @@ class CI(NamedTuple):
except ImportError:
INTERNAL_CI_SKIP_DYNAMIC_BATCH_ONLY = set()

try:
from pytorch.benchmark.fb.run_utils import trace_handler
except ImportError:
trace_handler = None


CI_SKIP_DYNAMIC_BATCH_ONLY = {
"sam",
# See https://github.com/mindee/doctr/blob/f2114758d529ed8d3d0030581638f0520b6b98d8/doctr/models/detection/core.py#L89
Expand Down Expand Up @@ -910,7 +916,7 @@ def maybe_mark_profile(*args, **kwargs):

times = args.iterations_per_run

with maybe_profile(args.export_profiler_trace) as p:
with maybe_profile(args.export_profiler_trace, **args.profile_details) as p:
for rep in trange(args.repeat, desc="running benchmark"):
inputs = (
randomize_input(copy.deepcopy(example_inputs))
Expand Down Expand Up @@ -1065,7 +1071,7 @@ def maybe_mark_profile(*args, **kwargs):
tolerance = args.xla_tolerance if args.trace_on_xla else 1e-4
torch._dynamo.config.repro_tolerance = tolerance

with maybe_profile(args.export_profiler_trace) as p:
with maybe_profile(args.export_profiler_trace, **args.profile_details) as p:
if args.export_aot_inductor:
frozen_model_iter_fn = export_aot_inductor(model, example_inputs)
else:
Expand Down Expand Up @@ -1114,9 +1120,13 @@ def maybe_mark_profile(*args, **kwargs):
name = args.profiler_trace_name + "_" + model.name
if hasattr(args, "rank"):
name += f"_rank_{args.rank}"
name += ".json"
name = os.path.join(torch._dynamo.config.base_dir, name)
p.export_chrome_trace(name)
if args.export_perfdoctor and trace_handler:
trace_handler(name, p)
else:
name += ".json"
name = os.path.join(torch._dynamo.config.base_dir, name)
p.export_chrome_trace(name)

median = np.median(timings, axis=0)
speedup = median[0] / median[1]
if args.dump_raw_metrics:
Expand Down Expand Up @@ -3918,6 +3928,14 @@ def get_example_inputs(self):
"--profiler_trace_name",
help="Overwrites exported trace name",
)
parser.add_argument(
"--profile-details", action="store_true", help="More detailed profiler trace."
)
parser.add_argument(
"--export-perfdoctor",
action="store_true",
help="Export Chrome trace to perf doctor. (internal only)",
)
parser.add_argument(
"--diff-branch",
default=diff_branch_default,
Expand Down Expand Up @@ -4754,7 +4772,16 @@ def model_iter_fn_and_mark_step(*args, **kwargs):
write_outputs(output_filename, [], [args.only, batch_size])
return

args.profile_details = {}
if args.export_profiler_trace:
if args.profile_details:
args.profile_details = {
"record_shapes": True,
"profile_memory": True,
"with_stack": True,
"with_modules": True,
}

if args.profiler_trace_name is None:
if args.backend:
args.profiler_trace_name = args.backend
Expand Down
Loading