[brief] Allow users to stop training after a training step.

[detailed]
- This covers the cases where the loss functions go to infinity or
  become NaN, in which case training should stop immediately.

src/helios/trainer.py

@@ -812,6 +812,9 @@ def _train_on_iteration(self, state: TrainingState) -> None:
                 if state.global_iteration % accumulation_steps == 0 and not pbar.update():
                     pbar.refresh()
                 state.dataset_iter += 1
+                if self.model.should_training_stop():
+                    training_done = True
+                    break
 
                 if (
                     val_freq is not None
@@ -939,6 +942,10 @@ def _train_on_epoch(self, state: TrainingState) -> None:
                     if not ite_pbar.update():
                         ite_pbar.refresh()
 
+                    if self.model.should_training_stop():
+                        training_done = True
+                        break
+
             state.dataset_iter = 0
 
             if val_freq is not None and state.global_epoch % val_freq == 0: