Raise clear errors when validation yields zero batches or zero valid tokens

tests/unit_tests/test_validate.py

@@ -210,3 +210,29 @@ def generate_image(**kwargs):
 
     assert generated["img_height"] == 6
     assert generated["img_width"] == 10
+
+
+def test_generic_validator_raises_on_zero_validation_batches(monkeypatch):
+    loader = _ClosableLoader([])
+    validator = _generic_validator(loader)
+    monkeypatch.setattr(validate_module.utils, "device_type", "cpu")
+
+    with pytest.raises(ValueError, match="zero batches"):
+        validator.validate([nn.Identity()], step=1)
+
+    assert loader.closed
+
+
+def test_generic_validator_raises_on_zero_valid_tokens(monkeypatch):
+    row = (
+        {"input": torch.ones(1, 1)},
+        torch.full((1, 1), validate_module.IGNORE_INDEX, dtype=torch.long),
+    )
+    loader = _ClosableLoader([row])
+    validator = _generic_validator(loader)
+    monkeypatch.setattr(validate_module.utils, "device_type", "cpu")
+
+    with pytest.raises(ValueError, match="zero valid tokens"):
+        validator.validate([nn.Identity()], step=1)
+
+    assert loader.closed

torchtitan/components/validate.py

@@ -329,8 +329,22 @@ def validate(
             total_global_valid_tokens.add_(global_valid_tokens)
             num_steps += 1
 
-        assert accumulated_loss is not None
+        if accumulated_loss is None:
+            raise ValueError(
+                "Validation ran zero batches on this rank. This happens when the "
+                "validation dataset supplies fewer than num_tokens_per_batch "
+                "tokens on this rank, because concat-then-split packing drops "
+                "partially filled batches. Decrease "
+                "training.num_tokens_per_microbatch_per_dp_rank or use a larger "
+                "validation dataset."
+            )
         num_global_valid_tokens = int(total_global_valid_tokens.item())
+        if num_global_valid_tokens == 0:
+            raise ValueError(
+                "Validation ran on zero valid tokens; cannot compute an average "
+                "validation loss. Ensure the validation batches contain unmasked "
+                "labels."
+            )
         if parallel_dims.dp_cp_enabled:
             global_loss_sum = dist_utils.dist_sum(
                 accumulated_loss, parallel_dims.get_optional_mesh("loss")