@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3badb0213b

Fork 0

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, 14:45 UTC

[Model] Add PLaMo2 (#14323)

Signed-off-by: Shinichi Hemmi <50256998+Alnusjaponica@users.noreply.github.com> Signed-off-by: shemmi <shemmi@preferred.jp> Co-authored-by: Kento Nozawa <nzw0301@preferred.jp> Co-authored-by: Hiroaki Mikami <mhiroaki@preferred.jp> Co-authored-by: Calvin Metzger <metzger@preferred.jp>

Shinichi Hemmi committed Apr 16, 2025 at 11:31 UTC 3badb0213b11dcd74db85dea04b883f3391c47a7

9 files changed +800 -24

.buildkite/test-pipeline.yaml

+6 -1

@@ -400,8 +400,9 @@ steps:

400 - pytest -v -s models/test_transformers.py

401 - pytest -v -s models/test_registry.py

402 # V1 Test: https://github.com/vllm-project/vllm/issues/14531

403 - - VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'not llama4'

403 + - VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'not llama4 and not plamo2'

404 - VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'llama4'

405 + - VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'plamo2'

406

407 - label: Language Models Test (Standard) # 32min

408 #mirror_hardwares: [amd]

@@ -411,6 +412,8 @@ steps:

412 - tests/models/embedding/language

413 - tests/models/encoder_decoder/language

414 commands:

415 + # Install causal-conv1d for plamo2 models here, as it is not compatible with pip-compile.

416 + - pip install causal-conv1d

417 - pytest -v -s models/decoder_only/language -m 'core_model or quant_model'

418 - pytest -v -s models/embedding/language -m core_model

419

@@ -422,6 +425,8 @@ steps:

425 - tests/models/embedding/language

426 - tests/models/encoder_decoder/language

427 commands:

428 + # Install causal-conv1d for plamo2 models here, as it is not compatible with pip-compile.

429 + - pip install causal-conv1d

430 - pytest -v -s models/decoder_only/language -m 'not core_model and not quant_model'

431 - pytest -v -s models/embedding/language -m 'not core_model'

432

docs/source/models/supported_models.md

@@ -497,6 +497,11 @@ See [this page](#generative-models) for more information on how to use generativ

497 * `adept/persimmon-8b-base`, `adept/persimmon-8b-chat`, etc.

498 *

499 * ✅︎

500 +- * `Plamo2ForCausalLM`

501 + * PLaMo2

502 + * `pfnet/plamo-2-1b`, `pfnet/plamo-2-8b`, etc.

503 + *

504 + *

505 - * `QWenLMHeadModel`

506 * Qwen

507 * `Qwen/Qwen-7B`, `Qwen/Qwen-7B-Chat`, etc.

requirements/test.in

@@ -27,6 +27,7 @@ torch==2.6.0

27 torchaudio==2.6.0

28 torchvision==0.21.0

29 transformers_stream_generator # required for qwen-vl test

30 +mamba_ssm # required for plamo2 test

31 matplotlib # required for qwen-vl test

32 mistral_common[opencv] >= 1.5.4 # required for pixtral test

33 num2words # required for smolvlm test

requirements/test.txt

@@ -111,6 +111,7 @@ einops==0.8.0

111 # via

112 # -r requirements/test.in

113 # encodec

114 + # mamba-ssm

115 # vector-quantize-pytorch

116 # vocos

117 einx==0.3.0

@@ -233,6 +234,8 @@ lxml==5.3.0

234 # via

235 # blobfile

236 # sacrebleu

237 +mamba-ssm==2.2.4

238 + # via -r requirements/test.in

239 markdown-it-py==3.0.0

240 # via rich

241 markupsafe==3.0.2

@@ -268,6 +271,8 @@ mypy-extensions==1.0.0

271 # via black

272 networkx==3.2.1

273 # via torch

274 +ninja==1.11.1.3

275 + # via mamba-ssm

276 nltk==3.9.1

277 # via rouge-score

278 num2words==0.5.14

@@ -360,6 +365,7 @@ packaging==24.1

365 # fastparquet

366 # huggingface-hub

367 # lazy-loader

368 + # mamba-ssm

369 # matplotlib

370 # peft

371 # plotly

@@ -571,6 +577,7 @@ sentencepiece==0.2.0

577 # via mistral-common

578 setuptools==75.8.0

579 # via

580 + # mamba-ssm

581 # pytablewriter

582 # torch

583 shellingham==1.5.4

@@ -627,6 +634,7 @@ torch==2.6.0

634 # encodec

635 # fastsafetensors

636 # lm-eval

637 + # mamba-ssm

638 # peft

639 # runai-model-streamer

640 # sentence-transformers

@@ -664,6 +672,7 @@ transformers==4.51.1

672 # -r requirements/test.in

673 # genai-perf

674 # lm-eval

675 + # mamba-ssm

676 # peft

677 # sentence-transformers

678 # transformers-stream-generator

tests/models/decoder_only/language/test_hybrid.py

+18 -23

@@ -9,9 +9,15 @@ from vllm.sampling_params import SamplingParams

9 from ...utils import check_outputs_equal

11 # This test is for the hybrid models

12 -MODELS = ["ai21labs/Jamba-tiny-dev", "Zyphra/Zamba2-1.2B-instruct"]

12 +MODELS = [

13 + "ai21labs/Jamba-tiny-dev", "Zyphra/Zamba2-1.2B-instruct",

14 + "pfnet/plamo-2-1b"

15 +]

16 # Bamba at Fp32 is too big for the CI (L4 GPU).

17 # MODELS = ["ai21labs/Jamba-tiny-dev", "ibm-ai-platform/Bamba-9B"]

18 +# Note: Running Plamo2 in transformers implementation requires to install

19 +# causal-conv1d package, which is not listed as a test dependency as it's

20 +# not compatible with pip-compile.

23 @pytest.mark.parametrize("model", MODELS)

@@ -25,21 +31,11 @@ def test_models(

31 dtype: str,

32 max_tokens: int,

33 ) -> None:

28 -

34 # numeric error produces different generation

35 if "Bamba" in model:

36 example_prompts.pop(3)

33 - model_kwargs = {

34 - "use_mamba_kernels": False, # mamba kernels are not installed so HF

35 - # don't use them

36 - }

37 - if "Zamba2" in model:

38 - # Zamba2 HF implementation automatically checks if mamba kernels are

39 - # installed

40 - model_kwargs = {}

41 -

42 - with hf_runner(model, dtype=dtype, model_kwargs=model_kwargs) as hf_model:

38 + with hf_runner(model, dtype=dtype) as hf_model:

39 hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)

41 with vllm_runner(model, dtype=dtype) as vllm_model:

@@ -94,6 +90,10 @@ def test_mamba_prefill_chunking_with_parallel_sampling(

90 # correctly for n > 1 decoding steps inside a

91 # chunked prefill forward pass (where we have both prefills

92 # and decoding together )

93 +

94 + if 'plamo-2' in model:

95 + dtype = "float" # use a different dtype for plamo

96 +

97 sampling_params = SamplingParams(n=3,

98 temperature=1,

99 seed=0,

@@ -125,20 +125,14 @@ def test_mamba_prefill_chunking(hf_runner, vllm_runner, example_prompts,

125 example_prompts.pop(3)

126 example_prompts.pop(2)

127 dtype = "half" # use a different dtype for Bamba

128 +

129 elif "Zamba2" in model:

130 example_prompts.pop(7)

131 dtype = "half"

132 + elif "plamo-2-1b" in model:

133 + example_prompts.pop(7)

134

132 - model_kwargs = {

133 - "use_mamba_kernels": False, # mamba kernels are not installed so HF

134 - # don't use them

135 - }

136 - if "Zamba2" in model:

137 - # Zamba2 HF implementation automatically checks if mamba kernels are

138 - # installed

139 - model_kwargs = {}

140 -

141 - with hf_runner(model, dtype=dtype, model_kwargs=model_kwargs) as hf_model:

135 + with hf_runner(model, dtype=dtype) as hf_model:

136 non_chunked = hf_model.generate_greedy(example_prompts, max_tokens)

137

138 with vllm_runner(model,

@@ -208,7 +202,8 @@ def test_mamba_cache_cg_padding(

202 # This test is for verifying that mamba cache is padded to CG captured

203 # batch size. If it's not, a torch RuntimeError will be raised because

204 # tensor dimensions aren't compatible

211 - vllm_config = EngineArgs(model=model).create_engine_config()

205 + vllm_config = EngineArgs(model=model,

206 + trust_remote_code=True).create_engine_config()

207 while len(example_prompts) == vllm_config.pad_for_cudagraph(

208 len(example_prompts)):

209 example_prompts.append(example_prompts[0])

tests/models/registry.py

@@ -204,6 +204,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

204 trust_remote_code=True),

205 "PhiMoEForCausalLM": _HfExamplesInfo("microsoft/Phi-3.5-MoE-instruct",

206 trust_remote_code=True),

207 + "Plamo2ForCausalLM": _HfExamplesInfo("pfnet/plamo-2-1b",

208 + trust_remote_code=True),

209 "QWenLMHeadModel": _HfExamplesInfo("Qwen/Qwen-7B-Chat",

210 trust_remote_code=True),

211 "Qwen2ForCausalLM": _HfExamplesInfo("Qwen/Qwen2-7B-Instruct",

vllm/config.py

+12

@@ -2838,6 +2838,13 @@ def _get_and_verify_dtype(

2838 else:

2839 torch_dtype = config_dtype

2840

2841 + if config.model_type == "plamo2":

2842 + logger.info(

2843 + "For PLaMo2, we cast models to bfloat16 instead of using "

2844 + "float16 by default. This is because float16 does not work."

2845 + )

2846 + torch_dtype = torch.bfloat16

2847 +

2848 from vllm.platforms import current_platform

2849 if (current_platform.is_cpu()

2850 and current_platform.get_cpu_architecture()

@@ -2867,6 +2874,11 @@ def _get_and_verify_dtype(

2874 "using float16 by default. Please specify `dtype` if you "

2875 "want to use float16.")

2876 torch_dtype = torch.bfloat16

2877 + elif dtype == "float16" and config.model_type == "plamo2":

2878 + logger.warning(

2879 + "For PLaMo2, using float16 is unstable and might cause "

2880 + "unexpected behavior. Please use bfloat16 or float32 instead.")

2881 + torch_dtype = torch.float16

2882 else:

2883 if dtype not in _STR_DTYPE_TO_TORCH_DTYPE:

2884 raise ValueError(f"Unknown dtype: {dtype}")

vllm/model_executor/models/plamo2.py new

+746

@@ -0,0 +1,746 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +"""Inference-only PLaMo2 model."""

3 +import math

4 +from typing import Iterable, Optional, Tuple

5 +

6 +import torch

7 +from torch import nn

8 +from transformers import PretrainedConfig, PreTrainedModel

9 +

10 +from vllm.attention.backends.abstract import AttentionMetadata

11 +from vllm.attention.layer import Attention

12 +from vllm.config import CacheConfig, VllmConfig

13 +from vllm.distributed import get_tensor_model_parallel_world_size

14 +from vllm.forward_context import get_forward_context

15 +from vllm.model_executor.layers.layernorm import RMSNorm

16 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,

17 + MergedColumnParallelLinear,

18 + QKVParallelLinear,

19 + RowParallelLinear)

20 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

21 +from vllm.model_executor.layers.mamba.ops.causal_conv1d import (

22 + causal_conv1d_fn, causal_conv1d_update)

23 +from vllm.model_executor.layers.mamba.ops.mamba_ssm import (

24 + selective_scan_fn, selective_state_update)

25 +from vllm.model_executor.layers.quantization import QuantizationConfig

26 +from vllm.model_executor.layers.rotary_embedding import get_rope

27 +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

28 +from vllm.model_executor.layers.vocab_parallel_embedding import (

29 + DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)

30 +from vllm.model_executor.model_loader.weight_utils import (

31 + composed_weight_loader, default_weight_loader, sharded_weight_loader)

32 +from vllm.model_executor.models.interfaces import (HasInnerState, IsHybrid,

33 + SupportsV0Only)

34 +from vllm.model_executor.models.mamba_cache import (MambaCacheManager,

35 + MambaCacheParams)

36 +from vllm.model_executor.models.utils import maybe_prefix

37 +from vllm.model_executor.sampling_metadata import SamplingMetadata

38 +from vllm.model_executor.utils import set_weight_attrs

39 +from vllm.sequence import IntermediateTensors

40 +from vllm.utils import LayerBlockType

41 +

42 +

43 +# Only used for type hinting.

44 +class Plamo2Config(PretrainedConfig): # type: ignore

45 + model_type: str = "plamo2"

46 +

47 + hidden_size: int

48 + num_hidden_layers: int

49 + rms_norm_eps: float

50 + # Attention

51 + num_attention_heads: int

52 + hidden_size_per_head: int

53 + num_key_value_heads: int

54 + # Mamba

55 + mamba_d_state: int

56 + mamba_d_conv: int

57 + mamba_num_heads: int

58 + mamba_step: int

59 + # MLP

60 + intermediate_size: int

61 + # Tokenizer

62 + vocab_size: int

63 +

64 +

65 +class Plamo2PreTrainedModel(PreTrainedModel): # type: ignore

66 +

67 + def _init_weights(self, module: torch.nn.Module) -> None:

68 + std = 0.02

69 + if isinstance(module, nn.Linear):

70 + module.weight.data.normal_(mean=0.0, std=std)

71 + if module.bias is not None:

72 + module.bias.data.zero_()

73 + elif isinstance(module, nn.Embedding):

74 + module.weight.data.normal_(mean=0.0, std=std)

75 + if module.padding_idx is not None:

76 + module.weight.data[module.padding_idx].zero_()

77 +

78 +

79 +def get_initial_dt_bias(num_heads: int) -> torch.Tensor:

80 + dt_min = 0.001

81 + dt_max = 0.1

82 + dt = torch.exp(

83 + torch.rand(num_heads) * (math.log(dt_max) - math.log(dt_min)) +

84 + math.log(dt_min))

85 + dt = torch.clamp(dt, 1e-4)

86 + inv_dt = dt + torch.log(-torch.expm1(-dt))

87 + return inv_dt

88 +

89 +

90 +def is_mamba(config: Plamo2Config, i: int) -> bool:

91 + assert config.mamba_step > 1

92 +

93 + if config.num_hidden_layers <= (config.mamba_step // 2):

94 + # use attention in last layer

95 + return i != config.num_hidden_layers - 1

96 + return (i % config.mamba_step) != (config.mamba_step // 2)

97 +

98 +

99 +# TODO(Shinichi): Replace this with RMSNorm.

100 +def _rms_norm(hidden_states: torch.Tensor, weight: torch.Tensor,

101 + eps: float) -> torch.Tensor:

102 + input_shape = hidden_states.shape

103 + hidden_states = hidden_states.reshape(input_shape[:-1] + weight.shape)

104 + input_dtype = hidden_states.dtype

105 + hidden_states = hidden_states.to(torch.float32)

106 + variance = hidden_states.pow(2).mean(-1, keepdim=True)

107 + hidden_states = hidden_states * torch.rsqrt(variance + eps)

108 + hidden_states = hidden_states.to(input_dtype)

109 + hidden_states = weight * hidden_states

110 + return hidden_states.reshape(input_shape)

111 +

112 +

113 +def _swiglu(h: torch.Tensor) -> torch.Tensor:

114 + h0, h1 = h.chunk(2, dim=-1)

115 + return torch.nn.functional.silu(h0) * h1

116 +

117 +

118 +# Adapted from transformers.models.mamba.modeling_mamba.MambaMixer

119 +class Plamo2MambaMixer(nn.Module):

120 + # TODO(Shinichi): Rebase on Mamba2 implementation.

121 +

122 + def __init__(self,

123 + config: Plamo2Config,

124 + cache_config: CacheConfig,

125 + quant_config: QuantizationConfig,

126 + max_model_len: int,

127 + prefix: str = "",

128 + **kwargs) -> None:

129 + super().__init__()

130 + self.config = config

131 + self.hidden_size = config.hidden_size

132 + self.ssm_state_size = config.mamba_d_state

133 + self.conv_kernel_size = config.mamba_d_conv

134 + self.intermediate_size = (config.mamba_num_heads *

135 + config.hidden_size_per_head)

136 + self.hidden_size_per_head = config.hidden_size_per_head

137 + self.num_heads = config.mamba_num_heads

138 + self.time_step_rank = max(64, self.hidden_size // 16)

139 + self.use_conv_bias = False

140 + self.use_bias = False

141 + self.conv1d = ColumnParallelLinear(

142 + input_size=self.conv_kernel_size,

143 + output_size=self.intermediate_size,

144 + bias=self.use_conv_bias,

145 + )

146 + # unsqueeze to fit conv1d weights shape into the linear weights shape.

147 + # Can't do this in `weight_loader` since it already exists in

148 + # `ColumnParallelLinear` and `set_weight_attrs`

149 + # doesn't allow to override it

150 + self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)

151 +

152 + self.in_proj = MergedColumnParallelLinear(

153 + self.hidden_size,

154 + [self.intermediate_size] * 2,

155 + bias=self.use_bias,

156 + prefix=f"{prefix}.in_proj",

157 + )

158 + # selective projection used to make dt, B and C input dependent

159 + self.bcdt_proj = RowParallelLinear(

160 + self.intermediate_size,

161 + self.time_step_rank + self.ssm_state_size * 2,

162 + bias=False,

163 + prefix=f"{prefix}.bcdt_proj",

164 + )

165 + # time step projection (discretization) -

166 + # In the forward we need to apply dt_proj without the bias,

167 + # as the bias is added in the selective scan kernel.

168 + self.dt_proj = ColumnParallelLinear(

169 + self.time_step_rank,

170 + self.num_heads,

171 + bias=False,

172 + prefix=f"{prefix}.dt_proj",

173 + )

174 + self.dt_bias = torch.nn.Parameter(get_initial_dt_bias(self.num_heads))

175 +

176 + tp_size = get_tensor_model_parallel_world_size()

177 + self.A = nn.Parameter(

178 + torch.empty(

179 + self.intermediate_size // tp_size,

180 + self.ssm_state_size,

181 + dtype=torch.float32,

182 + ))

183 + self.D = nn.Parameter(torch.ones(self.intermediate_size // tp_size))

184 +

185 + set_weight_attrs(self.D, {"weight_loader": sharded_weight_loader(0)})

186 + a_weight_loader = composed_weight_loader(

187 + sharded_weight_loader(0), lambda x: -torch.exp(x.float()))

188 + set_weight_attrs(self.A, {"weight_loader": a_weight_loader})

189 +

190 + self.out_proj = RowParallelLinear(

191 + self.intermediate_size,

192 + self.hidden_size,

193 + bias=self.use_bias,

194 + input_is_parallel=True,

195 + prefix=f"{prefix}.out_proj",

196 + )

197 + # The activation function is fixed to SiLU.

198 + self.activation = "silu"

199 +

200 + self.dt_norm = RMSNorm(self.time_step_rank, eps=config.rms_norm_eps)

201 + self.B_norm = RMSNorm(self.ssm_state_size, eps=config.rms_norm_eps)

202 + self.C_norm = RMSNorm(self.ssm_state_size, eps=config.rms_norm_eps)

203 +

204 + def forward(

205 + self,

206 + hidden_states: torch.Tensor,

207 + mamba_cache_params: MambaCacheParams,

208 + **kwargs,

209 + ) -> torch.Tensor:

210 +

211 + attn_metadata: AttentionMetadata = get_forward_context().attn_metadata

212 +

213 + # 1. Gated MLP's linear projection

214 + projected_states = self.in_proj(hidden_states)[0]

215 + # Reshaping the projected states as in modeling_plamo.py.

216 + length = len(hidden_states)

217 + projected_states = projected_states.reshape(length, self.num_heads, -1)

218 + gate, hidden_states = torch.split(

219 + projected_states,

220 + [self.hidden_size_per_head, self.hidden_size_per_head],

221 + dim=-1)

222 + hidden_states = hidden_states.reshape(length, -1).transpose(0, 1)

223 + gate = gate.reshape(length, -1).transpose(0, 1)

224 +

225 + # 2. Convolution sequence transformation

226 + conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0),

227 + self.conv1d.weight.size(2))

228 +

229 + if attn_metadata.query_start_loc is not None \

230 + and attn_metadata.context_lens_tensor is not None:

231 + # |---------- N-1 iteration --------|

232 + # |---------------- N iteration ---------------------|

233 + # |- tokenA -|......................|-- newTokens ---|

234 + # |---------- context_len ----------|

235 + # |-------------------- seq_len ---------------------|

236 + # |-- query_len ---|

237 + hidden_states = causal_conv1d_fn(

238 + hidden_states,

239 + conv_weights,

240 + self.conv1d.bias,

241 + activation=self.activation,

242 + conv_states=mamba_cache_params.conv_state,

243 + has_initial_state=attn_metadata.context_lens_tensor > 0,

244 + cache_indices=mamba_cache_params.state_indices_tensor,

245 + query_start_loc=attn_metadata.query_start_loc)

246 + else:

247 + hidden_states = causal_conv1d_update(

248 + hidden_states.transpose(0, 1),

249 + mamba_cache_params.conv_state,

250 + conv_weights,

251 + self.conv1d.bias,

252 + self.activation,

253 + conv_state_indices=mamba_cache_params.state_indices_tensor)

254 + hidden_states = hidden_states.transpose(0, 1)

255 +

256 + # 3. State Space Model sequence transformation

257 + # 3.a. input varying initialization of time_step, B and C

258 + ssm_parameters = self.bcdt_proj(hidden_states.transpose(-2, -1))[0]

259 +

260 + # Splitting the ssm_parameters as in modeling_plamo.py.

261 + B, C, time_step = torch.split(

262 + ssm_parameters,

263 + [self.ssm_state_size, self.ssm_state_size, self.time_step_rank],

264 + dim=-1,

265 + )

266 + time_step = self.dt_norm(time_step.contiguous())

267 + B = self.B_norm(B.contiguous())

268 + C = self.C_norm(C.contiguous())

269 +

270 + discrete_time_step = self.dt_proj(time_step)[0].transpose(-2, -1)

271 + # 3.c perform the recurrence y ← SSM(A, B, C)(x)

272 + time_proj_bias = (self.dt_bias.float() if hasattr(

273 + self.dt_proj, "bias") else None)

274 +

275 + # Broadcasting as in modeling_plamo.py.

276 + discrete_time_step = discrete_time_step.transpose(

277 + 0, 1)[..., None].expand(-1, -1, self.hidden_size_per_head)

278 + discrete_time_step = discrete_time_step.reshape(

279 + -1, self.intermediate_size).transpose(0, 1)

280 + time_proj_bias = time_proj_bias[...,

281 + None].expand(-1,

282 + self.hidden_size_per_head)

283 + time_proj_bias = time_proj_bias.reshape(self.intermediate_size)

284 +

285 + if attn_metadata.query_start_loc is not None \

286 + and attn_metadata.context_lens_tensor is not None:

287 + scan_outputs = selective_scan_fn(

288 + hidden_states,

289 + mamba_cache_params.ssm_state,

290 + discrete_time_step,

291 + self.A,

292 + B.transpose(-2, -1),

293 + C.transpose(-2, -1),

294 + self.D.float(),

295 + gate,

296 + time_proj_bias,

297 + delta_softplus=True,

298 + cache_indices=mamba_cache_params.state_indices_tensor,

299 + has_initial_state=attn_metadata.context_lens_tensor > 0,

300 + query_start_loc=attn_metadata.query_start_loc)

301 + else:

302 + scan_outputs = selective_state_update(

303 + mamba_cache_params.ssm_state,

304 + hidden_states.transpose(0, 1),

305 + discrete_time_step.transpose(0, 1),

306 + self.A,

307 + B,

308 + C,

309 + self.D,

310 + gate.transpose(0, 1),

311 + time_proj_bias,

312 + dt_softplus=True,

313 + state_batch_indices=mamba_cache_params.state_indices_tensor)

314 + scan_outputs = scan_outputs.transpose(0, 1)

315 +

316 + # 4. Final linear projection

317 + contextualized_states = self.out_proj(scan_outputs.transpose(-2,

318 + -1))[0]

319 + return contextualized_states

320 +

321 +

322 +class DenseMLP(nn.Module):

323 +

324 + def __init__(

325 + self,

326 + config: Plamo2Config,

327 + quant_config: Optional[QuantizationConfig] = None,

328 + prefix: str = "",

329 + ) -> None:

330 + super().__init__()

331 + self.hidden_size = config.hidden_size

332 + self.intermediate_size = config.intermediate_size

333 + self.gate_up_proj = MergedColumnParallelLinear(

334 + self.hidden_size, [self.intermediate_size] * 2,

335 + bias=False,

336 + prefix=f"{prefix}.gate_up_proj",

337 + quant_config=quant_config)

338 + self.down_proj = RowParallelLinear(self.intermediate_size,

339 + self.hidden_size,

340 + bias=False,

341 + prefix=f"{prefix}.down_proj",

342 + quant_config=quant_config)

343 +

344 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

345 + h = self.gate_up_proj(hidden_states)[0]

346 + h = _swiglu(h)

347 + output, _ = self.down_proj(h)

348 + return output # type: ignore

349 +

350 +

351 +class Plamo2AttentionMixer(nn.Module):

352 +

353 + def __init__(self,

354 + config: Plamo2Config,

355 + cache_config: CacheConfig,

356 + quant_config: QuantizationConfig,

357 + max_model_len: int | None = None,

358 + prefix: str = "",

359 + **kwargs) -> None:

360 + super().__init__()

361 + self.hidden_size = config.hidden_size

362 + tp_size = get_tensor_model_parallel_world_size()

363 + self.total_num_heads = config.num_attention_heads

364 + assert self.total_num_heads % tp_size == 0

365 + self.num_heads = self.total_num_heads // tp_size

366 + self.total_num_kv_heads = config.num_key_value_heads

367 + if self.total_num_kv_heads >= tp_size:

368 + # Number of KV heads is greater than TP size, so we partition

369 + # the KV heads across multiple tensor parallel GPUs.

370 + assert self.total_num_kv_heads % tp_size == 0

371 + else:

372 + # Number of KV heads is less than TP size, so we replicate

373 + # the KV heads across multiple tensor parallel GPUs.

374 + assert tp_size % self.total_num_kv_heads == 0

375 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

376 + self.head_dim = config.hidden_size_per_head

377 + self.q_size = self.num_heads * self.head_dim

378 + self.kv_size = self.num_kv_heads * self.head_dim

379 + self.scaling = self.head_dim**-0.5

380 +

381 + self.qkv_proj = QKVParallelLinear(

382 + config.hidden_size,

383 + self.head_dim,

384 + self.total_num_heads,

385 + self.total_num_kv_heads,

386 + bias=False,

387 + quant_config=quant_config,

388 + )

389 + self.o_proj = RowParallelLinear(self.total_num_heads * self.head_dim,

390 + config.hidden_size,

391 + bias=False,

392 + quant_config=quant_config)

393 +

394 + self.rope_theta = config.rope_theta if hasattr(config,

395 + "rope_theta") else 10000

396 + self.rope_scaling = config.rope_scaling if hasattr(

397 + config, "rope_scaling") else None

398 +

399 + assert max_model_len is not None, "max_model_len must be provided"

400 + self.rotary_emb = get_rope(

401 + self.head_dim,

402 + rotary_dim=self.head_dim,

403 + max_position=max_model_len,

404 + base=self.rope_theta,

405 + rope_scaling=self.rope_scaling,

406 + )

407 + self.q_weight = torch.nn.Parameter(

408 + torch.ones((self.num_heads, config.hidden_size_per_head)))

409 + self.k_weight = torch.nn.Parameter(

410 + torch.ones((self.num_kv_heads, config.hidden_size_per_head)))

411 +

412 + self.attn = Attention(

413 + self.num_heads,

414 + self.head_dim,

415 + self.scaling,

416 + num_kv_heads=self.num_kv_heads,

417 + cache_config=cache_config,

418 + prefix=f"{prefix}.attn",

419 + )

420 +

421 + def forward(

422 + self,

423 + positions: torch.Tensor,

424 + hidden_states: torch.Tensor,

425 + residual: Optional[torch.Tensor],

426 + **kwargs,

427 + ) -> torch.Tensor:

428 + qkv, _ = self.qkv_proj(hidden_states)

429 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

430 + q = _rms_norm(q, self.q_weight, 1e-6)

431 + k = _rms_norm(k, self.k_weight, 1e-6)

432 + q, k = self.rotary_emb(positions, q, k)

433 + attn_output = self.attn(q, k, v)

434 + output, _ = self.o_proj(attn_output)

435 + return output

436 +

437 +

438 +class Plamo2DecoderLayer(nn.Module):

439 +

440 + def __init__(self,

441 + vllm_config: VllmConfig,

442 + layer_idx: int,

443 + max_model_len: int | None = None,

444 + prefix: str = "",

445 + **kwargs) -> None:

446 + super().__init__()

447 + config = vllm_config.model_config.hf_config

448 + cache_config = vllm_config.cache_config

449 + quant_config = vllm_config.quant_config

450 + max_model_len = vllm_config.scheduler_config.max_model_len

451 +

452 + self.is_mamba = is_mamba(config, layer_idx)

453 + if self.is_mamba:

454 + self.mixer = Plamo2MambaMixer(config=config,

455 + cache_config=cache_config,

456 + quant_config=quant_config,

457 + max_model_len=max_model_len,

458 + prefix=f"{prefix}.mixer")

459 + else:

460 + self.mixer = Plamo2AttentionMixer(config=config,

461 + cache_config=cache_config,

462 + quant_config=quant_config,

463 + max_model_len=max_model_len,

464 + prefix=f"{prefix}.mixer")

465 +

466 + self.mlp = DenseMLP(config=config,

467 + quant_config=quant_config,

468 + prefix=f"{prefix}.mlp")

469 + self.pre_mixer_norm = RMSNorm(config.hidden_size,

470 + eps=config.rms_norm_eps)

471 + self.post_mixer_norm = RMSNorm(config.hidden_size,

472 + eps=config.rms_norm_eps)

473 + self.pre_mlp_norm = RMSNorm(config.hidden_size,

474 + eps=config.rms_norm_eps)

475 + self.post_mlp_norm = RMSNorm(config.hidden_size,

476 + eps=config.rms_norm_eps)

477 +

478 + def forward(

479 + self,

480 + positions: torch.Tensor,

481 + hidden_states: torch.Tensor,

482 + residual: Optional[torch.Tensor],

483 + mamba_cache_params: MambaCacheParams,

484 + **kwargs,

485 + ):

486 + if residual is None:

487 + residual = hidden_states

488 + hidden_states = self.pre_mixer_norm(hidden_states)

489 + else:

490 + hidden_states, residual = self.pre_mixer_norm(

491 + hidden_states, residual)

492 +

493 + hidden_states = self.mixer(positions=positions,

494 + hidden_states=hidden_states,

495 + residual=residual,

496 + mamba_cache_params=mamba_cache_params)

497 + hidden_states = self.post_mixer_norm(hidden_states)

498 + # Fully Connected

499 + hidden_states, residual = self.pre_mlp_norm(hidden_states, residual)

500 + hidden_states = self.mlp(hidden_states)

501 + hidden_states = self.post_mlp_norm(hidden_states)

502 + return hidden_states, residual

503 +

504 +

505 +class Plamo2Decoder(torch.nn.Module):

506 +

507 + def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:

508 + super().__init__()

509 + num_hidden_layers = vllm_config.model_config.hf_config.num_hidden_layers

510 +

511 + self.layers = nn.ModuleList([

512 + Plamo2DecoderLayer(vllm_config=vllm_config,

513 + layer_idx=i,

514 + prefix=f"{prefix}.layers.{i}")

515 + for i in range(num_hidden_layers)

516 + ])

517 +

518 + def forward(

519 + self,

520 + positions: torch.Tensor,

521 + hidden_states: torch.Tensor,

522 + residual: Optional[torch.Tensor],

523 + mamba_cache_params: MambaCacheParams,

524 + ) -> torch.Tensor:

525 + mamba_cache_index = 0

526 + for layer in self.layers:

527 + layer_mamba_cache_params = None

528 + if layer.is_mamba:

529 + layer_mamba_cache_params = mamba_cache_params.at_layer_idx(

530 + mamba_cache_index)

531 + mamba_cache_index += 1

532 +

533 + hidden_states, residual = layer(

534 + positions=positions,

535 + hidden_states=hidden_states,

536 + residual=residual,

537 + mamba_cache_params=layer_mamba_cache_params)

538 + return hidden_states, residual

539 +

540 +

541 +class Plamo2Model(Plamo2PreTrainedModel):

542 +

543 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

544 + super().__init__(vllm_config.model_config.hf_config)

545 +

546 + config = vllm_config.model_config.hf_config

547 +

548 + self.config = config

549 + self.padding_idx = config.pad_token_id

550 + self.vocab_size = config.vocab_size

551 + self.org_vocab_size = config.vocab_size

552 +

553 + self.embed_tokens = VocabParallelEmbedding(

554 + self.vocab_size,

555 + config.hidden_size,

556 + org_num_embeddings=config.vocab_size,

557 + prefix=f"{prefix}.embed_tokens",

558 + )

559 + self.layers = Plamo2Decoder(vllm_config, prefix=f"{prefix}.layers")

560 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

561 + self.post_init()

562 +

563 + def forward(

564 + self,

565 + input_ids: torch.Tensor,

566 + positions: torch.Tensor,

567 + mamba_cache_params: MambaCacheParams,

568 + intermediate_tensors: Optional[IntermediateTensors] = None,

569 + inputs_embeds: Optional[torch.Tensor] = None,

570 + ) -> torch.Tensor:

571 + # TODO(Shinichi): Implement pipeline parallelism.

572 + hidden_states = self.embed_tokens(input_ids)

573 + residual = None

574 +

575 + hidden_states, residual = self.layers(

576 + positions=positions,

577 + hidden_states=hidden_states,

578 + residual=residual,

579 + mamba_cache_params=mamba_cache_params)

580 + hidden_states, _ = self.norm(hidden_states, residual)

581 + return hidden_states

582 +

583 +

584 +class Plamo2ForCausalLM(Plamo2PreTrainedModel, HasInnerState, IsHybrid,

585 + SupportsV0Only):

586 + packed_modules_mapping = {

587 + "qkv_proj": [

588 + "q_proj",

589 + "k_proj",

590 + "v_proj",

591 + ],

592 + }

593 +

594 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:

595 + config = vllm_config.model_config.hf_config

596 + scheduler_config = vllm_config.scheduler_config

597 + assert not vllm_config.cache_config.enable_prefix_caching, \

598 + "PLaMo2 currently does not support prefix caching"

599 +

600 + super().__init__(config)

601 + self.config = config

602 + self.vllm_config = vllm_config

603 + self.model_config = vllm_config.model_config

604 + self.scheduler_config = scheduler_config

605 +

606 + # ModelConfig.get_head_size assumes head_dim is set or calculated as

607 + # hidden_size // num_attention_heads. However, this is not always

608 + # the case for PLaMo2, as indicated by the FIXME comment.

609 + self.config.head_dim = self.config.hidden_size_per_head

610 +

611 + self.model = Plamo2Model(vllm_config=vllm_config,

612 + prefix=maybe_prefix(prefix, "model"))

613 + self.vocab_size = self.config.vocab_size

614 + self.unpadded_vocab_size = self.config.vocab_size

615 + num_embeddings = ((self.vocab_size + 15) // 16) * 16

616 + self.lm_head = ParallelLMHead(

617 + num_embeddings,

618 + self.config.hidden_size,

619 + org_num_embeddings=self.config.vocab_size,

620 + padding_size=DEFAULT_VOCAB_PADDING_SIZE,

621 + prefix=f"{prefix}.lm_head",

622 + )

623 + if self.config.tie_word_embeddings:

624 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

625 +

626 + # Used to track and store by the Mamba cache between steps.

627 + self.mamba_cache: Optional[MambaCacheManager] = None

628 +

629 + self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,

630 + self.config.vocab_size)

631 + self.sampler = get_sampler()

632 +

633 + # Initialize weights and apply final processing

634 + self.post_init()

635 +

636 + def forward(self,

637 + input_ids: torch.Tensor,

638 + positions: torch.Tensor,

639 + intermediate_tensors: Optional[IntermediateTensors] = None,

640 + inputs_embeds: Optional[torch.Tensor] = None,

641 + **kwargs):

642 + if self.mamba_cache is None:

643 + num_mamba_layers = self.model_config.get_num_layers_by_block_type(

644 + self.vllm_config.parallel_config, LayerBlockType.mamba)

645 +

646 + self.mamba_cache = MambaCacheManager(

647 + self.vllm_config, self.lm_head.weight.dtype, num_mamba_layers,

648 + *self._get_mamba_cache_shape())

649 +

650 + mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)

651 +

652 + hidden_states = self.model(input_ids, positions, mamba_cache_params,

653 + intermediate_tensors, inputs_embeds)

654 + return hidden_states

655 +

656 + def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):

657 + return self.mamba_cache.copy_inputs_before_cuda_graphs(

658 + input_buffers, **kwargs)

659 +

660 + def get_seqlen_agnostic_capture_inputs(self, batch_size: int):

661 + return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)

662 +

663 + def _get_mamba_cache_shape(

664 + self) -> Tuple[Tuple[int, int], Tuple[int, int]]:

665 + world_size = get_tensor_model_parallel_world_size()

666 + hidden_size = (self.config.mamba_num_heads *

667 + self.config.hidden_size_per_head)

668 + conv_state_shape = (

669 + hidden_size // world_size,

670 + self.config.mamba_d_conv - 1,

671 + )

672 + temporal_state_shape = (

673 + hidden_size // world_size,

674 + self.config.mamba_d_state,

675 + )

676 + return conv_state_shape, temporal_state_shape

677 +

678 + def compute_logits(

679 + self,

680 + hidden_states: torch.Tensor,

681 + sampling_metadata: SamplingMetadata,

682 + ) -> Optional[torch.Tensor]:

683 + logits = self.logits_processor(self.lm_head, hidden_states,

684 + sampling_metadata)

685 + return logits

686 +

687 + def sample(

688 + self,

689 + logits: Optional[torch.Tensor],

690 + sampling_metadata: SamplingMetadata,

691 + ) -> Optional[SamplerOutput]:

692 + next_tokens = self.sampler(logits, sampling_metadata)

693 + return next_tokens

694 +

695 + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):

696 + params_dict = dict(self.named_parameters())

697 + for name, loaded_weight in weights:

698 +

699 + # Both tie_word_embeddings=True and lm_head.weight in the safetensor

700 + # at the same time causes dict key access error.

701 + if name == "lm_head.weight" and self.config.tie_word_embeddings:

702 + assert "lm_head.weight" not in params_dict

703 + continue

704 +

705 + # Update the weight names to be compatible with the vllm version

706 + # of the model.

707 + # Do not change the order of the replacements.

708 + replacements = {

709 + # Rename incompatible weight names.

710 + ".A_log": ".A",

711 + ".B_norm_weight": ".B_norm.weight",

712 + ".C_norm_weight": ".C_norm.weight",

713 + ".dt_norm_weight": ".dt_norm.weight",

714 + }

715 + # Apply replacements based on the defined mappings

716 + for old, new in replacements.items():

717 + if old in name:

718 + name = name.replace(old, new)

719 +

720 + # Broadcast the loaded weight to match the model's parameter shape.

721 + if ".A" in name:

722 + loaded_weight = loaded_weight[:, None, None].expand(

723 + -1, self.config.hidden_size_per_head,

724 + self.config.mamba_d_state)

725 + loaded_weight = loaded_weight.reshape(

726 + -1, self.config.mamba_d_state)

727 + elif ".D" in name:

728 + loaded_weight = loaded_weight[:, None].expand(

729 + -1, self.config.hidden_size_per_head)

730 + loaded_weight = loaded_weight.reshape(-1)

731 + # Offset parameter with vllm's RMSNorm haven't been supported yet.

732 + if ".pre_mixer_norm" in name:

733 + loaded_weight += 1.0

734 + elif ".post_mixer_norm" in name:

735 + loaded_weight += 1.0 / 5

736 + elif ".pre_mlp_norm" in name:

737 + loaded_weight += 1.0

738 + elif ".post_mlp_norm" in name:

739 + loaded_weight += 1.0 / (5**1.5)

740 + elif "model.norm.weight" in name:

741 + loaded_weight += 1.0

742 +

743 + param = params_dict[name]

744 + weight_loader = getattr(param, "weight_loader",

745 + default_weight_loader)

746 + weight_loader(param, loaded_weight)

vllm/model_executor/models/registry.py

@@ -99,6 +99,7 @@ _TEXT_GENERATION_MODELS = {

99 "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),

100 "Phi3SmallForCausalLM": ("phi3_small", "Phi3SmallForCausalLM"),

101 "PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"),

102 + "Plamo2ForCausalLM": ("plamo2", "Plamo2ForCausalLM"),

103 "QWenLMHeadModel": ("qwen", "QWenLMHeadModel"),

104 "Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),

105 "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"),