| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import os |
| 5 | |
| 6 | import pytest |
| 7 | |
| 8 | from vllm.distributed.utils import get_pp_indices |
| 9 | |
| 10 | |
| 11 | def test_custom_layer_partition(monkeypatch: pytest.MonkeyPatch): |
| 12 | with monkeypatch.context() as m: |
| 13 | |
| 14 | def _verify(partition_str, num_layers, pp_size, goldens): |
| 15 | bak = os.environ.get("VLLM_PP_LAYER_PARTITION", None) |
| 16 | m.setenv("VLLM_PP_LAYER_PARTITION", partition_str) |
| 17 | for pp_rank, golden in enumerate(goldens): |
| 18 | assert get_pp_indices(num_layers, pp_rank, pp_size) == golden |
| 19 | if bak is not None: |
| 20 | m.setenv("VLLM_PP_LAYER_PARTITION", bak) |
| 21 | |
| 22 | # Even partition |
| 23 | _verify("5,5,5,5", 20, 4, [(0, 5), (5, 10), (10, 15), (15, 20)]) |
| 24 | # Balanced partition |
| 25 | _verify("4,6,6,4", 20, 4, [(0, 4), (4, 10), (10, 16), (16, 20)]) |
| 26 | # Put reminder somewhere |
| 27 | _verify("5,6,5,6", 22, 4, [(0, 5), (5, 11), (11, 16), (16, 22)]) |
| 28 | # Invalid partition strings |
| 29 | with pytest.raises(ValueError): |
| 30 | _verify("5,5,5,5,", 20, 4, [(0, 5), (5, 10), (10, 15), (15, 20)]) |
| 31 | with pytest.raises(ValueError): |
| 32 | _verify("5,5,5,a", 20, 4, [(0, 5), (5, 10), (10, 15), (15, 20)]) |
| 33 | # Wrong number of partitions |
| 34 | with pytest.raises(ValueError): |
| 35 | _verify("5,5,5", 20, 4, [(0, 5), (5, 10), (10, 15), (15, 20)]) |
| 36 | # Wrong number of layers |
| 37 | with pytest.raises(ValueError): |
| 38 | _verify("5,5,5,5", 21, 4, [(0, 5), (5, 10), (10, 15), (15, 20)]) |
| 39 | |
| 40 | |
| 41 | @pytest.mark.parametrize( |
| 42 | "num_hidden_layers,pp_size,pp_rank,indices", |
| 43 | [ |
| 44 | # pp_size 2 |
| 45 | (2, 2, 0, (0, 1)), |
| 46 | (2, 2, 1, (1, 2)), |
| 47 | (3, 2, 0, (0, 2)), |
| 48 | (3, 2, 1, (2, 3)), |
| 49 | # pp_size 3 |
| 50 | (3, 3, 0, (0, 1)), |
| 51 | (3, 3, 1, (1, 2)), |
| 52 | (3, 3, 2, (2, 3)), |
| 53 | (4, 3, 0, (0, 1)), |
| 54 | (4, 3, 1, (1, 3)), |
| 55 | (4, 3, 2, (3, 4)), |
| 56 | (5, 3, 0, (0, 2)), |
| 57 | (5, 3, 1, (2, 4)), |
| 58 | (5, 3, 2, (4, 5)), |
| 59 | ], |
| 60 | ) |
| 61 | def test_uneven_auto_partition( |
| 62 | num_hidden_layers: int, |
| 63 | pp_size: int, |
| 64 | pp_rank: int, |
| 65 | indices: tuple[int, int], |
| 66 | ): |
| 67 | assert indices == get_pp_indices(num_hidden_layers, pp_rank, pp_size) |