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, 12:15 UTC
[Model] Add GraniteMoeHybrid 4.0 model (#17497)
Signed-off-by: Thomas Ortner <boh@zurich.ibm.com> Signed-off-by: Stanislaw Wozniak <stw@zurich.ibm.com> Co-authored-by: Thomas Ortner <boh@zurich.ibm.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com> Co-authored-by: Tyler Michael Smith <tysmith@redhat.com>
Stan Wozniak committed
May 6, 2025 at 06:00 UTC
999328be0d180c40417131252d7a0838c3f1f403
6 files changed
+637
docs/source/models/supported_models.md
+5
@@ -385,6 +385,11 @@ See [this page](#generative-models) for more information on how to use generativ
385
* `ibm-granite/granite-3.0-1b-a400m-base`, `ibm-granite/granite-3.0-3b-a800m-instruct`, `ibm/PowerMoE-3b`, etc.
386
* ✅︎
387
* ✅︎
388
+- * `GraniteMoeHybridForCausalLM`
389
+ * Granite 4.0 MoE Hybrid
390
+ * `ibm-granite/granite-4.0-tiny-preview`, etc.
391
+ * ✅︎
392
+ * ✅︎
393
- * `GraniteMoeSharedForCausalLM`
394
* Granite MoE Shared
395
* `ibm-research/moe-7b-1b-active-shared-experts` (test model)
tests/models/language/generation/test_granitemoehybrid.py
new
+41
@@ -0,0 +1,41 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+import pytest
4
+
5
+from ...utils import check_logprobs_close
6
+
7
+# Path of the checkpoints
8
+MODELS = [
9
+ "ibm-granite/granite-4.0-tiny-preview",
10
+]
11
+
12
+
13
+@pytest.mark.skip(
14
+ reason="Granite 4.0 is not yet available in huggingface transformers")
15
+@pytest.mark.parametrize("model", MODELS)
16
+@pytest.mark.parametrize("dtype", ["float16", "bfloat16"])
17
+@pytest.mark.parametrize("max_tokens", [64])
18
+@pytest.mark.parametrize("num_logprobs", [5])
19
+def test_model_equivalence_to_hf_greedy(
20
+ hf_runner,
21
+ vllm_runner,
22
+ example_prompts,
23
+ model: str,
24
+ dtype: str,
25
+ max_tokens: int,
26
+ num_logprobs: int,
27
+):
28
+ with vllm_runner(model, dtype=dtype) as vllm_model:
29
+ vllm_outputs = vllm_model.generate_greedy_logprobs(
30
+ example_prompts, max_tokens, num_logprobs)
31
+
32
+ with hf_runner(model, dtype=dtype) as hf_model:
33
+ hf_outputs = hf_model.generate_greedy_logprobs_limit(
34
+ example_prompts, max_tokens, num_logprobs)
35
+
36
+ check_logprobs_close(
37
+ outputs_0_lst=hf_outputs,
38
+ outputs_1_lst=vllm_outputs,
39
+ name_0="hf",
40
+ name_1="vllm",
41
+ )
tests/models/language/generation/test_hybrid.py
+3
@@ -23,6 +23,9 @@ SSM_MODELS = [
23
24
HYBRID_MODELS = [
25
"ai21labs/Jamba-tiny-dev",
26
+ # NOTE: ibm-granite/granite-4.0-tiny-preview are skipped currently as
27
+ # it is not yet available in huggingface transformers
28
+ # "ibm-granite/granite-4.0-tiny-preview",
29
# NOTE: Running Plamo2 in transformers implementation requires to install
30
# causal-conv1d package, which is not listed as a test dependency as it's
31
# not compatible with pip-compile.
tests/models/registry.py
+2
@@ -166,6 +166,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
166
{"1b": "EleutherAI/pythia-1.4b"}),
167
"GraniteForCausalLM": _HfExamplesInfo("ibm/PowerLM-3b"),
168
"GraniteMoeForCausalLM": _HfExamplesInfo("ibm/PowerMoE-3b"),
169
+ "GraniteMoeHybridForCausalLM": _HfExamplesInfo("ibm-granite/granite-4.0-tiny-preview", # noqa: E501
170
+ min_transformers_version="4.52.0"), # noqa: E501
171
"GraniteMoeSharedForCausalLM": _HfExamplesInfo("ibm-research/moe-7b-1b-active-shared-experts"), # noqa: E501
172
"Grok1ModelForCausalLM": _HfExamplesInfo("hpcai-tech/grok-1",
173
trust_remote_code=True),
vllm/model_executor/models/granitemoehybrid.py
new
+585
@@ -0,0 +1,585 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""Inference-only GraniteMoeHybrid model."""
3
+# Added by the IBM Team, 2025
4
+from typing import Iterable, Optional, Set, Tuple
5
+
6
+import torch
7
+from torch import nn
8
+from transformers import GraniteMoeHybridConfig
9
+
10
+from vllm.attention.layer import Attention
11
+from vllm.config import CacheConfig, VllmConfig
12
+from vllm.distributed import divide, get_tensor_model_parallel_world_size
13
+from vllm.distributed.parallel_state import get_pp_group
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 ReplicatedLinear
17
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
18
+from vllm.model_executor.layers.mamba.mamba2_metadata import (
19
+ Mamba2Metadata, prepare_mamba2_metadata)
20
+from vllm.model_executor.layers.mamba.mamba_mixer2 import (
21
+ MambaMixer2, extra_groups_for_head_shards)
22
+from vllm.model_executor.layers.quantization import QuantizationConfig
23
+from vllm.model_executor.layers.rotary_embedding import get_rope
24
+from vllm.model_executor.layers.vocab_parallel_embedding import (
25
+ DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
26
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
27
+from vllm.model_executor.models.mamba_cache import (MambaCacheManager,
28
+ MambaCacheParams)
29
+from vllm.model_executor.sampling_metadata import SamplingMetadata
30
+from vllm.sequence import IntermediateTensors
31
+from vllm.utils import LayerBlockType
32
+
33
+from .granitemoe import GraniteMoeMoE
34
+from .granitemoeshared import GraniteMoeSharedMLP
35
+from .interfaces import (HasInnerState, IsHybrid, SupportsLoRA, SupportsPP,
36
+ SupportsQuant, SupportsV0Only)
37
+from .utils import (AutoWeightsLoader, make_empty_intermediate_tensors_factory,
38
+ make_layers, maybe_prefix)
39
+
40
+
41
+class GraniteMoeHybridMambaDecoderLayer(nn.Module):
42
+
43
+ def __init__(self,
44
+ config: GraniteMoeHybridConfig,
45
+ layer_idx: int,
46
+ cache_config: Optional[CacheConfig] = None,
47
+ quant_config: Optional[QuantizationConfig] = None,
48
+ prefix: str = "") -> None:
49
+ super().__init__()
50
+ self.config = config
51
+ self.hidden_size = config.hidden_size
52
+ self.residual_multiplier = config.residual_multiplier
53
+
54
+ self.mamba = MambaMixer2(hidden_size= config.hidden_size,
55
+ ssm_state_size = config.mamba_d_state,
56
+ conv_kernel_size = config.mamba_d_conv,
57
+ intermediate_size = config.mamba_expand *\
58
+ config.hidden_size,
59
+ use_conv_bias = config.mamba_conv_bias,
60
+ use_bias = config.mamba_proj_bias,
61
+ n_groups=config.mamba_n_groups,
62
+ num_heads=config.mamba_n_heads,
63
+ head_dim=config.mamba_d_head,
64
+ rms_norm_eps=config.rms_norm_eps,
65
+ activation=config.hidden_act,
66
+ quant_config=quant_config)
67
+
68
+ self.block_sparse_moe = GraniteMoeMoE(
69
+ num_experts=config.num_local_experts,
70
+ top_k=config.num_experts_per_tok,
71
+ hidden_size=config.hidden_size,
72
+ intermediate_size=config.intermediate_size,
73
+ quant_config=quant_config,
74
+ prefix=f"{prefix}.block_sparse_moe")
75
+
76
+ self.shared_mlp = None if \
77
+ getattr(config, 'shared_intermediate_size', 0) == 0 \
78
+ else GraniteMoeSharedMLP(
79
+ config,
80
+ quant_config=quant_config,
81
+ prefix=f"{prefix}.shared_mlp"
82
+ )
83
+
84
+ self.input_layernorm = RMSNorm(config.hidden_size,
85
+ eps=config.rms_norm_eps)
86
+ self.post_attention_layernorm = RMSNorm(config.hidden_size,
87
+ eps=config.rms_norm_eps)
88
+
89
+ def forward(
90
+ self,
91
+ hidden_states: torch.Tensor,
92
+ residual: Optional[torch.Tensor],
93
+ mamba_cache_params: MambaCacheParams,
94
+ mamba2_metadata: Mamba2Metadata,
95
+ **kwargs,
96
+ ):
97
+ residual = hidden_states
98
+ hidden_states = self.input_layernorm(hidden_states)
99
+ hidden_states = self.mamba(hidden_states, mamba_cache_params,
100
+ mamba2_metadata)
101
+ hidden_states = residual + hidden_states * self.residual_multiplier
102
+
103
+ residual = hidden_states
104
+ hidden_states = self.post_attention_layernorm(hidden_states)
105
+ if self.shared_mlp is None:
106
+ hidden_states = self.block_sparse_moe(hidden_states)
107
+ else:
108
+ # create a copy since block_sparse_moe modifies in-place
109
+ moe_hidden_states = hidden_states.clone()
110
+ moe_hidden_states = self.block_sparse_moe(moe_hidden_states)
111
+ hidden_states = moe_hidden_states + self.shared_mlp(hidden_states)
112
+ del moe_hidden_states
113
+ hidden_states = residual + hidden_states * self.residual_multiplier
114
+
115
+ return hidden_states, residual
116
+
117
+
118
+class GraniteMoeHybridAttentionDecoderLayer(nn.Module):
119
+
120
+ def __init__(
121
+ self,
122
+ config: GraniteMoeHybridConfig,
123
+ layer_idx: int,
124
+ cache_config: Optional[CacheConfig] = None,
125
+ quant_config: Optional[QuantizationConfig] = None,
126
+ prefix: str = "",
127
+ ) -> None:
128
+ super().__init__()
129
+ self.hidden_size = config.hidden_size
130
+ self.residual_multiplier = config.residual_multiplier
131
+
132
+ self.self_attn = GraniteMoeHybridAttention(
133
+ config,
134
+ cache_config=cache_config,
135
+ quant_config=quant_config,
136
+ prefix=f"{prefix}.self_attn")
137
+
138
+ self.block_sparse_moe = GraniteMoeMoE(
139
+ num_experts=config.num_local_experts,
140
+ top_k=config.num_experts_per_tok,
141
+ hidden_size=config.hidden_size,
142
+ intermediate_size=config.intermediate_size,
143
+ quant_config=quant_config,
144
+ prefix=f"{prefix}.block_sparse_moe")
145
+
146
+ self.shared_mlp = None if \
147
+ getattr(config, 'shared_intermediate_size', 0) == 0 \
148
+ else GraniteMoeSharedMLP(
149
+ config,
150
+ quant_config=quant_config,
151
+ prefix=f"{prefix}.shared_mlp"
152
+ )
153
+
154
+ self.input_layernorm = RMSNorm(config.hidden_size,
155
+ eps=config.rms_norm_eps)
156
+ self.post_attention_layernorm = RMSNorm(config.hidden_size,
157
+ eps=config.rms_norm_eps)
158
+
159
+ def forward(
160
+ self,
161
+ positions: torch.Tensor,
162
+ hidden_states: torch.Tensor,
163
+ residual: Optional[torch.Tensor],
164
+ mamba_cache_params: MambaCacheParams,
165
+ mamba2_metadata: Mamba2Metadata,
166
+ ) -> torch.Tensor:
167
+ residual = hidden_states
168
+ hidden_states = self.input_layernorm(hidden_states)
169
+
170
+ hidden_states = self.self_attn(
171
+ positions=positions,
172
+ hidden_states=hidden_states,
173
+ )
174
+ hidden_states = residual + hidden_states * self.residual_multiplier
175
+
176
+ residual = hidden_states
177
+ hidden_states = self.post_attention_layernorm(hidden_states)
178
+ if self.shared_mlp is None:
179
+ hidden_states = self.block_sparse_moe(hidden_states)
180
+ else:
181
+ # create a copy since block_sparse_moe modifies in-place
182
+ moe_hidden_states = hidden_states.clone()
183
+ moe_hidden_states = self.block_sparse_moe(moe_hidden_states)
184
+ hidden_states = moe_hidden_states + self.shared_mlp(hidden_states)
185
+ del moe_hidden_states
186
+ hidden_states = residual + hidden_states * self.residual_multiplier
187
+
188
+ return hidden_states, residual
189
+
190
+
191
+class GraniteMoeHybridAttention(nn.Module):
192
+
193
+ def __init__(
194
+ self,
195
+ config: GraniteMoeHybridConfig,
196
+ cache_config: Optional[CacheConfig] = None,
197
+ quant_config: Optional[QuantizationConfig] = None,
198
+ prefix: str = "",
199
+ ) -> None:
200
+ super().__init__()
201
+ self.causal = True
202
+ self.hidden_size = config.hidden_size
203
+ self.attention_bias = config.attention_bias
204
+ self.attention_multiplier = config.attention_multiplier
205
+ self.num_heads = config.num_attention_heads
206
+ self.head_dim = self.hidden_size // self.num_heads
207
+ self.num_key_value_heads = config.num_key_value_heads
208
+
209
+ self.q_proj = ReplicatedLinear(self.hidden_size,
210
+ self.num_heads * self.head_dim,
211
+ bias=self.attention_bias,
212
+ quant_config=quant_config,
213
+ prefix=f"{prefix}.q_proj")
214
+
215
+ self.k_proj = ReplicatedLinear(self.hidden_size,
216
+ self.num_key_value_heads *
217
+ self.head_dim,
218
+ bias=self.attention_bias,
219
+ quant_config=quant_config,
220
+ prefix=f"{prefix}.k_proj")
221
+
222
+ self.v_proj = ReplicatedLinear(self.hidden_size,
223
+ self.num_key_value_heads *
224
+ self.head_dim,
225
+ bias=self.attention_bias,
226
+ quant_config=quant_config,
227
+ prefix=f"{prefix}.v_proj")
228
+
229
+ self.o_proj = ReplicatedLinear(self.hidden_size,
230
+ self.hidden_size,
231
+ bias=self.attention_bias,
232
+ quant_config=quant_config,
233
+ prefix=f"{prefix}.o_proj")
234
+
235
+ if config.position_embedding_type == "rope":
236
+ self.rotary_emb = get_rope(
237
+ self.head_dim,
238
+ rotary_dim=self.head_dim,
239
+ max_position=config.max_position_embeddings,
240
+ base=int(config.rope_theta),
241
+ rope_scaling=config.rope_scaling \
242
+ if hasattr(config, "rope_scaling") \
243
+ and config.rope_scaling is not None else None,
244
+ is_neox_style=True,
245
+ )
246
+ else:
247
+ self.rotary_emb = None
248
+
249
+ self.attn = Attention(self.num_heads,
250
+ self.head_dim,
251
+ self.attention_multiplier,
252
+ num_kv_heads=self.num_key_value_heads,
253
+ cache_config=cache_config,
254
+ quant_config=quant_config,
255
+ prefix=f"{prefix}.attn")
256
+
257
+ def forward(
258
+ self,
259
+ positions: torch.Tensor,
260
+ hidden_states: torch.Tensor,
261
+ ) -> torch.Tensor:
262
+
263
+ query = self.q_proj(hidden_states)[0]
264
+ key = self.k_proj(hidden_states)[0]
265
+ value = self.v_proj(hidden_states)[0]
266
+
267
+ if self.rotary_emb is not None:
268
+ query, key = self.rotary_emb(positions, query, key)
269
+
270
+ hidden_states = self.attn(query, key, value)
271
+ del query, key, value
272
+
273
+ hidden_states = self.o_proj(hidden_states)[0]
274
+ return hidden_states
275
+
276
+
277
+ALL_DECODER_LAYER_TYPES = {
278
+ "attention": GraniteMoeHybridAttentionDecoderLayer,
279
+ "mamba": GraniteMoeHybridMambaDecoderLayer,
280
+}
281
+
282
+
283
+class GraniteMoeHybridModel(nn.Module):
284
+
285
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
286
+ super().__init__()
287
+
288
+ config = vllm_config.model_config.hf_config
289
+ cache_config = vllm_config.cache_config
290
+ quant_config = vllm_config.quant_config
291
+ lora_config = vllm_config.lora_config
292
+
293
+ self.config = config
294
+ lora_vocab = ((lora_config.lora_extra_vocab_size *
295
+ (lora_config.max_loras or 1)) if lora_config else 0)
296
+ self.vocab_size = config.vocab_size + lora_vocab
297
+ self.org_vocab_size = config.vocab_size
298
+
299
+ self.embed_tokens = VocabParallelEmbedding(
300
+ self.vocab_size,
301
+ config.hidden_size,
302
+ org_num_embeddings=config.vocab_size,
303
+ )
304
+ self.embedding_multiplier = config.embedding_multiplier
305
+
306
+ def get_layer(prefix: str):
307
+ layer_idx = int(prefix.rsplit(".", 1)[1])
308
+ layer_class = ALL_DECODER_LAYER_TYPES[
309
+ config.layer_types[layer_idx]]
310
+ return layer_class(
311
+ config,
312
+ layer_idx,
313
+ cache_config,
314
+ quant_config=quant_config,
315
+ prefix=prefix,
316
+ )
317
+
318
+ self.start_layer, self.end_layer, self.layers = make_layers(
319
+ config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers")
320
+ self.make_empty_intermediate_tensors = (
321
+ make_empty_intermediate_tensors_factory(
322
+ ["hidden_states", "residual"], config.hidden_size))
323
+
324
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
325
+
326
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
327
+ return self.embed_tokens(input_ids)
328
+
329
+ def forward(
330
+ self,
331
+ input_ids: torch.Tensor,
332
+ positions: torch.Tensor,
333
+ mamba_cache_params: MambaCacheParams,
334
+ intermediate_tensors: Optional[IntermediateTensors] = None,
335
+ inputs_embeds: Optional[torch.Tensor] = None,
336
+ ) -> torch.Tensor:
337
+
338
+ attn_metadata = get_forward_context().attn_metadata
339
+ mamba2_metadata = prepare_mamba2_metadata(
340
+ chunk_size=self.config.mamba_chunk_size,
341
+ input_ids=input_ids,
342
+ attn_metadata=attn_metadata,
343
+ )
344
+
345
+ if get_pp_group().is_first_rank:
346
+ if inputs_embeds is not None:
347
+ hidden_states = inputs_embeds
348
+ else:
349
+ hidden_states = self.get_input_embeddings(input_ids)
350
+ hidden_states = hidden_states * self.embedding_multiplier
351
+ residual = None
352
+ else:
353
+ if intermediate_tensors is None:
354
+ raise RuntimeError('Intermediate tensors may not be None!')
355
+ hidden_states = intermediate_tensors["hidden_states"]
356
+ residual = intermediate_tensors["residual"]
357
+
358
+ num_attn = 0
359
+ for i in range(len(self.layers)):
360
+ layer = self.layers[i]
361
+ if isinstance(layer, GraniteMoeHybridAttentionDecoderLayer):
362
+ num_attn += 1
363
+
364
+ layer_mamba_cache_params = None
365
+ if isinstance(layer, GraniteMoeHybridMambaDecoderLayer):
366
+ layer_mamba_cache_params = mamba_cache_params.at_layer_idx(
367
+ i - num_attn)
368
+
369
+ hidden_states, residual = layer(
370
+ positions=positions,
371
+ hidden_states=hidden_states,
372
+ residual=residual,
373
+ mamba_cache_params=layer_mamba_cache_params,
374
+ mamba2_metadata=mamba2_metadata)
375
+
376
+ if not get_pp_group().is_last_rank:
377
+ return IntermediateTensors({
378
+ "hidden_states": hidden_states,
379
+ "residual": residual
380
+ })
381
+
382
+ hidden_states = self.norm(hidden_states)
383
+ return hidden_states
384
+
385
+ def load_weights(self, weights: Iterable[Tuple[str,
386
+ torch.Tensor]]) -> Set[str]:
387
+ params_dict = dict(self.named_parameters())
388
+ loaded_params: Set[str] = set()
389
+
390
+ def _load(n, p):
391
+ param = params_dict[n]
392
+ weight_loader = getattr(param, "weight_loader",
393
+ default_weight_loader)
394
+ weight_loader(param, p)
395
+ loaded_params.add(n)
396
+
397
+ def _load_expert(n, p, name, shard_id, expert_id):
398
+ param = params_dict[n]
399
+ weight_loader = getattr(param, "weight_loader",
400
+ default_weight_loader)
401
+ weight_loader(param,
402
+ p,
403
+ name,
404
+ shard_id=shard_id,
405
+ expert_id=expert_id)
406
+ loaded_params.add(n)
407
+
408
+ for n, p in weights:
409
+ if "A_log" in n:
410
+ n = n.replace("A_log", "A")
411
+
412
+ # Logic analogous to: https://github.com/vllm-project/vllm/blob/f49e5aff11c986ed4d45202b1716c5d74786efa9/vllm/model_executor/models/granitemoeshared.py#L215
413
+ # Mapping different experts' layout:
414
+ # from HF (input_linear, output_linear, router)
415
+ # to vLLM (experts_w13({e}.w1, {e}.w2), experts_w3({e}.w3), gate)
416
+ if n.endswith('.block_sparse_moe.input_linear.weight'):
417
+ for e in range(p.size(0)):
418
+ w1_name = n.replace(
419
+ '.block_sparse_moe.input_linear.weight',
420
+ f".block_sparse_moe.experts.{e}.w1.weight")
421
+ w3_name = n.replace(
422
+ '.block_sparse_moe.input_linear.weight',
423
+ f".block_sparse_moe.experts.{e}.w3.weight")
424
+ w1_param, w3_param = p[e].chunk(2, dim=0)
425
+ _load_expert(n.replace('.input_linear.', '.experts.w13_'),
426
+ w1_param,
427
+ w1_name,
428
+ shard_id='w1',
429
+ expert_id=e)
430
+ _load_expert(n.replace('.input_linear.', '.experts.w13_'),
431
+ w3_param,
432
+ w3_name,
433
+ shard_id='w3',
434
+ expert_id=e)
435
+ elif n.endswith('.block_sparse_moe.output_linear.weight'):
436
+ for e in range(p.size(0)):
437
+ w2_name = n.replace(
438
+ '.block_sparse_moe.output_linear.weight',
439
+ f".block_sparse_moe.experts.{e}.w2.weight")
440
+ w2_param = p[e]
441
+ _load_expert(n.replace('.output_linear.', '.experts.w2_'),
442
+ w2_param,
443
+ w2_name,
444
+ shard_id='w2',
445
+ expert_id=e)
446
+ elif n.endswith('.block_sparse_moe.router.layer.weight'):
447
+ gate_name = n.replace('.block_sparse_moe.router.layer.weight',
448
+ ".block_sparse_moe.gate.weight")
449
+ _load(gate_name, p)
450
+ else:
451
+ _load(n, p)
452
+
453
+ return loaded_params
454
+
455
+
456
+class GraniteMoeHybridForCausalLM(nn.Module, HasInnerState, SupportsLoRA,
457
+ SupportsPP, IsHybrid, SupportsV0Only,
458
+ SupportsQuant):
459
+ packed_modules_mapping = {}
460
+ embedding_modules = {
461
+ "embed_tokens": "input_embeddings",
462
+ "lm_head": "output_embeddings",
463
+ }
464
+ embedding_padding_modules = ["lm_head"]
465
+
466
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
467
+ super().__init__()
468
+
469
+ config = vllm_config.model_config.hf_config
470
+ self.vllm_config = vllm_config
471
+ self.model_config = vllm_config.model_config
472
+ cache_config = vllm_config.cache_config
473
+ lora_config = vllm_config.lora_config
474
+ scheduler_config = vllm_config.scheduler_config
475
+ if cache_config.enable_prefix_caching:
476
+ raise RuntimeError(
477
+ "GraniteMoeHybrid currently does not support prefix caching")
478
+
479
+ self.quant_config = vllm_config.quant_config
480
+ self.config = config
481
+ self.scheduler_config = scheduler_config
482
+ self.model = GraniteMoeHybridModel(vllm_config=vllm_config,
483
+ prefix=maybe_prefix(
484
+ prefix, "model"))
485
+ self.unpadded_vocab_size = config.vocab_size
486
+ if lora_config:
487
+ self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
488
+
489
+ self.lm_head = ParallelLMHead(
490
+ self.unpadded_vocab_size,
491
+ config.hidden_size,
492
+ org_num_embeddings=config.vocab_size,
493
+ padding_size=DEFAULT_VOCAB_PADDING_SIZE
494
+ # We need bigger padding if using lora for kernel
495
+ # compatibility
496
+ if not lora_config else lora_config.lora_vocab_padding_size,
497
+ quant_config=self.quant_config,
498
+ prefix=maybe_prefix(prefix, "lm_head"))
499
+ if config.tie_word_embeddings:
500
+ self.lm_head.weight = self.model.embed_tokens.weight
501
+ self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
502
+ config.vocab_size,
503
+ scale=1 /
504
+ self.config.logits_scaling)
505
+
506
+ # Used to track and store by the Mamba cache between steps.
507
+ self.mamba_cache: Optional[MambaCacheManager] = None
508
+
509
+ self.make_empty_intermediate_tensors = (
510
+ self.model.make_empty_intermediate_tensors)
511
+
512
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
513
+ return self.model.get_input_embeddings(input_ids)
514
+
515
+ def forward(self,
516
+ input_ids: torch.Tensor,
517
+ positions: torch.Tensor,
518
+ intermediate_tensors: Optional[IntermediateTensors] = None,
519
+ inputs_embeds: Optional[torch.Tensor] = None,
520
+ **kwargs):
521
+ if self.mamba_cache is None:
522
+ num_mamba_layers = self.model_config.get_num_layers_by_block_type(
523
+ self.vllm_config.parallel_config, LayerBlockType.mamba)
524
+ self.mamba_cache = MambaCacheManager(
525
+ self.vllm_config, self.model_config.dtype, num_mamba_layers,
526
+ *self._get_mamba_cache_shape())
527
+
528
+ mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
529
+ hidden_states = self.model(input_ids, positions, mamba_cache_params,
530
+ intermediate_tensors, inputs_embeds)
531
+
532
+ return hidden_states
533
+
534
+ def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
535
+ return self.mamba_cache.copy_inputs_before_cuda_graphs(
536
+ input_buffers, **kwargs)
537
+
538
+ def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
539
+ return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
540
+
541
+ def _get_mamba_cache_shape(
542
+ self) -> Tuple[Tuple[int, int], Tuple[int, int]]:
543
+ world_size = get_tensor_model_parallel_world_size()
544
+ hidden_size = self.config.hidden_size
545
+
546
+ conv_state_shape, temporal_state_shape = None, None
547
+
548
+ intermediate_size = self.config.mamba_expand * hidden_size
549
+
550
+ # if n_groups is not divisible by world_size, need to extend the shards
551
+ # to ensure all groups needed by a head is sharded along with it
552
+ n_groups = (self.config.mamba_n_groups + extra_groups_for_head_shards(
553
+ self.config.mamba_n_groups, world_size))
554
+
555
+ # - heads and n_groups are TP-ed
556
+ conv_dim = (intermediate_size +
557
+ 2 * n_groups * self.config.mamba_d_state)
558
+ conv_state_shape = (
559
+ divide(conv_dim, world_size),
560
+ self.config.mamba_d_conv - 1,
561
+ )
562
+
563
+ # These are not TP-ed as they depend on A, dt_bias, D
564
+ # - they are typically small
565
+ # e.g., (h_heads, d_head, d_state) = (128, 64, 128)
566
+ temporal_state_shape = (
567
+ divide(self.config.mamba_n_heads, world_size),
568
+ self.config.mamba_d_head,
569
+ self.config.mamba_d_state,
570
+ )
571
+ return conv_state_shape, temporal_state_shape
572
+
573
+ def compute_logits(
574
+ self,
575
+ hidden_states: torch.Tensor,
576
+ sampling_metadata: SamplingMetadata,
577
+ ) -> Optional[torch.Tensor]:
578
+ logits = self.logits_processor(self.lm_head, hidden_states,
579
+ sampling_metadata)
580
+ return logits
581
+
582
+ def load_weights(self, weights: Iterable[Tuple[str,
583
+ torch.Tensor]]) -> Set[str]:
584
+ loader = AutoWeightsLoader(self)
585
+ return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+1
@@ -64,6 +64,7 @@ _TEXT_GENERATION_MODELS = {
64
"GPTNeoXForCausalLM": ("gpt_neox", "GPTNeoXForCausalLM"),
65
"GraniteForCausalLM": ("granite", "GraniteForCausalLM"),
66
"GraniteMoeForCausalLM": ("granitemoe", "GraniteMoeForCausalLM"),
67
+ "GraniteMoeHybridForCausalLM": ("granitemoehybrid", "GraniteMoeHybridForCausalLM"), # noqa: E501
68
"GraniteMoeSharedForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"), # noqa: E501
69
"GritLM": ("gritlm", "GritLM"),
70
"Grok1ModelForCausalLM": ("grok1", "Grok1ForCausalLM"),