add max_steps for validation to avoid hang

Author: DNXie

apps/sft/llama3_8b.yaml

@@ -32,12 +32,15 @@ training:
   steps: 1000
   compile: false
 
+validation:
+  local_batch_size: 1
+  freq: -1  # Change to a positive number to enable validation
+  steps: 200  # Max steps to run validation. Validation disabled if negative.
+
 dataset:
   path: yahma/alpaca-cleaned
   split: train[:95%]
 
-# Validation
-run_val_every_n_steps: null  # Change to an integer to enable validation every N steps
 dataset_val:
   path: yahma/alpaca-cleaned
   split: train[95%:]

apps/sft/main.py

@@ -62,7 +62,6 @@ def __init__(self, job_config: ForgeJobConfig):
         self.current_step = 0
         self.num_training_steps = job_config.training.steps
         self.gradient_accumulation_steps = 1  # Example value, adjust as needed
-        self._run_val_every_n_steps = job_config.get("run_val_every_n_steps", None)
         super().__init__(job_config)
         self.metric_logger = None  # TODO: fix this
 
@@ -74,8 +73,7 @@ def setup(self):
 
         self.val_dataloader = self.setup_data(
             self.job_config.dataset_val,
-            batch_size=self.job_config.training.local_batch_size,
-            infinite=False,
+            batch_size=self.job_config.validation.local_batch_size,
         )
 
         # self.train_dataloader = self.setup_data(
@@ -236,19 +234,22 @@ def train(self) -> None:
             )
 
             if (
-                self._run_val_every_n_steps is not None
-                and self.current_step % self._run_val_every_n_steps == 0
+                self.job_config.validation.freq > 0
+                and self.job_config.validation.steps > 0
+                and self.current_step % self.job_config.validation.freq == 0
             ):
-                self.validate()
+                self.validate(self.job_config.validation.steps)
 
-    def validate(self) -> None:
+    def validate(self, max_steps: int) -> None:
         for m in self.model_parts:
             m.eval()
         total_val_loss = torch.tensor(0.0, device=self.device)
         total_val_tokens = torch.tensor(0.0, device=self.device)
         with torch.no_grad():
             val_pbar = tqdm(self.val_dataloader, desc="Validation", leave=False)
             for batch_idx, batch in enumerate(val_pbar):
+                if batch_idx >= max_steps:
+                    break
                 batch_to_device(batch, self.device)
                 current_num_tokens = (batch["labels"] != CROSS_ENTROPY_IGNORE_IDX).sum()
                 # Compute loss