Introduce validate_at_start

Author: cbalioglu

src/fairseq2/nn/utils/mask.py

@@ -28,7 +28,7 @@ def apply_mask(
 
     :returns: The input sequences with mask applied. *Shape:* Same as ``seqs``.
     """
-    unsqueeze(mask, dim=-1, count=seqs.ndim - mask.ndim)
+    mask = unsqueeze(mask, dim=-1, count=seqs.ndim - mask.ndim)
 
     return seqs.where(mask, fill_value)
 
@@ -95,7 +95,7 @@ def compute_row_mask(
         )
     else:
         # (N)
-        row_lens = row_lens.view(num_rows)
+        row_lens = row_lens.to(torch.int64).view(num_rows)
 
         # We only mask rows that are longer than the mask span length.
         if (span_len >= row_lens).any():

src/fairseq2/recipes/_trainer.py

@@ -120,7 +120,8 @@ class Trainer(Recipe, Generic[BatchT]):
     _seed: int
     _max_num_steps: int | None
     _max_num_data_epochs: int | None
-    _validator: Validator
+    _validator: Validator | None
+    _validate_at_start: bool
     _validate_after_n_steps: int
     _validate_every_n_steps: int | None
     _validate_after_n_data_epochs: int
@@ -181,6 +182,7 @@ def __init__(
         anomaly_detection: bool = False,
         max_num_steps: int | None = None,
         max_num_data_epochs: int | None = None,
+        validate_at_start: bool = False,
         validate_after_n_steps: int = 0,
         validate_every_n_steps: int | None = None,
         validate_after_n_data_epochs: int = 0,
@@ -328,6 +330,8 @@ def __init__(
 
         self._validator = validator
 
+        self._validate_at_start = validate_at_start
+
         if validate_every_n_steps is not None:
             if validate_every_n_steps <= 0:
                 raise ValueError(
@@ -507,6 +511,9 @@ def _maybe_restore_state(self) -> _TrainerState:
             ) from ex
 
         if step_nr is None:
+            if self._validate_at_start:
+                return _TrainerState.PRE_VALIDATION
+
             return _TrainerState.DATA_LOAD
 
         log.info("Restoring training from the last checkpoint at step {}.", step_nr)
@@ -574,6 +581,9 @@ def _do_run(self) -> None:
         with progress_task, self._lapse_watch:
             while self._state != _TrainerState.STOPPED:
                 match self._state:
+                    case _TrainerState.PRE_VALIDATION:
+                        self._state = self._pre_validate()
+
                     case _TrainerState.DATA_LOAD:
                         self._state = self._read_next_batches()
 
@@ -607,6 +617,12 @@ def _do_run(self) -> None:
 
                         self._state = self._stop()
 
+    def _pre_validate(self) -> _TrainerState:
+        if self._validate is not None:
+            self._validate()
+
+        return _TrainerState.DATA_LOAD
+
     def _read_next_batches(self) -> _TrainerState:
         with self._data_watch:
             try:
@@ -1045,6 +1061,9 @@ def _maybe_validate(self) -> float | None:
         return score
 
     def _should_validate(self) -> bool:
+        if self._validator is None:
+            return False
+
         return self._should_do(
             self._validate_after_n_steps,
             self._validate_every_n_steps,
@@ -1053,7 +1072,13 @@ def _should_validate(self) -> bool:
         )
 
     def _validate(self) -> float | None:
-        log.info("Starting validation after step {}.", self._step_nr)
+        if self._validator is None:
+            raise InternalError("`_validator` is `None`.")
+
+        if self._step_nr == 0:
+            log.info("Starting pre-validation before training.")
+        else:
+            log.info("Starting validation after step {}.", self._step_nr)
 
         self._model.module.eval()
 
@@ -1135,9 +1160,12 @@ def _should_do(
         after_n_data_epochs: int,
         every_n_data_epochs: int | None,
     ) -> bool:
+        if self._state == _TrainerState.PRE_VALIDATION:
+            return False
+
         def should_do_at_step() -> bool:
             if every_n_steps is not None:
-                if self._step_nr >= after_n_steps:
+                if self._step_nr > after_n_steps:
                     if self._step_nr % every_n_steps == 0:
                         return True
 
@@ -1158,7 +1186,7 @@ def should_do_at_step() -> bool:
 
         if self._state == _TrainerState.END_OF_DATA_EPOCH:
             if every_n_data_epochs is not None:
-                if self._data_epoch_nr >= after_n_data_epochs:
+                if self._data_epoch_nr > after_n_data_epochs:
                     if self._data_epoch_nr % every_n_data_epochs == 0:
                         already_done = should_do_at_step()
 
@@ -1191,16 +1219,17 @@ def step_nr(self) -> int:
 
 class _TrainerState(Enum):
     NOT_STARTED = 0
-    DATA_LOAD = 1
-    STEP = 2
-    POST_STEP = 3
-    END_OF_DATA_EPOCH = 4
-    END_OF_TRAINING = 5
-    END_OF_DATA = 6
-    GRADIENT_OVERFLOW = 7
-    EARLY_STOP = 8
-    STOP_REQUESTED = 9
-    STOPPED = 10
+    PRE_VALIDATION = 1
+    DATA_LOAD = 2
+    STEP = 3
+    POST_STEP = 4
+    END_OF_DATA_EPOCH = 5
+    END_OF_TRAINING = 6
+    END_OF_DATA = 7
+    GRADIENT_OVERFLOW = 8
+    EARLY_STOP = 9
+    STOP_REQUESTED = 10
+    STOPPED = 11
 
 
 T = TypeVar("T")

src/fairseq2/recipes/_validator.py

@@ -538,8 +538,8 @@ def reset(self) -> None:
 @final
 class NoopValidator(Validator):
     @override
-    def run(self, train_step_nr: int, train_data_epoch_nr: int) -> float:
-        return -torch.inf
+    def run(self, train_step_nr: int, train_data_epoch_nr: int) -> float | None:
+        return None
 
     @override
     def reset(self) -> None:

src/fairseq2/recipes/common/_trainer.py

@@ -119,6 +119,7 @@ def create_trainer(
         max_num_steps=regime_section.num_steps,
         max_num_data_epochs=regime_section.num_data_epochs,
         validator=validator,
+        validate_at_start=regime_section.validate_at_start,
         validate_after_n_steps=regime_section.validate_after_n_steps,
         validate_every_n_steps=regime_section.validate_every_n_steps,
         validate_after_n_data_epochs=regime_section.validate_after_n_data_epochs,

src/fairseq2/recipes/config.py

@@ -234,6 +234,9 @@ class RegimeSection:
     num_data_epochs: int | None = None
     """The maximum number of data epochs to train for."""
 
+    validate_at_start: bool = False
+    """If ``True``, runs validation before starting training."""
+
     validate_after_n_steps: int = 0
     """The number of steps after which to start validating the model."""
 

src/fairseq2/recipes/wav2vec2/asr/_train.py

@@ -105,7 +105,7 @@ class Wav2Vec2AsrTrainConfig:
     regime: RegimeSection = field(
         default_factory=lambda: RegimeSection(
             num_steps=20_000,
-            validate_after_n_steps=10_000,
+            validate_after_n_steps=9999,
             validate_every_n_steps=1_000,
             publish_metrics_every_n_steps=200,
         )