[Model] Add GraniteSWA and GraniteMoeSWA via existing Granite (#52706)
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Davis Wertheimer committed
Aug 18, 2026 at 23:32 UTC
31305736307ecdfe3398b643240770651b523a84
7 files changed
+186
-207
docs/models/supported_models.md
+2
@@ -405,6 +405,8 @@ th {
405
| `GraniteMoeForCausalLM` | Granite 3.0 MoE, PowerMoE | `ibm-granite/granite-3.0-1b-a400m-base`, `ibm-granite/granite-3.0-3b-a800m-instruct`, `ibm/PowerMoE-3b`, etc. | ✅︎ | ✅︎ |
406
| `GraniteMoeHybridForCausalLM` | Granite 4.0 MoE Hybrid | `ibm-granite/granite-4.0-tiny-preview`, etc. | ✅︎ | ✅︎ |
407
| `GraniteMoeSharedForCausalLM` | Granite MoE Shared | `ibm-research/moe-7b-1b-active-shared-experts` (test model) | ✅︎ | ✅︎ |
408
+| `GraniteMoeSWAForCausalLM` | Granite MoE SWA | `ibm-granite/granite-swash-3b-a600m` | ✅︎ | ✅︎ |
409
+| `GraniteSWAForCausalLM` | Granite SWA | `ibm-granite/granite-swash-2b` | ✅︎ | ✅︎ |
410
| `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ |
411
| `HrmTextForCausalLM` | HRM-Text | `sapientinc/HRM-Text-1B`, etc. | | |
412
| `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
tests/models/language/generation/test_granite.py
+44
-5
@@ -1,14 +1,20 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import pytest
4
+from transformers import GraniteConfig
5
5
-from ...utils import check_logprobs_close
6
+from vllm.model_executor.models.granite import granite_layer_attn_params
7
7
-MODELS = [
8
+from ...utils import check_logprobs_close, check_transformers_version
9
+
10
+# model -> minimum transformers version, or None if unconstrained
11
+MODELS = {
12
# TODO(sang): Sliding window should be tested separately.
9
- "ibm/PowerLM-3b",
10
- "ibm/PowerMoE-3b",
11
-]
13
+ "ibm/PowerLM-3b": None,
14
+ "ibm/PowerMoE-3b": None,
15
+ "ibm-granite/granite-swash-2b": "5.15.1",
16
+ "ibm-granite/granite-swash-3b-a600m": "5.15.1",
17
+}
18
19
20
@pytest.mark.parametrize("model", MODELS)
@@ -25,6 +31,8 @@ def test_models(
31
max_tokens: int,
32
num_logprobs: int,
33
) -> None:
34
+ check_transformers_version(model, min_transformers_version=MODELS[model])
35
+
36
with hf_runner(model, dtype=dtype) as hf_model:
37
hf_outputs = hf_model.generate_greedy_logprobs_limit(
38
example_prompts, max_tokens, num_logprobs
@@ -40,3 +48,34 @@ def test_models(
48
name_0="hf",
49
name_1="vllm",
50
)
51
+
52
+
53
+def test_granite_swa_features_are_off_without_swa_config():
54
+ """A plain Granite config must not pick up sliding windows or sinks."""
55
+ config = GraniteConfig(num_hidden_layers=4)
56
+ theta = config.rope_parameters["rope_theta"]
57
+
58
+ assert [granite_layer_attn_params(config, i) for i in range(4)] == [
59
+ (None, theta, False)
60
+ ] * 4
61
+
62
+
63
+@pytest.mark.parametrize(
64
+ "attention_sinks, expected_sink", [(None, True), (False, False)]
65
+)
66
+def test_granite_swa_features_resolve_per_layer(attention_sinks, expected_sink):
67
+ """`layer_types`/`layer_rope_theta` apply per layer; theta 0 means NoPE."""
68
+ kwargs = {} if attention_sinks is None else {"attention_sinks": attention_sinks}
69
+ config = GraniteConfig(
70
+ num_hidden_layers=3,
71
+ sliding_window=128,
72
+ layer_types=["full_attention", "sliding_attention", "sliding_attention"],
73
+ layer_rope_theta=[10000.0, 0.0, 1000000.0],
74
+ **kwargs,
75
+ )
76
+
77
+ assert [granite_layer_attn_params(config, i) for i in range(3)] == [
78
+ (None, 10000.0, expected_sink),
79
+ (128, 0.0, expected_sink),
80
+ (128, 1000000.0, expected_sink),
81
+ ]
tests/models/registry.py
+6
@@ -320,6 +320,12 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
320
"GraniteMoeSharedForCausalLM": _HfExamplesInfo(
321
"ibm-research/moe-7b-1b-active-shared-experts"
322
),
323
+ "GraniteMoeSWAForCausalLM": _HfExamplesInfo(
324
+ "ibm-granite/granite-swash-3b-a600m", min_transformers_version="5.15.1"
325
+ ),
326
+ "GraniteSWAForCausalLM": _HfExamplesInfo(
327
+ "ibm-granite/granite-swash-2b", min_transformers_version="5.15.1"
328
+ ),
329
"HrmTextForCausalLM": _HfExamplesInfo(
330
"sapientinc/HRM-Text-1B",
331
min_transformers_version="5.9.0",
vllm/model_executor/models/granite.py
+66
-10
@@ -22,7 +22,13 @@
22
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
23
# See the License for the specific language governing permissions and
24
# limitations under the License.
25
-"""Inference-only IBM Granite model compatible with HuggingFace weights."""
25
+"""Inference-only IBM Granite model compatible with HuggingFace weights.
26
+
27
+Also serves the `granite_swa` checkpoints (`GraniteSWAForCausalLM`), supporting
28
+three additional features: per-layer sliding window attention (`layer_types`), a
29
+learnable per-head attention sink (`self_attn.sinks`), and a per-layer RoPE base
30
+(`layer_rope_theta`, with 0 for NoPE).
31
+"""
32
33
from collections.abc import Iterable
34
from itertools import islice
@@ -30,6 +36,7 @@ from itertools import islice
36
import torch
37
from torch import nn
38
from transformers import GraniteConfig
39
+from transformers.configuration_utils import PretrainedConfig
40
41
from vllm.compilation.decorators import support_torch_compile
42
from vllm.config import CacheConfig, VllmConfig
@@ -49,6 +56,8 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
56
ParallelLMHead,
57
VocabParallelEmbedding,
58
)
59
+from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader
60
+from vllm.model_executor.utils import set_weight_attrs
61
from vllm.sequence import IntermediateTensors
62
63
from .interfaces import SupportsLoRA, SupportsPP, SupportsQuant
@@ -56,11 +65,44 @@ from .utils import (
65
AutoWeightsLoader,
66
PPMissingLayer,
67
WeightsMapper,
68
+ extract_layer_index,
69
make_layers,
70
maybe_prefix,
71
)
72
73
74
+def granite_layer_attn_params(
75
+ config: PretrainedConfig, layer_idx: int
76
+) -> tuple[int | None, float, bool]:
77
+ """Resolve one layer's sliding window, RoPE base and sink usage.
78
+
79
+ Plain Granite configs carry no SWA fields and fall back to full
80
+ attention, global RoPE base and no sink. HF SWA checkpoints use
81
+ sinks without a dedicated flag, so assume true when `layer_types`
82
+ is used, and allow `attention_sinks` to override that decision.
83
+
84
+ Returns:
85
+ Sliding window size (`None` for full attention), RoPE base theta (`0`
86
+ for NoPE), and attention sink presence/absence.
87
+ """
88
+ layer_types = getattr(config, "layer_types", None)
89
+ sliding_window = (
90
+ config.sliding_window
91
+ if layer_types is not None and layer_types[layer_idx] == "sliding_attention"
92
+ else None
93
+ )
94
+
95
+ layer_rope_theta = getattr(config, "layer_rope_theta", None)
96
+ rope_theta = (
97
+ layer_rope_theta[layer_idx]
98
+ if layer_rope_theta is not None
99
+ else config.rope_parameters["rope_theta"]
100
+ )
101
+
102
+ has_sink = getattr(config, "attention_sinks", layer_types is not None)
103
+ return sliding_window, rope_theta, has_sink
104
+
105
+
106
class GraniteMLP(nn.Module):
107
def __init__(
108
self,
@@ -154,11 +196,25 @@ class GraniteAttention(nn.Module):
196
prefix=f"{prefix}.o_proj",
197
)
198
157
- self.rotary_emb = get_rope(
158
- self.head_dim,
159
- max_position=max_position_embeddings,
160
- rope_parameters=config.rope_parameters,
199
+ sliding_window, rope_theta, has_sink = granite_layer_attn_params(
200
+ config, extract_layer_index(prefix)
201
)
202
+
203
+ self.use_rope = rope_theta != 0
204
+ if self.use_rope:
205
+ self.rotary_emb = get_rope(
206
+ self.head_dim,
207
+ max_position=max_position_embeddings,
208
+ rope_parameters={**config.rope_parameters, "rope_theta": rope_theta},
209
+ )
210
+
211
+ # Per-head sink, applied in backend as extra logit in the softmax denom
212
+ if has_sink:
213
+ self.sinks = nn.Parameter(torch.empty(self.num_heads), requires_grad=False)
214
+ set_weight_attrs(self.sinks, {"weight_loader": sharded_weight_loader(0)})
215
+ else:
216
+ self.sinks = None
217
+
218
self.attn = Attention(
219
self.num_heads,
220
self.head_dim,
@@ -166,7 +222,9 @@ class GraniteAttention(nn.Module):
222
num_kv_heads=self.num_kv_heads,
223
cache_config=cache_config,
224
quant_config=quant_config,
225
+ per_layer_sliding_window=sliding_window,
226
prefix=f"{prefix}.attn",
227
+ sinks=self.sinks,
228
)
229
230
def forward(
@@ -176,7 +234,8 @@ class GraniteAttention(nn.Module):
234
) -> torch.Tensor:
235
qkv, _ = self.qkv_proj(hidden_states)
236
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
179
- q, k = self.rotary_emb(positions, q, k)
237
+ if self.use_rope:
238
+ q, k = self.rotary_emb(positions, q, k)
239
attn_output = self.attn(q, k, v)
240
output, _ = self.o_proj(attn_output)
241
return output
@@ -413,8 +472,5 @@ class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):
472
# processed with quantization, LoRA, fine-tuning, etc.
473
skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None
474
416
- loader = AutoWeightsLoader(
417
- self,
418
- skip_prefixes=skip_prefixes,
419
- )
475
+ loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
476
return loader.load_weights(weights)
vllm/model_executor/models/granitemoe.py
+57
-150
@@ -26,10 +26,10 @@
26
27
from collections.abc import Iterable
28
from itertools import islice
29
-from typing import Any
29
30
import torch
31
from torch import nn
32
+from transformers.configuration_utils import PretrainedConfig
33
34
from vllm.compilation.decorators import support_torch_compile
35
from vllm.config import CacheConfig, VllmConfig
@@ -39,10 +39,7 @@ from vllm.distributed import (
39
tensor_model_parallel_all_gather,
40
)
41
from vllm.model_executor.layers.attention import Attention
42
-from vllm.model_executor.layers.fused_moe import (
43
- FusedMoEFactory,
44
- fused_moe_make_expert_params_mapping,
45
-)
42
+from vllm.model_executor.layers.fused_moe import FusedMoEFactory
43
from vllm.model_executor.layers.layernorm import RMSNorm
44
from vllm.model_executor.layers.linear import (
45
QKVParallelLinear,
@@ -56,15 +53,20 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
53
ParallelLMHead,
54
VocabParallelEmbedding,
55
)
59
-from vllm.model_executor.model_loader.weight_utils import (
60
- default_weight_loader,
61
- maybe_remap_kv_scale_name,
62
-)
56
+from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader
57
from vllm.model_executor.models.utils import sequence_parallel_chunk
58
+from vllm.model_executor.utils import set_weight_attrs
59
from vllm.sequence import IntermediateTensors
60
61
+from .granite import granite_layer_attn_params
62
from .interfaces import SupportsLoRA, SupportsPP
67
-from .utils import AutoWeightsLoader, is_pp_missing_parameter, make_layers, maybe_prefix
63
+from .utils import (
64
+ AutoWeightsLoader,
65
+ WeightsMapper,
66
+ extract_layer_index,
67
+ make_layers,
68
+ maybe_prefix,
69
+)
70
71
72
class GraniteMoeMoE(nn.Module):
@@ -139,11 +141,11 @@ class GraniteMoeMoE(nn.Module):
141
class GraniteMoeAttention(nn.Module):
142
def __init__(
143
self,
144
+ config: PretrainedConfig,
145
hidden_size: int,
146
num_heads: int,
147
num_kv_heads: int,
148
max_position: int = 4096 * 32,
146
- rope_parameters: dict[str, Any] | None = None,
149
cache_config: CacheConfig | None = None,
150
quant_config: QuantizationConfig | None = None,
151
attention_multiplier: float | None = None,
@@ -190,12 +192,25 @@ class GraniteMoeAttention(nn.Module):
192
quant_config=quant_config,
193
prefix=f"{prefix}.o_proj",
194
)
193
- self.rotary_emb = get_rope(
194
- self.head_dim,
195
- max_position=max_position,
196
- rope_parameters=rope_parameters,
197
- is_neox_style=True,
195
+ sliding_window, rope_theta, has_sink = granite_layer_attn_params(
196
+ config, extract_layer_index(prefix)
197
)
198
+
199
+ self.use_rope = rope_theta != 0
200
+ if self.use_rope:
201
+ self.rotary_emb = get_rope(
202
+ self.head_dim,
203
+ max_position=max_position,
204
+ rope_parameters={**config.rope_parameters, "rope_theta": rope_theta},
205
+ is_neox_style=True,
206
+ )
207
+
208
+ if has_sink:
209
+ self.sinks = nn.Parameter(torch.empty(self.num_heads), requires_grad=False)
210
+ set_weight_attrs(self.sinks, {"weight_loader": sharded_weight_loader(0)})
211
+ else:
212
+ self.sinks = None
213
+
214
self.attn = Attention(
215
self.num_heads,
216
self.head_dim,
@@ -203,7 +218,9 @@ class GraniteMoeAttention(nn.Module):
218
num_kv_heads=self.num_kv_heads,
219
cache_config=cache_config,
220
quant_config=quant_config,
221
+ per_layer_sliding_window=sliding_window,
222
prefix=f"{prefix}.attn",
223
+ sinks=self.sinks,
224
)
225
226
def forward(
@@ -213,7 +230,8 @@ class GraniteMoeAttention(nn.Module):
230
) -> torch.Tensor:
231
qkv, _ = self.qkv_proj(hidden_states)
232
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
216
- q, k = self.rotary_emb(positions, q, k)
233
+ if self.use_rope:
234
+ q, k = self.rotary_emb(positions, q, k)
235
attn_output = self.attn(q, k, v)
236
output, _ = self.o_proj(attn_output)
237
return output
@@ -234,11 +252,11 @@ class GraniteMoeDecoderLayer(nn.Module):
252
253
self.hidden_size = config.hidden_size
254
self.self_attn = GraniteMoeAttention(
255
+ config=config,
256
hidden_size=self.hidden_size,
257
num_heads=config.num_attention_heads,
258
max_position=config.max_position_embeddings,
259
num_kv_heads=config.num_key_value_heads,
241
- rope_parameters=config.rope_parameters,
260
cache_config=cache_config,
261
quant_config=quant_config,
262
prefix=f"{prefix}.self_attn",
@@ -284,6 +302,23 @@ class GraniteMoeDecoderLayer(nn.Module):
302
303
@support_torch_compile
304
class GraniteMoeModel(nn.Module):
305
+ hf_to_vllm_mapper: WeightsMapper = WeightsMapper(
306
+ orig_to_new_suffix={
307
+ # Legacy names to new names
308
+ "moe.input_linear.weight": "moe.experts.gate_up_proj",
309
+ "moe.output_linear.weight": "moe.experts.down_proj",
310
+ ".router.layer.weight": ".gate.weight",
311
+ # Checkpoint name to vLLM name
312
+ ".router.weight": ".gate.weight",
313
+ },
314
+ orig_to_new_stacked={
315
+ # weight_name: (param_name, shard_id)
316
+ ".q_proj": (".qkv_proj", "q"),
317
+ ".k_proj": (".qkv_proj", "k"),
318
+ ".v_proj": (".qkv_proj", "v"),
319
+ },
320
+ )
321
+
322
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
323
super().__init__()
324
@@ -339,135 +374,9 @@ class GraniteMoeModel(nn.Module):
374
hidden_states = self.norm(hidden_states)
375
return hidden_states
376
342
- def _load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
343
- """
344
- This function is copied from `MixtralModel.load_weights`, mainly to
345
- decouple from mixtral, avoiding impact on support like BNB
346
- quantization.
347
- """
348
- stacked_params_mapping = [
349
- # (param_name, shard_name, shard_id)
350
- ("qkv_proj", "q_proj", "q"),
351
- ("qkv_proj", "k_proj", "k"),
352
- ("qkv_proj", "v_proj", "v"),
353
- ]
354
-
355
- # Params for weights, fp8 weight scales, fp8 activation scales
356
- # (param_name, weight_name, expert_id, shard_id)
357
- expert_params_mapping = fused_moe_make_expert_params_mapping(
358
- self,
359
- ckpt_gate_proj_name="w1",
360
- ckpt_down_proj_name="w2",
361
- ckpt_up_proj_name="w3",
362
- num_experts=self.config.num_local_experts,
363
- )
364
-
365
- params_dict = dict(self.named_parameters())
366
- loaded_params: set[str] = set()
367
- for name, loaded_weight in weights:
368
- for param_name, weight_name, shard_id in stacked_params_mapping:
369
- if weight_name not in name:
370
- continue
371
- name = name.replace(weight_name, param_name)
372
- # Skip loading extra bias for GPTQ models.
373
- if (
374
- name.endswith(".bias") or name.endswith("_bias")
375
- ) and name not in params_dict:
376
- continue
377
- # Skip layers on other devices.
378
- if is_pp_missing_parameter(name, self):
379
- continue
380
- if name.endswith("scale"):
381
- # Remapping the name of FP8 kv-scale.
382
- name = maybe_remap_kv_scale_name(name, params_dict)
383
- if name is None:
384
- continue
385
- param = params_dict[name]
386
- weight_loader = param.weight_loader
387
- weight_loader(param, loaded_weight, shard_id)
388
- break
389
- else:
390
- for mapping in expert_params_mapping:
391
- param_name, weight_name, expert_id, shard_id = mapping
392
- if weight_name not in name:
393
- continue
394
- name = name.replace(weight_name, param_name)
395
- # Skip layers on other devices.
396
- if is_pp_missing_parameter(name, self):
397
- continue
398
- if (
399
- name.endswith(".bias") or name.endswith("_bias")
400
- ) and name not in params_dict:
401
- continue
402
- param = params_dict[name]
403
- weight_loader = param.weight_loader
404
- weight_loader(
405
- param,
406
- loaded_weight,
407
- name,
408
- shard_id=shard_id,
409
- expert_id=expert_id,
410
- )
411
- break
412
- else:
413
- # Skip loading extra bias for GPTQ models.
414
- if (
415
- name.endswith(".bias") or name.endswith("_bias")
416
- ) and name not in params_dict:
417
- continue
418
- # Skip layers on other devices.
419
- if is_pp_missing_parameter(name, self):
420
- continue
421
- # Remapping the name of FP8 kv-scale.
422
- name = maybe_remap_kv_scale_name(name, params_dict)
423
- if name is None:
424
- continue
425
-
426
- param = params_dict[name]
427
- weight_loader = getattr(
428
- param, "weight_loader", default_weight_loader
429
- )
430
- weight_loader(param, loaded_weight)
431
- loaded_params.add(name)
432
- return loaded_params
433
-
377
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
435
- new_weights = {}
436
- for n, p in weights:
437
- if n.endswith(".block_sparse_moe.input_linear.weight"):
438
- for e in range(p.size(0)):
439
- w1_name = n.replace(
440
- ".block_sparse_moe.input_linear.weight",
441
- f".block_sparse_moe.experts.{e}.w1.weight",
442
- )
443
- w3_name = n.replace(
444
- ".block_sparse_moe.input_linear.weight",
445
- f".block_sparse_moe.experts.{e}.w3.weight",
446
- )
447
- w1_param, w3_param = p[e].chunk(2, dim=0)
448
- assert w1_name not in new_weights
449
- assert w3_name not in new_weights
450
- new_weights[w1_name] = w1_param
451
- new_weights[w3_name] = w3_param
452
- elif n.endswith(".block_sparse_moe.output_linear.weight"):
453
- for e in range(p.size(0)):
454
- w2_name = n.replace(
455
- ".block_sparse_moe.output_linear.weight",
456
- f".block_sparse_moe.experts.{e}.w2.weight",
457
- )
458
- w2_param = p[e]
459
- assert w2_name not in new_weights
460
- new_weights[w2_name] = w2_param
461
- elif n.endswith(".block_sparse_moe.router.layer.weight"):
462
- gate_name = n.replace(
463
- ".block_sparse_moe.router.layer.weight",
464
- ".block_sparse_moe.gate.weight",
465
- )
466
- assert gate_name not in new_weights
467
- new_weights[gate_name] = p
468
- else:
469
- new_weights[n] = p
470
- return self._load_weights(new_weights.items())
378
+ loader = AutoWeightsLoader(self)
379
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
380
381
382
class GraniteMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
@@ -543,8 +452,6 @@ class GraniteMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
452
)
453
454
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
546
- loader = AutoWeightsLoader(
547
- self,
548
- skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
549
- )
455
+ skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None
456
+ loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
457
return loader.load_weights(weights)
vllm/model_executor/models/granitemoeshared.py
+9
-42
@@ -4,6 +4,10 @@
4
5
The architecture is the same as granitemoe but with the addition of shared
6
experts.
7
+
8
+Also serves the `granitemoe_swa` checkpoints (`GraniteMoeSWAForCausalLM`), which
9
+add the same per-layer sliding window, attention sink and per-layer RoPE support
10
+as `granite_swa` (see `granite.py`).
11
"""
12
13
from collections.abc import Iterable
@@ -85,11 +89,11 @@ class GraniteMoeSharedDecoderLayer(nn.Module):
89
super().__init__()
90
self.hidden_size = config.hidden_size
91
self.self_attn = GraniteMoeAttention(
92
+ config=config,
93
hidden_size=self.hidden_size,
94
num_heads=config.num_attention_heads,
95
max_position=config.max_position_embeddings,
96
num_kv_heads=config.num_key_value_heads,
92
- rope_parameters=config.rope_parameters,
97
cache_config=cache_config,
98
quant_config=quant_config,
99
prefix=f"{prefix}.self_attn",
@@ -207,43 +211,8 @@ class GraniteMoeSharedModel(nn.Module):
211
hidden_states = self.norm(hidden_states)
212
return hidden_states
213
210
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
211
- new_weights = {}
212
- for n, p in weights:
213
- if n.endswith(".block_sparse_moe.input_linear.weight"):
214
- for e in range(p.size(0)):
215
- w1_name = n.replace(
216
- ".block_sparse_moe.input_linear.weight",
217
- f".block_sparse_moe.experts.{e}.w1.weight",
218
- )
219
- w3_name = n.replace(
220
- ".block_sparse_moe.input_linear.weight",
221
- f".block_sparse_moe.experts.{e}.w3.weight",
222
- )
223
- w1_param, w3_param = p[e].chunk(2, dim=0)
224
- assert w1_name not in new_weights
225
- assert w3_name not in new_weights
226
- new_weights[w1_name] = w1_param
227
- new_weights[w3_name] = w3_param
228
- elif n.endswith(".block_sparse_moe.output_linear.weight"):
229
- for e in range(p.size(0)):
230
- w2_name = n.replace(
231
- ".block_sparse_moe.output_linear.weight",
232
- f".block_sparse_moe.experts.{e}.w2.weight",
233
- )
234
- w2_param = p[e]
235
- assert w2_name not in new_weights
236
- new_weights[w2_name] = w2_param
237
- elif n.endswith(".block_sparse_moe.router.layer.weight"):
238
- gate_name = n.replace(
239
- ".block_sparse_moe.router.layer.weight",
240
- ".block_sparse_moe.gate.weight",
241
- )
242
- assert gate_name not in new_weights
243
- new_weights[gate_name] = p
244
- else:
245
- new_weights[n] = p
246
- return GraniteMoeModel._load_weights(self, new_weights.items())
214
+ hf_to_vllm_mapper = GraniteMoeModel.hf_to_vllm_mapper
215
+ load_weights = GraniteMoeModel.load_weights
216
217
218
class GraniteMoeSharedForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
@@ -320,8 +289,6 @@ class GraniteMoeSharedForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
289
)
290
291
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
323
- loader = AutoWeightsLoader(
324
- self,
325
- skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
326
- )
292
+ skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None
293
+ loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
294
return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+2
@@ -124,6 +124,8 @@ _TEXT_GENERATION_MODELS = {
124
"GraniteMoeForCausalLM": ("granitemoe", "GraniteMoeForCausalLM"),
125
"GraniteMoeHybridForCausalLM": ("granitemoehybrid", "GraniteMoeHybridForCausalLM"),
126
"GraniteMoeSharedForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"),
127
+ "GraniteMoeSWAForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"),
128
+ "GraniteSWAForCausalLM": ("granite", "GraniteForCausalLM"),
129
"GritLM": ("gritlm", "GritLM"),
130
"HrmTextForCausalLM": ("hrm_text", "HrmTextForCausalLM"),
131
"HunYuanMoEV1ForCausalLM": ("hunyuan_v1", "HunYuanMoEV1ForCausalLM"),