[trainer, data] fix: do not drop_last SFT validation batches

tests/trainer/test_sft_val_batch_on_cpu.py

@@ -0,0 +1,35 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from verl.trainer.sft_val_utils import resolve_sft_val_batch_size
+
+
+def test_resolve_prefers_explicit_val_batch_size():
+    assert resolve_sft_val_batch_size({"val_batch_size": 16, "micro_batch_size_per_gpu": 4}, 256, 200) == 16
+
+
+def test_resolve_falls_back_to_micro_batch_not_train_batch():
+    assert resolve_sft_val_batch_size({"micro_batch_size_per_gpu": 4}, 256, 200) == 4
+
+
+def test_resolve_uses_val_len_when_no_micro_batch():
+    assert resolve_sft_val_batch_size({}, 256, 200) == 200
+
+
+def test_small_val_set_is_not_empty_with_resolved_batch():
+    train_bs = 256
+    n = 200
+    assert n // train_bs == 0  # old train-batch + drop_last path
+    batch = resolve_sft_val_batch_size({"micro_batch_size_per_gpu": 4}, train_bs, n)
+    assert (n + batch - 1) // batch > 0

verl/trainer/config/sft_trainer_engine.yaml

@@ -15,7 +15,11 @@ defaults:
 
 data:
   train_batch_size: 256 # global batch size
-  micro_batch_size_per_gpu: 4  # this is also val batch size
+  micro_batch_size_per_gpu: 4
+
+  # Validation batch size per DP rank. Null uses micro_batch_size_per_gpu.
+  val_batch_size: null
+
   max_token_len_per_gpu: 8192
   use_dynamic_bsz: True
   train_files: ~/data/gsm8k/train.parquet

verl/trainer/sft_trainer.py

@@ -31,6 +31,7 @@
 from torchdata.stateful_dataloader import StatefulDataLoader
 from tqdm import tqdm
 
+from verl.trainer.sft_val_utils import resolve_sft_val_batch_size
 from verl.utils import tensordict_utils as tu
 from verl.utils.checkpoint import CheckpointHandler
 from verl.utils.dataset.dataset_utils import SFTTensorCollator
@@ -253,17 +254,20 @@ def _build_dataloader(self):
         )
 
         if self.val_dataset:
+            val_batch_size = resolve_sft_val_batch_size(
+                config.data, self.train_batch_size_per_dp, len(self.val_dataset)
+            )
             self.val_sampler = DistributedSampler(
-                self.val_dataset, shuffle=False, num_replicas=dp_size, rank=dp_rank, drop_last=True
+                self.val_dataset, shuffle=False, num_replicas=dp_size, rank=dp_rank, drop_last=False
             )
             self.val_dataloader = StatefulDataLoader(
                 dataset=self.val_dataset,
-                batch_size=self.train_batch_size_per_dp,
+                batch_size=val_batch_size,
                 sampler=self.val_sampler,
                 collate_fn=self.collate_fn,
                 num_workers=self.config.data.num_workers,
                 pin_memory=False,
-                drop_last=True,
+                drop_last=False,
                 pin_memory_device=device_name,
             )
         else:
@@ -425,16 +429,26 @@ def fit(self):
                             val_losses.append(metrics["loss"])
 
                     if self.engine.is_mp_src_rank_with_outputs():
-                        val_loss = torch.mean(torch.tensor(val_losses, device=self.device_name))
-                        # average over data parallel group
+                        n_val = torch.tensor(float(len(val_losses)), device=self.device_name)
+                        sum_val = torch.tensor(
+                            float(sum(val_losses)) if val_losses else 0.0, device=self.device_name
+                        )
                         dp_group = self.engine.get_data_parallel_group()
                         if dp_group is not None:
-                            torch.distributed.all_reduce(val_loss, op=torch.distributed.ReduceOp.AVG, group=dp_group)
-
-                    if is_logging:
-                        metric = {"val/loss": val_loss.detach().item()}
-                        tracking.log(data=metric, step=global_step)
-                        last_valid_metric = metric
+                            torch.distributed.all_reduce(n_val, op=torch.distributed.ReduceOp.SUM, group=dp_group)
+                            torch.distributed.all_reduce(sum_val, op=torch.distributed.ReduceOp.SUM, group=dp_group)
+                        if n_val.item() <= 0:
+                            log_with_rank(
+                                "Validation produced no batches; skip val/loss rather than logging NaN.",
+                                logger=logger,
+                                rank=self.rank,
+                                level=logging.WARNING,
+                                log_only_rank_0=True,
+                            )
+                        elif is_logging:
+                            metric = {"val/loss": (sum_val / n_val).detach().item()}
+                            tracking.log(data=metric, step=global_step)
+                            last_valid_metric = metric
                     torch.distributed.barrier()
 
                 if is_last_step or (self.save_freq > 0 and is_save_step):

verl/trainer/sft_trainer_ray.py

@@ -32,6 +32,7 @@
 from torchdata.stateful_dataloader import StatefulDataLoader
 from tqdm import tqdm
 
+from verl.trainer.sft_val_utils import resolve_sft_val_batch_size
 from verl.utils import tensordict_utils as tu
 from verl.utils.checkpoint import CheckpointHandler, OrchestrationMode
 from verl.utils.dataset.dataset_utils import SFTTensorCollator
@@ -201,17 +202,20 @@ def _build_dataloader(self):
         )
 
         if self.val_dataset:
+            val_batch_size = resolve_sft_val_batch_size(
+                config.data, self.train_batch_size_per_dp, len(self.val_dataset)
+            )
             self.val_sampler = DistributedSampler(
-                self.val_dataset, shuffle=False, num_replicas=dp_size, rank=dp_rank, drop_last=True
+                self.val_dataset, shuffle=False, num_replicas=dp_size, rank=dp_rank, drop_last=False
             )
             self.val_dataloader = StatefulDataLoader(
                 dataset=self.val_dataset,
-                batch_size=self.train_batch_size_per_dp,
+                batch_size=val_batch_size,
                 sampler=self.val_sampler,
                 collate_fn=self.collate_fn,
                 num_workers=8,
                 pin_memory=False,
-                drop_last=True,
+                drop_last=False,
                 pin_memory_device=device_name,
             )
         else:
@@ -364,11 +368,19 @@ def fit(self):
                         metrics = tu.get(output, "metrics")
                         val_losses.append(metrics["loss"])
 
-                    val_loss = torch.mean(torch.tensor(val_losses, device=self.device_name))
-
-                    metric = {"val/loss": val_loss.detach().item()}
-                    tracking.log(data=metric, step=global_step)
-                    last_valid_metric = metric
+                    if not val_losses:
+                        log_with_rank(
+                            "Validation produced no batches; skip val/loss rather than logging NaN.",
+                            logger=logger,
+                            rank=0,
+                            level=logging.WARNING,
+                            log_only_rank_0=True,
+                        )
+                    else:
+                        val_loss = torch.mean(torch.tensor(val_losses, device=self.device_name))
+                        metric = {"val/loss": val_loss.detach().item()}
+                        tracking.log(data=metric, step=global_step)
+                        last_valid_metric = metric
 
                 if is_last_step or (self.save_freq > 0 and is_save_step):
                     self.ckpt_handler.save_checkpoint(step=global_step)

verl/trainer/sft_val_utils.py

@@ -0,0 +1,28 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+def resolve_sft_val_batch_size(data_config, train_batch_size_per_dp: int, val_dataset_len: int) -> int:
+    """Pick a validation batch size that does not swallow small val sets.
+
+    Preference: ``data.val_batch_size`` > ``data.micro_batch_size_per_gpu`` >
+    the full val set. Falls back to the train per-DP batch only if nothing else
+    is available.
+    """
+    val_batch_size = data_config.get("val_batch_size", None)
+    if val_batch_size is None:
+        val_batch_size = data_config.get("micro_batch_size_per_gpu", None)
+    if val_batch_size is None:
+        val_batch_size = val_dataset_len if val_dataset_len > 0 else train_batch_size_per_dp
+    return max(1, int(val_batch_size))