Remove extraneous evaluation & export which happens on the final epoch.

Preserves current behavior for submodels of exporting duplicates (e.g., for 100 total steps, my_submodel_step_100.ckpt and my_submodel.ckpt are equivalent and both written after the final epoch).

PiperOrigin-RevId: 306629202

official/nlp/bert/model_training_utils.py

@@ -443,17 +443,19 @@ def run_customized_training_loop(
           train_summary_writer.flush()
       logging.info(training_status)
 
-      # Saves model checkpoints and run validation steps at every epoch end.
       if current_step % steps_per_epoch == 0:
-        # To avoid repeated model saving, we do not save after the last
-        # step of training.
-        if current_step < total_training_steps:
-          _save_checkpoint(strategy, checkpoint, model_dir,
-                           checkpoint_name.format(step=current_step))
+        # Save a submodel with the step in the file name after each epoch.
         if sub_model_export_name:
           _save_checkpoint(
               strategy, sub_model_checkpoint, model_dir,
               '%s_step_%d.ckpt' % (sub_model_export_name, current_step))
+
+        # Save model checkpoints and run validation steps after each epoch
+        # (with the exception of the final epoch which is handled after the
+        # training loop).
+        if current_step < total_training_steps:
+          _save_checkpoint(strategy, checkpoint, model_dir,
+                           checkpoint_name.format(step=current_step))
           if eval_input_fn:
             logging.info('Running evaluation after step: %s.', current_step)
             _run_evaluation(current_step,
@@ -462,12 +464,13 @@ def run_customized_training_loop(
             for metric in eval_metrics + model.metrics:
               metric.reset_states()
 
-    _save_checkpoint(strategy, checkpoint, model_dir,
-                     checkpoint_name.format(step=current_step))
     if sub_model_export_name:
       _save_checkpoint(strategy, sub_model_checkpoint, model_dir,
                        '%s.ckpt' % sub_model_export_name)
 
+    _save_checkpoint(strategy, checkpoint, model_dir,
+                     checkpoint_name.format(step=current_step))
+
     if eval_input_fn:
       logging.info('Running final evaluation after training is complete.')
       _run_evaluation(current_step,

official/nlp/bert/model_training_utils_test.py

@@ -156,6 +156,7 @@ class ModelTrainingUtilsTest(tf.test.TestCase, parameterized.TestCase):
         eval_input_fn=input_fn,
         eval_steps=10,
         init_checkpoint=None,
+        sub_model_export_name='my_submodel_name',
         metric_fn=metric_fn,
         custom_callbacks=None,
         run_eagerly=run_eagerly)
@@ -188,7 +189,20 @@ class ModelTrainingUtilsTest(tf.test.TestCase, parameterized.TestCase):
         distribution, model_dir, steps_per_loop=10, run_eagerly=False)
 
     # Two checkpoints should be saved after two epochs.
-    self.assertNotEmpty(tf.io.gfile.glob(os.path.join(model_dir, 'ctl_step_*')))
+    files = map(os.path.basename,
+                tf.io.gfile.glob(os.path.join(model_dir, 'ctl_step_*index')))
+    self.assertCountEqual(['ctl_step_20.ckpt-1.index',
+                           'ctl_step_40.ckpt-2.index'], files)
+
+    # Three submodel checkpoints should be saved after two epochs (one after
+    # each epoch plus one final).
+    files = map(os.path.basename,
+                tf.io.gfile.glob(os.path.join(model_dir,
+                                              'my_submodel_name*index')))
+    self.assertCountEqual(['my_submodel_name.ckpt-3.index',
+                           'my_submodel_name_step_20.ckpt-1.index',
+                           'my_submodel_name_step_40.ckpt-2.index'], files)
+
     self.assertNotEmpty(
         tf.io.gfile.glob(
             os.path.join(model_dir, 'summaries/training_summary*')))