Use benchmark_cls for checking precision.

[ysiraichi]

This PR makes it so we don't have to call load_benchmark only for checking the precision to be used.

cc @miladm @JackCaoG

[zpcore]

Refer to the issue here for the context: #6286
Thanks for making the fix.

The key point I think is to prevent leaving behind a dangling object which e.g., moved a model to xla device. del benchmark doesn't resolve the issue because it has already claimed the PJRT runtime. This will trigger the stackdump error: RuntimeError: Bad StatusOr access: UNKNOWN: TPU initialization failed: open(/dev/vfio/0): Device or resource busy: Device or resource busy; Couldn't open iommu group /dev/vfio/0.

[zpcore]

Since we only need to detect the precision, we can fetch the information directly without invoking

benchmark_cls(
        test=self.benchmark_experiment.test,
        device=device,
        batch_size=self.benchmark_experiment.batch_size,
    )

I think we can call the following load_benchmark_precision instead of load_benchmark to get the precision directly.

  def load_benchmark_precision(self):
   try:
     module = importlib.import_module(
         f"torchbenchmark.models.{self.model_name}")
   except ModuleNotFoundError:
     module = importlib.import_module(
         f"torchbenchmark.models.fb.{self.model_name}")
   benchmark_train_precision = getattr(module.Model, "DEFAULT_TRAIN_CUDA_PRECISION", None)
   benchmark_eval_precision = getattr(module.Model, "DEFAULT_EVAL_CUDA_PRECISION", None)
   return benchmark_train_precision, benchmark_eval_precision

WDYT?

[ysiraichi]

Right. Correct if I'm misunderstanding things, but isn't that exactly what I'm doing here?

[zpcore]

Right. Correct if I'm misunderstanding things, but isn't that exactly what I'm doing here?

Hah, you are right. I didn't notice that you called benchmark_cls instead.

Now it LGTM!