diff --git a/tests/unit_tests/test_validate.py b/tests/unit_tests/test_validate.py index 55e5ecb751..b1c14544db 100644 --- a/tests/unit_tests/test_validate.py +++ b/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 diff --git a/torchtitan/components/validate.py b/torchtitan/components/validate.py index 585377369c..b8d889a296 100644 --- a/torchtitan/components/validate.py +++ b/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")