Merge pull request #4547 from AI-Hypercomputer:zxhe/eval_start_step

PiperOrigin-RevId: 952443710

Author: Google-ML-Automation

src/maxtext/configs/base.yml

@@ -984,6 +984,7 @@ decode_sampling_nucleus_p: -1 # set if you're doing nucleus / top-p
 decode_sampling_top_k: 0 # set if you're doing top-k
 decode_sampling_temperature: 1.
 
+eval_start_step: 0 # start eval after train step is >= eval_start_step
 eval_interval: -1  # the specific number of train step between eval_step
 eval_steps: -1  # run this number of steps for eval, recommend setting this to prevent error due to running out of evel data
 target_eval_loss: 0.  # early stop once reaching target eval_loss

src/maxtext/configs/types.py

@@ -1585,6 +1585,10 @@ class TrainingLoop(BaseModel):
       description="Total number of training steps. -1 defaults to learning_rate_schedule_steps.",
   )
   log_period: int = Field(100, description="Frequency (in steps) to log metrics and flush Tensorboard.")
+  eval_start_step: int = Field(
+      0,
+      description="Start evaluation after training step is >= eval_start_step.",
+  )
   eval_interval: int = Field(
       -1,
       description="Run evaluation every N training steps. -1 disables interval-based evaluation.",

src/maxtext/trainers/pre_train/train.py

@@ -674,6 +674,7 @@ def training_loop_iteration(
   eval_interval = immutable_data["eval_interval"]
   eval_steps = immutable_data["eval_steps"]
   start_step = immutable_data["start_step"]
+  eval_start_step = immutable_data["eval_start_step"]
 
   # HLO dump config
   dump_hlo = immutable_data["dump_hlo"]
@@ -717,7 +718,7 @@ def training_loop_iteration(
         all_host_upload=dump_hlo_upload_all,
     )
 
-  if eval_interval > 0 and step > start_step and (step + 1) % eval_interval == 0:
+  if eval_interval > 0 and step >= start_step and step >= eval_start_step and (step + 1) % eval_interval == 0:
     assert eval_data_iterator
     # Explicitly reset the eval iterator and counters before starting the eval loop
     eval_data_iterator.reset()
@@ -871,6 +872,7 @@ def train_loop(config, recorder, state=None):
       "steps": config.steps,
       "eval_interval": config.eval_interval,
       "eval_steps": config.eval_steps,
+      "eval_start_step": config.eval_start_step,
       "save_checkpoint_on_completion": config.save_checkpoint_on_completion,
       "start_step": start_step,
       "dump_hlo": config.dump_hlo,

tests/unit/pyconfig_test.py

@@ -317,6 +317,21 @@ def test_local_sa_flags_explicit_override(self):
     self.assertEqual(config.local_sa_v_layout, "SEQ_MINOR")
     self.assertFalse(config.local_use_splash_scheduler)
 
+  def test_eval_start_step_config(self):
+    """Verifies that eval_start_step defaults to 0 and can be overridden via pyconfig."""
+    config_default = pyconfig.initialize(
+        [os.path.join(MAXTEXT_PKG_DIR, "train.py"), get_test_config_path()],
+        skip_jax_distributed_system=True,
+    )
+    self.assertEqual(config_default.eval_start_step, 0)
+
+    config_override = pyconfig.initialize(
+        [os.path.join(MAXTEXT_PKG_DIR, "train.py"), get_test_config_path()],
+        skip_jax_distributed_system=True,
+        eval_start_step=50,
+    )
+    self.assertEqual(config_override.eval_start_step, 50)
+
 
 if __name__ == "__main__":
   unittest.main()