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, 17:15 UTC
Add pipeline parallel support to `TransformersModel` (#12832)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Isotr0py <2037008807@qq.com> Co-authored-by: Isotr0py <2037008807@qq.com>
Harry Mellor committed
Mar 25, 2025 at 02:41 UTC
97cfa65df7b8ae8978880817b9bf000683c2b262
4 files changed
+245
-88
docs/source/models/supported_models.md
+1
-1
@@ -73,7 +73,7 @@ The Transformers fallback explicitly supports the following features:
73
74
- <project:#quantization-index> (except GGUF)
75
- <project:#lora-adapter>
76
-- <project:#distributed-serving> (pipeline parallel coming soon <gh-pr:12832>!)
76
+- <project:#distributed-serving> (requires `transformers>=4.49.0`)
77
78
#### Remote code
79
tests/distributed/test_pipeline_parallel.py
+3
@@ -175,6 +175,8 @@ TEXT_GENERATION_MODELS = {
175
"inceptionai/jais-13b-chat": PPTestSettings.fast(),
176
"ai21labs/Jamba-tiny-dev": PPTestSettings.fast(),
177
"meta-llama/Llama-3.2-1B-Instruct": PPTestSettings.detailed(),
178
+ # Tests TransformersModel
179
+ "ArthurZ/Ilama-3.2-1B": PPTestSettings.fast(),
180
"openbmb/MiniCPM-2B-sft-bf16": PPTestSettings.fast(),
181
"openbmb/MiniCPM3-4B": PPTestSettings.fast(),
182
# Uses Llama
@@ -243,6 +245,7 @@ TEST_MODELS = [
245
# [LANGUAGE GENERATION]
246
"microsoft/Phi-3.5-MoE-instruct",
247
"meta-llama/Llama-3.2-1B-Instruct",
248
+ # "ArthurZ/Ilama-3.2-1B", NOTE: Uncomment after #13905
249
"ibm/PowerLM-3b",
250
# [LANGUAGE EMBEDDING]
251
"intfloat/e5-mistral-7b-instruct",
vllm/model_executor/models/transformers.py
+230
-86
@@ -15,21 +15,25 @@
15
# limitations under the License.
16
"""Wrapper around `transformers` models"""
17
import re
18
+from itertools import chain
19
from typing import Iterable, Literal, Optional, Union
20
21
import torch
22
from torch import nn
22
-from transformers import AutoModel, PreTrainedModel
23
+from transformers import AutoModel, PretrainedConfig, PreTrainedModel
24
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
25
26
from vllm.attention import Attention
26
-from vllm.config import VllmConfig
27
-from vllm.distributed import get_tensor_model_parallel_world_size
27
+from vllm.config import (CacheConfig, DeviceConfig, ModelConfig,
28
+ ParallelConfig, VllmConfig)
29
+from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
30
+from vllm.distributed.utils import get_pp_indices
31
from vllm.logger import init_logger
32
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
33
ReplicatedLinear,
34
RowParallelLinear)
35
from vllm.model_executor.layers.logits_processor import LogitsProcessor
36
+from vllm.model_executor.layers.quantization import QuantizationConfig
37
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
38
from vllm.model_executor.layers.vocab_parallel_embedding import (
39
ParallelLMHead, VocabParallelEmbedding)
@@ -37,8 +41,9 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
41
from vllm.model_executor.sampling_metadata import SamplingMetadata
42
from vllm.sequence import IntermediateTensors
43
40
-from .interfaces import SupportsLoRA, SupportsQuant
41
-from .utils import maybe_prefix
44
+from .interfaces import SupportsLoRA, SupportsPP, SupportsQuant
45
+from .utils import (PPMissingLayer, is_pp_missing_parameter,
46
+ make_empty_intermediate_tensors_factory, maybe_prefix)
47
48
logger = init_logger(__name__)
49
@@ -53,7 +58,7 @@ def vllm_flash_attention_forward(
58
# Transformers kwargs
59
scaling: Optional[float] = None,
60
# vLLM kwargs
56
- attention_instances: Optional[list[Attention]] = None,
61
+ attention_instances: Optional[dict[Attention]] = None,
62
**kwargs):
63
self_attn = attention_instances[module.layer_idx]
64
if scaling is not None:
@@ -72,13 +77,12 @@ def log_replacement(name: str, old_module: nn.Module, new_module: nn.Module):
77
78
79
def replace_linear_class(
75
- linear: nn.Linear,
76
- style: Literal["colwise", "rowwise"],
77
- quant_config=None) -> Union[ColumnParallelLinear, RowParallelLinear]:
80
+ linear: nn.Linear, style: Literal["colwise", "rowwise"],
81
+ quant_config: QuantizationConfig
82
+) -> Union[ColumnParallelLinear, RowParallelLinear]:
83
"""
84
Replace nn.Linear with one of vLLM's tensor parallel linear classes.
85
81
- `quant_config` is not yet supported.
86
Args:
87
linear (nn.Linear): `nn.Linear` to be replaced.
88
style (str): Tensor parallel style of the new linear, e.g. "colwise".
@@ -105,7 +109,7 @@ def replace_linear_class(
109
)
110
111
108
-class TransformersModel(nn.Module, SupportsQuant, SupportsLoRA):
112
+class TransformersModel(nn.Module, SupportsQuant, SupportsLoRA, SupportsPP):
113
embedding_padding_modules = ["lm_head"]
114
embedding_modules = ["embed_tokens"
115
] # TODO transformers will have a util to get it
@@ -114,31 +118,175 @@ class TransformersModel(nn.Module, SupportsQuant, SupportsLoRA):
118
super().__init__()
119
logger.info("Using Transformers backend.")
120
117
- config = vllm_config.model_config.hf_config
118
- cache_config = vllm_config.cache_config
119
- model_config = vllm_config.model_config
120
- parallel_config = vllm_config.parallel_config
121
+ config: PretrainedConfig = vllm_config.model_config.hf_config
122
+ cache_config: CacheConfig = vllm_config.cache_config
123
+ device_config: DeviceConfig = vllm_config.device_config
124
+ model_config: ModelConfig = vllm_config.model_config
125
+ parallel_config: ParallelConfig = vllm_config.parallel_config
126
+ quant_config: QuantizationConfig = vllm_config.quant_config
127
128
self.config = config
129
+ self.cache_config = cache_config
130
+ self.device_config = device_config
131
+ self.model_config = model_config
132
+ self.parallel_config = parallel_config
133
+ self.quant_config = quant_config
134
+
135
self.vocab_size = model_config.get_vocab_size()
136
self.unpadded_vocab_size = model_config.get_vocab_size()
137
126
- self.model: PreTrainedModel = AutoModel.from_config(
127
- self.config,
128
- attn_implementation="vllm",
129
- torch_dtype=vllm_config.model_config.dtype,
130
- trust_remote_code=vllm_config.model_config.trust_remote_code,
131
- )
138
+ self.pp_group = get_pp_group()
139
+ self.pp_size = self.pp_group.world_size
140
+ self.pp_rank = self.pp_group.rank_in_group
141
+ self.tp_size = get_tensor_model_parallel_world_size()
142
+
143
+ # Use meta device to delay allocating GPU tensors
144
+ with torch.device("meta"):
145
+ self.model: PreTrainedModel = AutoModel.from_config(
146
+ config,
147
+ attn_implementation="vllm",
148
+ torch_dtype=model_config.dtype,
149
+ trust_remote_code=model_config.trust_remote_code,
150
+ )
151
prefix = self.model.base_model_prefix
152
134
- # MLP modifications
135
- self.apply_base_model_tp_plan(self.model)
153
+ self.pipeline_parallel()
154
+ self.tensor_parallel()
155
+
156
+ # Input embeddings
157
+ if not isinstance(self.model.get_input_embeddings(), PPMissingLayer):
158
+ self.model.set_input_embeddings(
159
+ VocabParallelEmbedding(
160
+ config.vocab_size,
161
+ config.hidden_size,
162
+ org_num_embeddings=config.vocab_size,
163
+ quant_config=quant_config,
164
+ ))
165
+
166
+ # Attention layers
167
+ self.attention_instances = self.create_attention_instances()
168
+
169
+ # Output embeddings
170
+ if not isinstance(getattr(self, "lm_head", None), PPMissingLayer):
171
+ self.unpadded_vocab_size = config.vocab_size
172
+ self.lm_head = ParallelLMHead(
173
+ config.vocab_size,
174
+ config.hidden_size,
175
+ quant_config=quant_config,
176
+ prefix=maybe_prefix(prefix, "lm_head"),
177
+ )
178
+ if config.tie_word_embeddings:
179
+ self.lm_head = self.lm_head.tie_weights(
180
+ self.model.get_input_embeddings())
181
+
182
+ logit_scale = getattr(config, "logit_scale", 1.0)
183
+ self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
184
+ config.vocab_size,
185
+ logit_scale)
186
+
187
+ # Initialize buffers (e.g. rotary embedding inverse frequency)
188
+ self.init_buffers(self.model)
189
+
190
+ # Move remaining meta tensors to device (should happen last)
191
+ self.meta_to_empty(self.model)
192
+
193
+ self.sampler = get_sampler()
194
+
195
+ self.make_empty_intermediate_tensors = (
196
+ make_empty_intermediate_tensors_factory(["hidden_states"],
197
+ config.hidden_size))
198
+
199
+ def pipeline_parallel(self):
200
+ """
201
+ Apply the model's pipeline parallelization plan.
202
+ """
203
+ if self.pp_size <= 1:
204
+ return
205
137
- # Attention modifications (assumes 1 attention op per hidden layer)
138
- num_heads = model_config.get_num_attention_heads(parallel_config)
139
- head_size = model_config.get_head_size()
140
- num_kv_heads = model_config.get_num_kv_heads(parallel_config)
141
- self.attention_instances = [
206
+ if not self.model.supports_pp_plan:
207
+ raise ValueError(
208
+ f"{type(self.model)} does not support pipeline parallel yet!")
209
+
210
+ module_lists = []
211
+ module_list_idx = None
212
+ pp_plan = list(self.model._pp_plan.keys())
213
+ for i, name in enumerate(pp_plan):
214
+ if isinstance(getattr(self.model, name), nn.ModuleList):
215
+ module_lists.append(name)
216
+ module_list_idx = i
217
+
218
+ if len(module_lists) > 1:
219
+ raise ValueError(
220
+ "Pipeline parallel of models with multiple `ModuleList`s "
221
+ "in the base model are not supported yet!")
222
+ if module_list_idx is None:
223
+ raise ValueError(
224
+ f"Could not find `ModuleList` in {type(self.model)}")
225
+
226
+ # Layers before module list
227
+ for name in pp_plan[:module_list_idx]:
228
+ if self.pp_group.is_first_rank or (self.config.tie_word_embeddings
229
+ and self.pp_group.is_last_rank):
230
+ continue
231
+ setattr(self.model, name, PPMissingLayer())
232
+
233
+ # Module list
234
+ start_layer, end_layer = get_pp_indices(self.config.num_hidden_layers,
235
+ self.pp_rank, self.pp_size)
236
+ layers_name = pp_plan[module_list_idx]
237
+ layers = getattr(self.model, layers_name)
238
+ for i in range(len(layers)):
239
+ if start_layer <= i and i < end_layer:
240
+ continue
241
+ layers[i] = PPMissingLayer(return_tuple=True)
242
+
243
+ # Layers after module list
244
+ for name in pp_plan[module_list_idx + 1:]:
245
+ # Modules that should be on last rank
246
+ if not self.pp_group.is_last_rank:
247
+ setattr(self.model, name, PPMissingLayer())
248
+
249
+ if not self.pp_group.is_last_rank:
250
+ self.lm_head = PPMissingLayer()
251
+
252
+ def tensor_parallel(self):
253
+ """
254
+ Apply the model's tensor parallelization plan.
255
+ Currently only supports linear layers.
256
+ """
257
+ if self.tp_size > 1 and self.config.base_model_tp_plan is None:
258
+ raise ValueError(
259
+ f"{type(self.model)} does not support tensor parallel yet!")
260
+
261
+ tp_plan = self.model._tp_plan
262
+
263
+ def _tensor_parallel(module: nn.Module, prefix: str = ""):
264
+ for child_name, child_module in module.named_children():
265
+ qual_name = maybe_prefix(prefix, child_name)
266
+ for pattern, style in tp_plan.items():
267
+ if re.match(pattern, qual_name) and isinstance(
268
+ child_module, nn.Linear):
269
+ new_module = replace_linear_class(
270
+ child_module, style, self.quant_config)
271
+ setattr(module, child_name, new_module)
272
+ log_replacement(qual_name, child_module, new_module)
273
+ else:
274
+ _tensor_parallel(child_module, prefix=qual_name)
275
+
276
+ _tensor_parallel(self.model)
277
+
278
+ def create_attention_instances(self) -> dict[int, Attention]:
279
+ """
280
+ Create `Attention` instances to inform KV cache allocation.
281
+ """
282
+ num_heads = self.model_config.get_num_attention_heads(
283
+ self.parallel_config)
284
+ head_size = self.model_config.get_head_size()
285
+ num_kv_heads = self.model_config.get_num_kv_heads(self.parallel_config)
286
+ start, end = get_pp_indices(self.config.num_hidden_layers,
287
+ self.pp_rank, self.pp_size)
288
+ return {
289
+ i:
290
Attention(
291
num_heads=num_heads,
292
head_size=head_size,
@@ -146,77 +294,70 @@ class TransformersModel(nn.Module, SupportsQuant, SupportsLoRA):
294
# Transformers, it's updated in vllm_flash_attention_forward
295
scale=head_size**-0.5,
296
num_kv_heads=num_kv_heads,
149
- cache_config=cache_config,
297
+ cache_config=self.cache_config,
298
quant_config=self.quant_config,
151
- prefix=f"{i}.attn") for i in range(config.num_hidden_layers)
152
- ]
153
-
154
- # Model modifications
155
- self.replace_vocab_embed_class(self.model)
156
-
157
- # ForCausalLM modifications
158
- self.lm_head = ParallelLMHead(self.vocab_size,
159
- config.hidden_size,
160
- quant_config=self.quant_config,
161
- prefix=maybe_prefix(prefix, "lm_head"))
162
- if config.tie_word_embeddings:
163
- self.lm_head.weight = self.model.get_input_embeddings().weight
164
-
165
- logit_scale = getattr(config, "logit_scale", 1.0)
166
- self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
167
- self.vocab_size, logit_scale)
168
- self.sampler = get_sampler()
299
+ prefix=f"{i}.attn")
300
+ for i in range(start, end)
301
+ }
302
170
- def apply_base_model_tp_plan(self, module: nn.Module, prefix: str = ""):
303
+ def init_buffers(self, module: nn.Module):
304
"""
172
- Apply the base model tensor parallelization plan to a module.
173
- Currently only supports linear layers.
305
+ If a `buffer` is on the `meta` device, then its parent
306
+ `module` is the original module created by:
307
+
308
+ ```python
309
+ with torch.device("meta"):
310
+ self.model: PreTrainedModel = AutoModel.from_config(...)
311
+ ```
312
+
313
+ This means that:
314
+ - `type(module)` is a class from `transformers`
315
+ - This class is constructed using a `PretrainedConfig`
316
"""
175
- if (self.config.base_model_tp_plan is None
176
- and get_tensor_model_parallel_world_size() > 1):
177
- raise ValueError(
178
- "Trying to run tensor parallelization but the model does not "
179
- "support it yet!")
180
-
181
- for child_name, child_module in module.named_children():
182
- qual_name = maybe_prefix(prefix, child_name)
183
- for pattern, style in self.config.base_model_tp_plan.items():
184
- if re.match(pattern, qual_name) and isinstance(
185
- child_module, nn.Linear):
186
- new_module = replace_linear_class(child_module, style,
187
- self.quant_config)
188
- setattr(module, child_name, new_module)
189
- log_replacement(qual_name, child_module, new_module)
190
- else:
191
- self.apply_base_model_tp_plan(child_module, prefix=qual_name)
192
-
193
- def replace_vocab_embed_class(self, module: nn.Module):
194
- # Use native set input embeddings
195
- new_module = VocabParallelEmbedding(
196
- self.vocab_size,
197
- self.config.hidden_size,
198
- org_num_embeddings=self.vocab_size,
199
- quant_config=None,
200
- )
201
- log_replacement("input embedding", self.model.get_input_embeddings(),
202
- new_module)
203
- module.set_input_embeddings(new_module)
317
+ for name, buffer in module.named_buffers(recurse=False):
318
+ if buffer.device == torch.device("meta"):
319
+ new_buffer = getattr(type(module)(self.config), name)
320
+ setattr(module, name, new_buffer)
321
+ for child in module.children():
322
+ self.init_buffers(child)
323
+
324
+ def meta_to_empty(self, module: nn.Module):
325
+ tensors = list(chain(module.buffers(), module.parameters()))
326
+ if tensors and all(t.device == torch.device("meta") for t in tensors):
327
+ module.to_empty(device=self.device_config.device)
328
+ return # We can stop recursing because to_empty is recursive
329
+ for child in module.children():
330
+ self.meta_to_empty(child)
331
332
def forward(
333
self,
207
- input_ids: torch.Tensor,
334
+ input_ids: Optional[torch.Tensor],
335
positions: torch.Tensor,
336
intermediate_tensors: Optional[IntermediateTensors] = None,
337
inputs_embeds: Optional[torch.Tensor] = None,
338
) -> Union[torch.Tensor, IntermediateTensors]:
212
- model_output = self.model(
213
- input_ids[None, ...],
339
+ if not get_pp_group().is_first_rank:
340
+ assert intermediate_tensors is not None
341
+ input_ids = None
342
+ inputs_embeds = intermediate_tensors["hidden_states"]
343
+
344
+ if input_ids is not None:
345
+ input_ids = input_ids[None, ...]
346
+ if inputs_embeds is not None:
347
+ inputs_embeds = inputs_embeds[None, ...]
348
+
349
+ hidden_states = self.model(
350
+ input_ids=input_ids,
351
+ inputs_embeds=inputs_embeds,
352
use_cache=False,
353
position_ids=positions[None, ...],
216
- intermediate_tensors=intermediate_tensors,
354
attention_instances=self.attention_instances,
355
return_dict=False)[0][0, ...] # we remove batch dimension for now
219
- return model_output
356
+
357
+ if not get_pp_group().is_last_rank:
358
+ return IntermediateTensors({"hidden_states": hidden_states})
359
+
360
+ return hidden_states
361
362
def compute_logits(
363
self,
@@ -238,8 +379,11 @@ class TransformersModel(nn.Module, SupportsQuant, SupportsLoRA):
379
params_dict = dict(self.named_parameters())
380
loaded_params = set[str]()
381
for name, loaded_weight in weights:
241
- if name not in params_dict:
242
- name = f"{self.model.base_model_prefix}.{name}"
382
+ # Necessary for some models which use remote code
383
+ if not name.startswith(prefix := self.model.base_model_prefix):
384
+ name = maybe_prefix(prefix, name)
385
+ if is_pp_missing_parameter(name, self):
386
+ continue
387
if name in params_dict:
388
param = params_dict[name]
389
weight_loader = getattr(param, "weight_loader",
vllm/model_executor/models/utils.py
+11
-1
@@ -472,6 +472,16 @@ class PPMissingLayer(torch.nn.Identity):
472
473
def __init__(self, *args, **kwargs):
474
super().__init__()
475
+ self.return_tuple = kwargs.get("return_tuple", False)
476
+
477
+ def forward(self, *args, **kwargs):
478
+ """
479
+ Return the first arg from args or the first value from kwargs.
480
+
481
+ Wraps the input in a tuple if `self.return_tuple` is True.
482
+ """
483
+ input = args[0] if args else next(iter(kwargs.values()))
484
+ return (input, ) if self.return_tuple else input
485
486
487
_CPU_OFFLOAD_BYTES = 0
@@ -650,4 +660,4 @@ def cast_overflow_tensors(
660
if tensors.isinf().any() or tensors.isnan().any():
661
clamp_value = torch.finfo(tensors.dtype).max - offset
662
tensors = torch.clamp(tensors, min=-clamp_value, max=clamp_value)
653
- return tensors
\ No newline at end of file
663
+ return tensors