Customize loss logging frequency

GitOrigin-RevId: 98677f4

Author: changlan

axlearn/common/launch_trainer.py

@@ -61,6 +61,12 @@
     "If the trainer hangs for longer than this interval, "
     "the trainer will crash to prevent indefinite hanging.",
 )
+flags.DEFINE_integer(
+    "trainer_log_every_n_steps",
+    None,
+    "Logging frequency for the loss value during training. "
+    "If None, defaults to every 100 steps.",
+)
 flags.DEFINE_enum(
     "device_monitor",
     "none",
@@ -120,6 +126,8 @@ def get_trainer_config(
         trainer_config.crash_on_hang_timeout_seconds = (
             flag_values.trainer_crash_on_hang_timeout_seconds
         )
+    if trainer_config.log_every_n_steps is None:
+        trainer_config.log_every_n_steps = flag_values.trainer_log_every_n_steps
     for eval_cfg in trainer_config.evalers.values():
         eval_cfg.trace_at_iters = [int(el) for el in flag_values.eval_trace_at_iters]
     if flag_values.device_monitor == "tpu":

axlearn/common/trainer.py

@@ -217,6 +217,10 @@ class Config(Module.Config):
         # Defaults to None which is interpreted as True.
         cache_compiled_train_step: Optional[bool] = None
 
+        # Log the loss value every n steps. Defaults to None which is interpreted as every
+        # 100 steps.
+        log_every_n_steps: Optional[int] = None
+
     def __init__(
         self,
         cfg: Config,
@@ -1111,7 +1115,8 @@ def _run_step(
             # Run the compiled function.
             self._trainer_state, outputs = compiled_train_step_fn(self.trainer_state, input_batch)
 
-        if self.step % 100 == 0 or 0 <= self.step <= 5:
+        n = self._config.log_every_n_steps or 100
+        if self.step % n == 0 or 0 <= self.step <= 5:
             self._step_log(
                 "loss=%s aux=%s",
                 outputs["loss"],