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 12, 21:00 UTC
[MODEL] Implement plamo3 (#28834)
Signed-off-by: Shinichi Hemmi <50256998+Alnusjaponica@users.noreply.github.com>
Shinichi Hemmi committed
Nov 20, 2025 at 20:00 UTC
c9e093116c00781dda86df7a77e976c614b35d51
5 files changed
+438
docs/models/supported_models.md
+1
@@ -436,6 +436,7 @@ th {
436
| `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ |
437
| `PersimmonForCausalLM` | Persimmon | `adept/persimmon-8b-base`, `adept/persimmon-8b-chat`, etc. | | ✅︎ |
438
| `Plamo2ForCausalLM` | PLaMo2 | `pfnet/plamo-2-1b`, `pfnet/plamo-2-8b`, etc. | | ✅︎ |
439
+| `Plamo3ForCausalLM` | PLaMo3 | `pfnet/plamo-3-nict-2b-base`, `pfnet/plamo-3-nict-8b-base`, etc. | | ✅︎ |
440
| `QWenLMHeadModel` | Qwen | `Qwen/Qwen-7B`, `Qwen/Qwen-7B-Chat`, etc. | ✅︎ | ✅︎ |
441
| `Qwen2ForCausalLM` | QwQ, Qwen2 | `Qwen/QwQ-32B-Preview`, `Qwen/Qwen2-7B-Instruct`, `Qwen/Qwen2-7B`, etc. | ✅︎ | ✅︎ |
442
| `Qwen2MoeForCausalLM` | Qwen2MoE | `Qwen/Qwen1.5-MoE-A2.7B`, `Qwen/Qwen1.5-MoE-A2.7B-Chat`, etc. | ✅︎ | ✅︎ |
tests/distributed/test_pipeline_parallel.py
+1
@@ -130,6 +130,7 @@ TEXT_GENERATION_MODELS = {
130
"inceptionai/jais-13b-chat": PPTestSettings.fast(),
131
"ai21labs/Jamba-tiny-dev": PPTestSettings.fast(),
132
"pfnet/plamo-2-1b": PPTestSettings.fast(),
133
+ "pfnet/plamo-3-nict-2b-base": PPTestSettings.fast(),
134
"meta-llama/Llama-3.2-1B-Instruct": PPTestSettings.detailed(),
135
# Tests TransformersForCausalLM
136
"hmellor/Ilama-3.2-1B": PPTestSettings.fast(),
tests/models/registry.py
+4
@@ -402,6 +402,10 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
402
"pfnet/plamo-2-1b",
403
trust_remote_code=True,
404
),
405
+ "Plamo3ForCausalLM": _HfExamplesInfo(
406
+ "pfnet/plamo-3-nict-2b-base",
407
+ trust_remote_code=True,
408
+ ),
409
"QWenLMHeadModel": _HfExamplesInfo(
410
"Qwen/Qwen-7B-Chat",
411
max_transformers_version="4.53",
vllm/model_executor/models/plamo3.py
new
+431
@@ -0,0 +1,431 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Inference-only PLaMo3 model."""
4
+
5
+from collections.abc import Iterable
6
+from itertools import islice
7
+from typing import Any
8
+
9
+import torch
10
+from torch import nn
11
+from transformers import PretrainedConfig
12
+
13
+from vllm.attention.layer import Attention
14
+from vllm.compilation.decorators import support_torch_compile
15
+from vllm.config import VllmConfig
16
+from vllm.distributed import get_tensor_model_parallel_world_size
17
+from vllm.distributed.parallel_state import get_pp_group
18
+from vllm.model_executor.layers.activation import SiluAndMul
19
+from vllm.model_executor.layers.layernorm import RMSNorm
20
+from vllm.model_executor.layers.linear import (
21
+ MergedColumnParallelLinear,
22
+ QKVParallelLinear,
23
+ RowParallelLinear,
24
+)
25
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
26
+from vllm.model_executor.layers.quantization import QuantizationConfig
27
+from vllm.model_executor.layers.rotary_embedding import get_rope
28
+from vllm.model_executor.layers.vocab_parallel_embedding import (
29
+ DEFAULT_VOCAB_PADDING_SIZE,
30
+ ParallelLMHead,
31
+ VocabParallelEmbedding,
32
+)
33
+from vllm.model_executor.model_loader.weight_utils import (
34
+ LoaderFunction,
35
+ composed_weight_loader,
36
+ default_weight_loader,
37
+)
38
+from vllm.model_executor.models.interfaces import SupportsPP
39
+from vllm.model_executor.models.utils import (
40
+ AutoWeightsLoader,
41
+ extract_layer_index,
42
+ make_empty_intermediate_tensors_factory,
43
+ make_layers,
44
+ maybe_prefix,
45
+)
46
+from vllm.model_executor.utils import set_weight_attrs
47
+from vllm.sequence import IntermediateTensors
48
+
49
+
50
+# Only used for type hinting.
51
+class Plamo3Config(PretrainedConfig): # type: ignore
52
+ model_type: str = "plamo3"
53
+
54
+ hidden_size: int
55
+ num_hidden_layers: int
56
+ rms_norm_eps: float
57
+ # Attention
58
+ num_attention_heads: int
59
+ head_dim: int
60
+ num_key_value_heads: int
61
+ # vllm rename `sliding_window` attr to `interleaved_sliding_window`
62
+ # if `sliding_window` is list
63
+ interleaved_sliding_window: list[int | None]
64
+ sliding_window_pattern: int
65
+ rope_theta: int
66
+ rope_local_theta: int
67
+ # MLP
68
+ intermediate_size: int
69
+ # Tokenizer
70
+ vocab_size: int
71
+
72
+
73
+def rms_norm_weight_loader(offset: float) -> LoaderFunction:
74
+ return composed_weight_loader(
75
+ default_weight_loader,
76
+ lambda x: x + offset,
77
+ )
78
+
79
+
80
+class DenseMLP(nn.Module):
81
+ def __init__(
82
+ self,
83
+ config: Plamo3Config,
84
+ quant_config: QuantizationConfig | None = None,
85
+ prefix: str = "",
86
+ ) -> None:
87
+ super().__init__()
88
+ self.hidden_size = config.hidden_size
89
+ self.intermediate_size = config.intermediate_size
90
+ self.gate_up_proj = MergedColumnParallelLinear(
91
+ self.hidden_size,
92
+ [self.intermediate_size] * 2,
93
+ bias=False,
94
+ prefix=f"{prefix}.gate_up_proj",
95
+ quant_config=quant_config,
96
+ return_bias=False,
97
+ )
98
+ self.act = SiluAndMul()
99
+ self.down_proj = RowParallelLinear(
100
+ self.intermediate_size,
101
+ self.hidden_size,
102
+ bias=False,
103
+ prefix=f"{prefix}.down_proj",
104
+ quant_config=quant_config,
105
+ return_bias=False,
106
+ )
107
+
108
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
109
+ h = self.gate_up_proj(hidden_states)
110
+ h = self.act(h)
111
+ return self.down_proj(h)
112
+
113
+
114
+class Plamo3AttentionMixer(nn.Module):
115
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs) -> None:
116
+ super().__init__()
117
+ config = vllm_config.model_config.hf_config
118
+ quant_config = vllm_config.quant_config
119
+
120
+ self.hidden_size = config.hidden_size
121
+ tp_size = get_tensor_model_parallel_world_size()
122
+ self.total_num_heads = config.num_attention_heads
123
+ assert self.total_num_heads % tp_size == 0
124
+ self.num_heads = self.total_num_heads // tp_size
125
+ self.total_num_kv_heads = config.num_key_value_heads
126
+ if self.total_num_kv_heads >= tp_size:
127
+ # Number of KV heads is greater than TP size, so we partition
128
+ # the KV heads across multiple tensor parallel GPUs.
129
+ assert self.total_num_kv_heads % tp_size == 0
130
+ else:
131
+ # Number of KV heads is less than TP size, so we replicate
132
+ # the KV heads across multiple tensor parallel GPUs.
133
+ assert tp_size % self.total_num_kv_heads == 0
134
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
135
+ self.head_dim = config.head_dim
136
+ self.q_size = self.num_heads * self.head_dim
137
+ self.kv_size = self.num_kv_heads * self.head_dim
138
+ self.scaling = self.head_dim**-0.5
139
+
140
+ self.qkv_proj = QKVParallelLinear(
141
+ config.hidden_size,
142
+ self.head_dim,
143
+ self.total_num_heads,
144
+ self.total_num_kv_heads,
145
+ bias=False,
146
+ quant_config=quant_config,
147
+ prefix=f"{prefix}.qkv_proj",
148
+ )
149
+ self.o_proj = RowParallelLinear(
150
+ self.total_num_heads * self.head_dim,
151
+ config.hidden_size,
152
+ bias=False,
153
+ quant_config=quant_config,
154
+ prefix=f"{prefix}.o_proj",
155
+ )
156
+ layer_idx = extract_layer_index(prefix)
157
+ full_attn = config.interleaved_sliding_window[layer_idx] is None
158
+
159
+ self.rope_theta = config.rope_theta if full_attn else config.rope_local_theta
160
+ self.rope_scaling = (
161
+ config.rope_scaling if hasattr(config, "rope_scaling") else None
162
+ )
163
+ max_position = config.max_position_embeddings
164
+ if hasattr(vllm_config.model_config, "max_model_len") and isinstance(
165
+ vllm_config.model_config.max_model_len, int
166
+ ):
167
+ max_position = min(max_position, vllm_config.model_config.max_model_len)
168
+
169
+ self.rotary_emb = get_rope(
170
+ self.head_dim,
171
+ rotary_dim=self.head_dim,
172
+ max_position=max_position,
173
+ base=self.rope_theta,
174
+ rope_scaling=self.rope_scaling,
175
+ )
176
+ self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
177
+ set_weight_attrs(
178
+ self.q_norm.weight, {"weight_loader": rms_norm_weight_loader(offset=1.0)}
179
+ )
180
+ self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
181
+ set_weight_attrs(
182
+ self.k_norm.weight, {"weight_loader": rms_norm_weight_loader(offset=1.0)}
183
+ )
184
+ self.attn = Attention(
185
+ self.num_heads,
186
+ self.head_dim,
187
+ self.scaling,
188
+ num_kv_heads=self.num_kv_heads,
189
+ cache_config=vllm_config.cache_config,
190
+ per_layer_sliding_window=config.interleaved_sliding_window[layer_idx],
191
+ prefix=f"{prefix}.attn",
192
+ )
193
+
194
+ def forward(
195
+ self,
196
+ positions: torch.Tensor,
197
+ hidden_states: torch.Tensor,
198
+ residual: torch.Tensor | None,
199
+ **kwargs: Any,
200
+ ) -> torch.Tensor:
201
+ qkv, _ = self.qkv_proj(hidden_states)
202
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
203
+
204
+ q_shape = q.shape
205
+ q = q.reshape(q_shape[:-1] + (q_shape[-1] // self.head_dim, self.head_dim))
206
+ q = self.q_norm.forward_native(q).reshape(q_shape)
207
+ k_shape = k.shape
208
+ k = k.reshape(k_shape[:-1] + (k_shape[-1] // self.head_dim, self.head_dim))
209
+ k = self.k_norm.forward_native(k).reshape(k_shape)
210
+
211
+ q, k = self.rotary_emb(positions, q, k)
212
+ attn_output = self.attn(q, k, v)
213
+ output, _ = self.o_proj(attn_output)
214
+ return output
215
+
216
+
217
+class Plamo3DecoderLayer(nn.Module):
218
+ def __init__(
219
+ self, vllm_config: VllmConfig, prefix: str = "", **kwargs: Any
220
+ ) -> None:
221
+ super().__init__()
222
+ config = vllm_config.model_config.hf_config
223
+ quant_config = vllm_config.quant_config
224
+
225
+ self.mixer = Plamo3AttentionMixer(
226
+ vllm_config=vllm_config,
227
+ prefix=f"{prefix}.mixer",
228
+ )
229
+
230
+ self.mlp = DenseMLP(
231
+ config=config, quant_config=quant_config, prefix=f"{prefix}.mlp"
232
+ )
233
+ self.pre_mixer_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
234
+ set_weight_attrs(
235
+ self.pre_mixer_norm.weight,
236
+ {"weight_loader": rms_norm_weight_loader(offset=1.0)},
237
+ )
238
+ self.post_mixer_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
239
+ set_weight_attrs(
240
+ self.post_mixer_norm.weight,
241
+ {"weight_loader": rms_norm_weight_loader(offset=1.0 / 5)},
242
+ )
243
+ self.pre_mlp_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
244
+ set_weight_attrs(
245
+ self.pre_mlp_norm.weight,
246
+ {"weight_loader": rms_norm_weight_loader(offset=1.0)},
247
+ )
248
+ self.post_mlp_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
249
+ set_weight_attrs(
250
+ self.post_mlp_norm.weight,
251
+ {"weight_loader": rms_norm_weight_loader(offset=1.0 / (5**1.5))},
252
+ )
253
+
254
+ def forward(
255
+ self,
256
+ positions: torch.Tensor,
257
+ hidden_states: torch.Tensor,
258
+ residual: torch.Tensor | None,
259
+ **kwargs: Any,
260
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
261
+ if residual is None:
262
+ residual = hidden_states
263
+ hidden_states = self.pre_mixer_norm(hidden_states)
264
+ else:
265
+ hidden_states, residual = self.pre_mixer_norm(hidden_states, residual)
266
+
267
+ hidden_states = self.mixer(
268
+ positions=positions, hidden_states=hidden_states, residual=residual
269
+ )
270
+ hidden_states = self.post_mixer_norm(hidden_states)
271
+ # Fully Connected
272
+ hidden_states, residual = self.pre_mlp_norm(hidden_states, residual)
273
+ hidden_states = self.mlp(hidden_states)
274
+ hidden_states = self.post_mlp_norm(hidden_states)
275
+ return hidden_states, residual
276
+
277
+
278
+class Plamo3Decoder(torch.nn.Module):
279
+ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
280
+ super().__init__()
281
+ num_hidden_layers = vllm_config.model_config.hf_config.num_hidden_layers
282
+
283
+ self.start_layer, self.end_layer, self.layers = make_layers(
284
+ num_hidden_layers,
285
+ lambda prefix: Plamo3DecoderLayer(vllm_config, prefix=prefix),
286
+ prefix=f"{prefix}.layers",
287
+ )
288
+
289
+ def forward(
290
+ self,
291
+ positions: torch.Tensor,
292
+ hidden_states: torch.Tensor,
293
+ residual: torch.Tensor | None,
294
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
295
+ for layer in islice(self.layers, self.start_layer, self.end_layer):
296
+ hidden_states, residual = layer(
297
+ positions=positions,
298
+ hidden_states=hidden_states,
299
+ residual=residual,
300
+ )
301
+ return hidden_states, residual
302
+
303
+
304
+@support_torch_compile
305
+class Plamo3Model(nn.Module):
306
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
307
+ super().__init__()
308
+ config = vllm_config.model_config.hf_config
309
+
310
+ self.config = config
311
+ self.padding_idx = config.pad_token_id
312
+ self.vocab_size = config.vocab_size
313
+ self.org_vocab_size = config.vocab_size
314
+
315
+ self.embed_tokens = VocabParallelEmbedding(
316
+ self.vocab_size,
317
+ config.hidden_size,
318
+ org_num_embeddings=config.vocab_size,
319
+ prefix=f"{prefix}.embed_tokens",
320
+ )
321
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
322
+ ["hidden_states", "residual"], config.hidden_size
323
+ )
324
+ self.layers = Plamo3Decoder(vllm_config, prefix=f"{prefix}.layers")
325
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
326
+ set_weight_attrs(
327
+ self.norm.weight,
328
+ {"weight_loader": rms_norm_weight_loader(offset=1.0)},
329
+ )
330
+
331
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
332
+ return self.embed_tokens(input_ids)
333
+
334
+ def forward(
335
+ self,
336
+ input_ids: torch.Tensor,
337
+ positions: torch.Tensor,
338
+ intermediate_tensors: IntermediateTensors | None = None,
339
+ inputs_embeds: torch.Tensor | None = None,
340
+ ) -> torch.Tensor:
341
+ if get_pp_group().is_first_rank:
342
+ if inputs_embeds is not None:
343
+ hidden_states = inputs_embeds
344
+ else:
345
+ hidden_states = self.embed_input_ids(input_ids)
346
+ residual = None
347
+ else:
348
+ assert intermediate_tensors is not None
349
+ hidden_states = intermediate_tensors["hidden_states"]
350
+ residual = intermediate_tensors["residual"]
351
+
352
+ hidden_states, residual = self.layers(
353
+ positions=positions, hidden_states=hidden_states, residual=residual
354
+ )
355
+ if not get_pp_group().is_last_rank:
356
+ return IntermediateTensors(
357
+ {"hidden_states": hidden_states, "residual": residual}
358
+ )
359
+ hidden_states, _ = self.norm(hidden_states, residual)
360
+ return hidden_states
361
+
362
+
363
+class Plamo3ForCausalLM(nn.Module, SupportsPP):
364
+ packed_modules_mapping = {
365
+ "qkv_proj": [
366
+ "q_proj",
367
+ "k_proj",
368
+ "v_proj",
369
+ ],
370
+ }
371
+
372
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
373
+ super().__init__()
374
+ self.config = vllm_config.model_config.hf_config
375
+ self.vllm_config = vllm_config
376
+ self.model_config = vllm_config.model_config
377
+ self.scheduler_config = vllm_config.scheduler_config
378
+
379
+ self.model = Plamo3Model(
380
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
381
+ )
382
+
383
+ self.vocab_size = self.config.vocab_size
384
+ self.unpadded_vocab_size = self.config.vocab_size
385
+
386
+ num_embeddings = ((self.vocab_size + 15) // 16) * 16
387
+ self.lm_head = ParallelLMHead(
388
+ num_embeddings,
389
+ self.config.hidden_size,
390
+ org_num_embeddings=self.config.vocab_size,
391
+ padding_size=DEFAULT_VOCAB_PADDING_SIZE,
392
+ prefix=f"{prefix}.lm_head",
393
+ )
394
+ if self.config.tie_word_embeddings:
395
+ self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
396
+
397
+ self.logits_processor = LogitsProcessor(
398
+ self.unpadded_vocab_size, self.config.vocab_size
399
+ )
400
+ self.make_empty_intermediate_tensors = (
401
+ self.model.make_empty_intermediate_tensors
402
+ )
403
+
404
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
405
+ return self.model.embed_input_ids(input_ids)
406
+
407
+ def forward(
408
+ self,
409
+ input_ids: torch.Tensor,
410
+ positions: torch.Tensor,
411
+ intermediate_tensors: IntermediateTensors | None = None,
412
+ inputs_embeds: torch.Tensor | None = None,
413
+ ) -> torch.Tensor:
414
+ hidden_states = self.model(
415
+ input_ids, positions, intermediate_tensors, inputs_embeds
416
+ )
417
+ return hidden_states
418
+
419
+ def compute_logits(
420
+ self,
421
+ hidden_states: torch.Tensor,
422
+ ) -> torch.Tensor | None:
423
+ logits = self.logits_processor(self.lm_head, hidden_states)
424
+ return logits
425
+
426
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
427
+ loader = AutoWeightsLoader(
428
+ self,
429
+ skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
430
+ )
431
+ return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+1
@@ -157,6 +157,7 @@ _TEXT_GENERATION_MODELS = {
157
"Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
158
"PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"),
159
"Plamo2ForCausalLM": ("plamo2", "Plamo2ForCausalLM"),
160
+ "Plamo3ForCausalLM": ("plamo3", "Plamo3ForCausalLM"),
161
"QWenLMHeadModel": ("qwen", "QWenLMHeadModel"),
162
"Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
163
"Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"),