[benchmarks] Increase compilation cache. (pytorch#6509)

benchmarks/torchbench_model.py

@@ -76,6 +76,7 @@
         {
             "test": "eval",
             "xla": "PJRT",
+            "dynamo": None,
         },  # TIMEOUT
     ],
     "hf_T5_generate": [
@@ -85,6 +86,7 @@
         {
             "test": "eval",
             "xla": "PJRT",
+            "dynamo": None,
         },  # TIMEOUT
     ],
     "doctr_det_predictor": [{
@@ -135,6 +137,13 @@
     ],
 }
 
+# Models that had more graphs to be compiled than the actual size of
+# the cache.
+NEED_LARGER_CACHE = {
+    "cm3leon_generate",
+    "hf_T5_generate",
+}
+
 
 class TorchBenchModelLoader(ModelLoader):
 
@@ -401,6 +410,9 @@ def update_process_env(self, process_env):
     if precision_flag is not None:
       process_env[precision_flag] = '1'
 
+    if self.model_name in NEED_LARGER_CACHE:
+      process_env["XLA_COMPILATION_CACHE_SIZE"] = "2048"
+
   def pick_grad(self):
     # special case
     if self.model_name in ("maml",):