diff --git a/userbenchmark/dynamo/dynamobench/common.py b/userbenchmark/dynamo/dynamobench/common.py index c1d11c596..ed724c69d 100644 --- a/userbenchmark/dynamo/dynamobench/common.py +++ b/userbenchmark/dynamo/dynamobench/common.py @@ -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 = lambda *args: None + + CI_SKIP_DYNAMIC_BATCH_ONLY = { "sam", # See https://github.com/mindee/doctr/blob/f2114758d529ed8d3d0030581638f0520b6b98d8/doctr/models/detection/core.py#L89 @@ -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)) @@ -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: @@ -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: + 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: @@ -3918,6 +3928,16 @@ 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, @@ -4755,6 +4775,15 @@ def model_iter_fn_and_mark_step(*args, **kwargs): return if args.export_profiler_trace: + if args.profile_details: + args.profile_details = { + "record_shapes": True, + "profile_memory": True, + "with_stack": True, + "with_modules": True, + } + else: + args.profile_details = {} if args.profiler_trace_name is None: if args.backend: args.profiler_trace_name = args.backend