Enable some models with inline_inbuilt_nn_modules (#128315)

Summary:
For all models, graph breaks/recompiles reduce.
For drq, it increases and this is a legit one.

X-link: pytorch/pytorch#128315
Approved by: https://github.com/jansel

Reviewed By: izaitsevfb

Differential Revision: D58700178

Pulled By: anijain2305

fbshipit-source-id: 69af5ca4682486b9a418fc59c531e555e89da713

Co-authored-by: Laith Sakka <[email protected]>

userbenchmark/dynamo/dynamobench/common.py

@@ -2214,6 +2214,10 @@ def skip_models_due_to_control_flow(self):
  def guard_on_nn_module_models(self):
  return set()
 
+ @property
+ def inline_inbuilt_nn_modules_models(self):
+ return set()
+
  def get_tolerance_and_cosine_flag(self, is_training, current_device, name):
  raise NotImplementedError
 
@@ -4218,16 +4222,21 @@ def detect_and_mark_batch(t):
  if name in runner.guard_on_nn_module_models:
  guard_ctx = torch._dynamo.config.patch(guard_nn_modules=True)
 
+ inline_ctx = contextlib.nullcontext()
+ if name in runner.inline_inbuilt_nn_modules_models:
+ inline_ctx = torch._dynamo.config.patch(inline_inbuilt_nn_modules=True)
+
  with guard_ctx:
- runner.run_one_model(
- name,
- model,
- example_inputs,
- optimize_ctx,
- experiment,
- explain=args.explain,
- tag=args.tag,
- )
+ with inline_ctx:
+ runner.run_one_model(
+ name,
+ model,
+ example_inputs,
+ optimize_ctx,
+ experiment,
+ explain=args.explain,
+ tag=args.tag,
+ )
  if args.generate_aot_autograd_stats:
  stats_file = output_filename.split(".csv")[0] + "_stats.csv"
  output_csv(

userbenchmark/dynamo/dynamobench/torchbench.py

@@ -217,6 +217,19 @@ def guard_on_nn_module_models(self):
  "vision_maskrcnn",
  }
 
+ @property
+ def inline_inbuilt_nn_modules_models(self):
+ return {
+ "basic_gnn_edgecnn",
+ "drq",
+ "hf_Reformer",
+ "DALLE2_pytorch",
+ "hf_BigBird",
+ "detectron2_maskrcnn_r_50_fpn",
+ "detectron2_maskrcnn_r_101_fpn",
+ "vision_maskrcnn",
+ }
+
  def load_model(
  self,
  device,