diff --git a/areal/utils/data.py b/areal/utils/data.py index 1560742a1a..8d16752d17 100644 --- a/areal/utils/data.py +++ b/areal/utils/data.py @@ -1423,6 +1423,11 @@ def _build_group_slices( slices.append(slice(offset, offset + sz)) offset += sz return slices + if bs % self.group_size != 0: + raise ValueError( + f"batch size ({bs}) must be divisible by group_size " + f"({self.group_size}) when group_sizes is not provided" + ) return [ slice(i * self.group_size, (i + 1) * self.group_size) for i in range(bs // self.group_size) diff --git a/tests/test_reward_norm_variable_group.py b/tests/test_reward_norm_variable_group.py index c6ce65232f..8b89bff5dc 100644 --- a/tests/test_reward_norm_variable_group.py +++ b/tests/test_reward_norm_variable_group.py @@ -143,6 +143,15 @@ def test_group_sizes_non_positive_raises(): norm(x, group_sizes=[4, 0]) +def test_group_sizes_none_rejects_non_divisible_group_batch(): + """Implicit fixed-size grouping must not leave a partial tail unnormalized.""" + norm = Normalization(_group_norm_config()) # group_size=2 + x = torch.tensor([0.0, 1.0, 2.0], dtype=torch.float32) # bs=3 + + with pytest.raises(ValueError, match="divisible by group_size"): + norm(x) + + def test_adv_norm_style_2d_variable_groups_normalize_per_group(): """Token-level (2D) advantages normalize per variable-size group on dim 0.