A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 22:15 UTC
Improve pipeline partitioning (#13839)
Harry Mellor committed
Feb 26, 2025 at 02:53 UTC
145944cb94a6fc663c05451763315d45e771a285
2 files changed
+46
-8
tests/distributed/test_pipeline_partition.py
+24
@@ -34,3 +34,27 @@ def test_custom_layer_partition():
34
# Wrong number of layers
35
with pytest.raises(ValueError):
36
_verify("5,5,5,5", 21, 4, [(0, 5), (5, 10), (10, 15), (15, 20)])
37
+
38
+
39
+@pytest.mark.parametrize(
40
+ "num_hidden_layers,pp_size,pp_rank,indices",
41
+ [
42
+ # pp_size 2
43
+ (2, 2, 0, (0, 1)),
44
+ (2, 2, 1, (1, 2)),
45
+ (3, 2, 0, (0, 2)),
46
+ (3, 2, 1, (2, 3)),
47
+ # pp_size 3
48
+ (3, 3, 0, (0, 1)),
49
+ (3, 3, 1, (1, 2)),
50
+ (3, 3, 2, (2, 3)),
51
+ (4, 3, 0, (0, 1)),
52
+ (4, 3, 1, (1, 3)),
53
+ (4, 3, 2, (3, 4)),
54
+ (5, 3, 0, (0, 2)),
55
+ (5, 3, 1, (2, 4)),
56
+ (5, 3, 2, (4, 5)),
57
+ ])
58
+def test_uneven_auto_partition(num_hidden_layers: int, pp_size: int,
59
+ pp_rank: int, indices: tuple[int, int]):
60
+ assert indices == get_pp_indices(num_hidden_layers, pp_rank, pp_size)
vllm/distributed/utils.py
+22
-8
@@ -67,8 +67,17 @@ def split_tensor_along_last_dim(
67
def get_pp_indices(num_hidden_layers: int, pp_rank: int,
68
pp_size: int) -> Tuple[int, int]:
69
"""Try to evenly distribute layers across partitions.
70
+
71
If the number of layers is not divisible by the number of partitions,
71
- the last partition will have the remaining layers.
72
+ the remaining layers are evenly distributed across all but the last
73
+ partition. The last partition is excluded because it often contains an
74
+ additional norm layer and we are attempting to balance compute.
75
+
76
+ If `pp_size > 2` and the number of remaining layers is
77
+ `0 < x <= pp_size - 2` then the remaining layers are evenly distributed
78
+ across the middle partitions. The first and last partitions are excluded
79
+ because they contain the input and output embeddings respectively and we
80
+ are attempting to reduce maximum memory consumption across partitions.
81
"""
82
partition_list_str = envs.VLLM_PP_LAYER_PARTITION
83
if partition_list_str is not None:
@@ -84,15 +93,20 @@ def get_pp_indices(num_hidden_layers: int, pp_rank: int,
93
if sum(partitions) != num_hidden_layers:
94
raise ValueError(
95
f"{sum(partitions)=} does not match {num_hidden_layers=}.")
87
- start_layer = sum(partitions[:pp_rank])
88
- end_layer = start_layer + partitions[pp_rank]
96
else:
97
layers_per_partition = num_hidden_layers // pp_size
91
- start_layer = pp_rank * layers_per_partition
92
- end_layer = start_layer + layers_per_partition
93
-
94
- if pp_rank == pp_size - 1:
95
- end_layer = num_hidden_layers
98
+ partitions = [layers_per_partition for _ in range(pp_size)]
99
+
100
+ if remaining_layers := num_hidden_layers % pp_size:
101
+ for i in range(2, remaining_layers + 2):
102
+ partitions[-i] += 1
103
+ logger.info("Hidden layers were unevenly partitioned: %s",
104
+ ",".join(str(p) for p in partitions))
105
+ logger.info("This can be manually overridden using the "
106
+ "VLLM_PP_LAYER_PARTITION environment variable")
107
+
108
+ start_layer = sum(partitions[:pp_rank])
109
+ end_layer = start_layer + partitions[pp_rank]
110
111
return (start_layer, end_layer)
112