Internal change

PiperOrigin-RevId: 307950230

official/benchmark/transformer_benchmark.py

@@ -112,7 +112,7 @@ class TransformerBenchmark(PerfZeroBenchmark):
                         'max_value': bleu_max})
 
     if (warmup and 'step_timestamp_log' in stats and
-        len(stats['step_timestamp_log']) > warmup):
+        len(stats['step_timestamp_log']) > warmup + 1):
       # first entry in the time_log is start of step 1. The rest of the
       # entries are the end of each step recorded
       time_log = stats['step_timestamp_log']
@@ -127,6 +127,11 @@ class TransformerBenchmark(PerfZeroBenchmark):
       metrics.append({'name': 'avg_exp_per_second',
                       'value': stats['avg_exp_per_second']})
 
+    if 'step_timestamp_log' in stats:
+      time_log = stats['step_timestamp_log']
+      metrics.append({'name': 'startup_time',
+                      'value': time_log[0].timestamp - start_time_sec})
+
     flags_str = flags_core.get_nondefault_flags_as_str()
     self.report_benchmark(iters=-1, wall_time=wall_time_sec, metrics=metrics,
                           extras={'flags': flags_str})
@@ -686,6 +691,7 @@ class TransformerBigKerasBenchmarkReal(TransformerKerasBenchmark):
     self._setup()
     FLAGS.model_dir = self._get_model_dir('benchmark_2x2_tpu')
     FLAGS.train_steps = 300
+    FLAGS.log_steps = 300
     FLAGS.distribution_strategy = 'tpu'
     FLAGS.static_batch = True
     FLAGS.use_ctl = True
@@ -700,6 +706,26 @@ class TransformerBigKerasBenchmarkReal(TransformerKerasBenchmark):
         total_batch_size=FLAGS.batch_size,
         log_steps=FLAGS.log_steps)
 
+  def benchmark_4x4_tpu(self):
+    """Port of former GCP transformer_big model on 4x4."""
+    self._setup()
+    FLAGS.model_dir = self._get_model_dir('benchmark_4x4_tpu')
+    FLAGS.train_steps = 300
+    FLAGS.log_steps = 300
+    FLAGS.distribution_strategy = 'tpu'
+    FLAGS.static_batch = True
+    FLAGS.use_ctl = True
+    FLAGS.batch_size = 24576
+    FLAGS.max_length = 64
+    FLAGS.decode_batch_size = 32
+    FLAGS.decode_max_length = 97
+    FLAGS.padded_decode = True
+    FLAGS.enable_checkpointing = False
+
+    self._run_and_report_benchmark(
+        total_batch_size=FLAGS.batch_size,
+        log_steps=FLAGS.log_steps)
+
 
 if __name__ == '__main__':
   tf.test.main()