Refactor val cadence

[anna-grim]

No description provided.

[carshadi]

Only a couple issues remain, please see the inline comments.

On the whole, I would make validation a single guarded operation, then call it periodically and once after training only when the final step was not already validated.

The key changes are:

1. Never pass None to save_model() or numeric formatting.
2. Track the last validated step.
3. Dedent final validation outside the epoch loop.
4. Skip duplicate final validation when the last step was already a cadence boundary.

Conceptually:

  step = 0
  best_score = np.inf
  last_validated_step = None

  def validate_and_checkpoint(step, label):
      nonlocal best_score, last_validated_step

      val_loss, val_cratio, val_score = self.validate(
          val_dataloader, step
      )
      last_validated_step = step

      if val_score is None:
          print(f"{label}: validation skipped (empty validation dataset)")
          return False

      is_best = val_score < best_score
      best_score = min(best_score, val_score)
      self.save_model(step, val_score)

      print(
          f"{label}: "
          f"val_loss={val_loss:.5f}, "
          f"val_cratio={val_cratio:.5f}, "
          f"val_score={val_score:.5f}"
          f"{' - New Best!' if is_best else ''}"
      )
      return is_best

Use it at periodic boundaries:

  if step % self.val_every == 0:
      avg_train_loss = running_loss / running_steps
      self.writer.add_scalar("train_loss", avg_train_loss, step)

      validate_and_checkpoint(step, f"Step {step}")

      running_loss = 0
      running_steps = 0
      self.model.train()

Then, crucially, place this after the entire outer epoch loop:

  # Outside: for epoch in range(self.max_epochs)
  if last_validated_step != step:
      validate_and_checkpoint(step, "Final model validation")

This produces exactly one final validation, avoids duplicate checkpoints at cadence boundaries, and cleanly skips checkpoint selection when validation is empty.

It would also be nice if we can add regression tests for:

• A run shorter than val_every still validates and saves once.
• A run ending exactly on val_every does not validate twice.
• Multiple epochs do not trigger unconditional epoch-end validation.
• An empty validation dataset neither saves nor raises an exception.

[carshadi]

This is close. In general, we should try to avoid introducing breaking public API changes that are unrelated to the scope of a PR, and if so they should be tested.

[carshadi]

Thanks @anna-grim !