Set model name early to keep warmup and main model same (#159231)

Summary:
X-link: pytorch/pytorch#159231
Approved by: https://github.com/williamwen42
ghstack dependencies: #159209

Reviewed By: ZainRizvi

Differential Revision: D79112571

fbshipit-source-id: fbaf66b7e28f639684fe4e59f6c54b8a69736b6a

Author: anijain2305

userbenchmark/dynamo/dynamobench/common.py

@@ -2425,6 +2425,8 @@ def warmup(fn, model, example_inputs, mode, niters=10):
         # Use distributed wrapping as necessary
         model = self.deepcopy_and_maybe_parallelize(model)
 
+        if not hasattr(model, name):
+            model.name = name
         self.init_optimizer(name, current_device, model.parameters())
 
         # The self.autocast context is needed for the model we export with aot_compile,
@@ -2528,8 +2530,6 @@ def warmup(fn, model, example_inputs, mode, niters=10):
             result_summary = latency_experiment_summary(
                 self.suite_name, self.args, model, timings, **experiment_kwargs
             )
-            if not hasattr(model, name):
-                model.name = name
             results.append(result_summary)
             return " ".join(map(str, results))
 
@@ -2586,6 +2586,9 @@ def warmup(fn, model, example_inputs, mode, niters=5):
         # Use distributed wrapping as necessary
         model = self.deepcopy_and_maybe_parallelize(model)
 
+        if not hasattr(model, name):
+            model.name = name
+
         self.init_optimizer(name, current_device, model.parameters())
 
         # The self.autocast context is needed for the model we export with aot_compile,
@@ -2699,8 +2702,6 @@ def warmup(fn, model, example_inputs, mode, niters=5):
                     f"{ok:3}/{total:3} +{frames_third_pass} frames {compilation_time:3.0f}s"
                 )
 
-            if not hasattr(model, name):
-                model.name = name
             results.append(experiment(model, example_inputs, **experiment_kwargs))
             return " ".join(map(str, results))