From 405a53cbfcd3f7b9412ebc89efb1f18013b154f3 Mon Sep 17 00:00:00 2001 From: Jeremy Schoemaker Date: Sun, 23 Aug 2026 19:55:38 -0500 Subject: [PATCH] Replace bare assert with clear errors when validation yields zero batches With the default concat-then-split packing, a per-rank validation split holding fewer than num_tokens_per_batch tokens produces zero packed rows, because partially filled rows are dropped. The validate loop then breaks on the first StopIteration with accumulated_loss still None, and the bare assert kills the run with a message-less AssertionError. If the assert were removed, the token-count division below would raise ZeroDivisionError instead. Raise ValueError with the cause and remedy for both conditions: zero validation batches on a rank, and zero valid (unmasked) tokens overall. No numerical path changes; the loss computation is untouched. --- tests/unit_tests/test_validate.py | 26 ++++++++++++++++++++++++++ torchtitan/components/validate.py | 16 +++++++++++++++- 2 files changed, 41 insertions(+), 1 deletion(-) 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")