Change validation to use StatefulDataLoader

Author: bzz

recipes/full_finetune_distributed.py

@@ -361,9 +361,9 @@ def setup(self, cfg: DictConfig) -> None:
         )
 
         # Setup validation dataloader if validation dataset is provided
-        self._val_sampler, self._val_dataloader = None, None
+        self._val_dataloader = None
         if cfg.get("dataset_validation") is not None:
-            self._val_sampler, self._val_dataloader = self._setup_data(
+            self._val_dataloader = self._setup_data(
                 cfg_dataset=cfg.dataset_validation,
                 batch_size=cfg.batch_size,
                 collate_fn=collate_name,
@@ -788,8 +788,10 @@ def validate(self) -> float:
 
         with torch.no_grad():
             for batch_idx, batch in enumerate(self._val_dataloader):
-                if (self._max_validation_batches > 0 and
-                    batch_idx >= self._max_validation_batches):
+                if (
+                    self._max_validation_batches > 0
+                    and batch_idx >= self._max_validation_batches
+                ):
                     break
 
                 utils.batch_to_device(batch, self._device)
@@ -813,7 +815,7 @@ def validate(self) -> float:
         avg_val_loss = (
             (total_val_loss / total_val_tokens).item()
             if total_val_tokens > 0
-            else float('inf')
+            else float("inf")
         )
 
         self._model.train()
@@ -965,8 +967,10 @@ def train(self) -> None:
                     self._profiler.step()
 
                     # Run validation after gradient update
-                    if (self._run_val_every_n_steps is not None and
-                        self.global_step % self._run_val_every_n_steps == 0):
+                    if (
+                        self._run_val_every_n_steps is not None
+                        and self.global_step % self._run_val_every_n_steps == 0
+                    ):
                         pbar.refresh()
                         val_loss = self.validate()
                         if self._is_rank_zero:

recipes/lora_finetune_distributed.py

@@ -321,9 +321,9 @@ def setup(self, cfg: DictConfig) -> None:
             ),
         )
 
-        self._val_sampler, self._val_dataloader = None, None
+        self._val_dataloader = None
         if cfg.get("dataset_validation") is not None:
-            self._val_sampler, self._val_dataloader = self._setup_data(
+            self._val_dataloader = self._setup_data(
                 cfg_dataset=cfg.dataset_validation,
                 batch_size=cfg.batch_size,
                 collate_fn=collate_name,
@@ -897,8 +897,10 @@ def train(self) -> None:
                     self._profiler.step()
 
                     # Run validation after gradient update
-                    if (self._run_val_every_n_steps is not None and
-                        self.global_step % self._run_val_every_n_steps == 0):
+                    if (
+                        self._run_val_every_n_steps is not None
+                        and self.global_step % self._run_val_every_n_steps == 0
+                    ):
                         pbar.refresh()
                         val_loss = self.validate()
                         if self._is_rank_zero:
@@ -931,8 +933,10 @@ def validate(self) -> float:
 
         with torch.no_grad():
             for batch_idx, batch in enumerate(self._val_dataloader):
-                if (self._max_validation_batches > 0 and
-                    batch_idx >= self._max_validation_batches):
+                if (
+                    self._max_validation_batches > 0
+                    and batch_idx >= self._max_validation_batches
+                ):
                     break
 
                 utils.batch_to_device(batch, self._device)
@@ -942,7 +946,6 @@ def validate(self) -> float:
                     batch["labels"] != self._loss_fn.ignore_index
                 ).sum()
 
-
                 labels = batch.pop("labels")
 
                 with self.activations_handling_ctx:
@@ -972,7 +975,7 @@ def validate(self) -> float:
         avg_val_loss = (
             (total_val_loss / total_val_tokens).item()
             if total_val_tokens > 0
-            else float('inf')
+            else float("inf")
         )
 
         self._model.train()