Add GraniteMoeHybrid support for 4.0 (#37658)

* initial config and MLA layer

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* first pass at decoder

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* completion of layers

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* modeling class

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* adding hybrid class to imports

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* fix imports granitemoehybrid

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* fix granitehybrid imports

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* fix granitehybrid import

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* fix generated modeling file

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* add some comments

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* minor fixes in layers

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* add sharedMLP layer

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* correct layer names

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* fixes in mamba config

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* fix mamba config

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* change name of MLP layer

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* fix seq mizer layers

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* correct mamba config

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* fixes in param names

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* enable hybrid model

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* update config

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* fix config granite hybrid

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* fix attention layer

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* cleanup to re-use mamba code

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* keep layer types

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* attention bias cleanup

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* update mamba layer name

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* first pass at tests

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* first pass at tests

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* use granite attention

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* fix: self attn weights

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* pass at making pos_emb optional

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* initialize self_attn only as needed

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* overwrite forward to create HybridMambaCache

Signed-off-by: Sukriti-Sharma4 <[email protected]>

* Log invalid layer types

* Add attention outputs test

* Only emit attentions/logits if not None

* Fix config test hidden size divisibility

* mark granitmoehybrid as stateful

* Initialize mamba convolutional layers

* Formatting fixes

* config docstring, removed some unused attrs

* Fix missing arg in models test

* Fix create and check decoder model test

* support logits to keep in granitemoe

* regen to pass logits_to_keep

* Allow None or rope

* Fix gradient checkpointing

* Add granitemoehybrid as special cache for generate check

* Remove unused MLA refs

* Fix mamba layer mask

* Remove logits to keep from config

* Minor docstring nits

* Update licenses

* Enable cache by default

* map layer types to layer block type

* First pass at granite moe hybrid docs

* Ignore granite moe hybrid in valid checkpoint check

* Align attention interfaces

* regenerate modular granitemoeshared attention interface

* Align granite moe hybrid attn interface

* run formatting

* Handle mamba initialization

* avoid conditional attr defs

* Move hybrid layer validation to config

* Add placeholder integration tests

* Docs nits / Update model names

* Clean up forward conditions

* Use gradient checkpointing layer

* Remove some copied bamba tests + inherit

align test init

delete more tests

Use common layer init with bamba tests

finish test consolidation

* avoid redundant intermediate std var

* use @can_return_tuple

* Remove unused moe state

* make skipped test names consistent

* Fix docstring order

* Add missing toc

* Always create the shared mlp

* Fix name in docstring

* link preview model in docs

---------

Signed-off-by: Sukriti-Sharma4 <[email protected]>
Co-authored-by: Alex-Brooks <[email protected]>

Author: Ssukriti

docs/source/en/_toctree.yml

@@ -495,6 +495,8 @@
         title: Granite
       - local: model_doc/granitemoe
         title: GraniteMoe
+      - local: model_doc/granitemoehybrid
+        title: GraniteMoeHybrid
       - local: model_doc/granitemoeshared
         title: GraniteMoeShared
       - local: model_doc/helium

docs/source/en/model_doc/granitemoehybrid.md

@@ -0,0 +1,64 @@
+<!--Copyright 2025 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+
+-->
+
+# GraniteMoeHybrid
+
+## Overview
+
+
+The `GraniteMoeHybrid` model builds on top of `GraniteMoeSharedModel` and `Bamba`. Its decoding layers consist of state space layers or MoE attention layers with shared experts. By default, the attention layers do not use positional encoding.
+
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_path = "ibm-granite/granite-4.0-tiny-preview"
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+# drop device_map if running on CPU
+model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto")
+model.eval()
+
+# change input text as desired
+prompt = "Write a code to find the maximum value in a list of numbers."
+
+# tokenize the text
+input_tokens = tokenizer(prompt, return_tensors="pt")
+# generate output tokens
+output = model.generate(**input_tokens, max_new_tokens=100)
+# decode output tokens into text
+output = tokenizer.batch_decode(output)
+# loop over the batch to print, in this example the batch size is 1
+for i in output:
+    print(i)
+```
+
+This HF implementation is contributed by [Sukriti Sharma](https://huggingface.co/SukritiSharma) and [Alexander Brooks](https://huggingface.co/abrooks9944).
+
+
+## GraniteMoeHybridConfig
+
+[[autodoc]] GraniteMoeHybridConfig
+
+## GraniteMoeHybridModel
+
+[[autodoc]] GraniteMoeHybridModel
+    - forward
+
+## GraniteMoeHybridForCausalLM
+
+[[autodoc]] GraniteMoeHybridForCausalLM
+    - forward

src/transformers/models/__init__.py

@@ -129,6 +129,7 @@
     from .granite import *
     from .granite_speech import *
     from .granitemoe import *
+    from .granitemoehybrid import *
     from .granitemoeshared import *
     from .grounding_dino import *
     from .groupvit import *

src/transformers/models/auto/configuration_auto.py

@@ -146,6 +146,7 @@
         ("granite", "GraniteConfig"),
         ("granite_speech", "GraniteSpeechConfig"),
         ("granitemoe", "GraniteMoeConfig"),
+        ("granitemoehybrid", "GraniteMoeHybridConfig"),
         ("granitemoeshared", "GraniteMoeSharedConfig"),
         ("granitevision", "LlavaNextConfig"),
         ("graphormer", "GraphormerConfig"),
@@ -509,6 +510,7 @@
         ("granite", "Granite"),
         ("granite_speech", "GraniteSpeech"),
         ("granitemoe", "GraniteMoeMoe"),
+        ("granitemoehybrid", "GraniteMoeHybrid"),
         ("granitemoeshared", "GraniteMoeSharedMoe"),
         ("granitevision", "LLaVA-NeXT"),
         ("graphormer", "Graphormer"),

src/transformers/models/auto/modeling_auto.py

@@ -138,6 +138,7 @@
         ("gptsan-japanese", "GPTSanJapaneseForConditionalGeneration"),
         ("granite", "GraniteModel"),
         ("granitemoe", "GraniteMoeModel"),
+        ("granitemoehybrid", "GraniteMoeHybridModel"),
         ("granitemoeshared", "GraniteMoeSharedModel"),
         ("graphormer", "GraphormerModel"),
         ("grounding-dino", "GroundingDinoModel"),
@@ -558,6 +559,7 @@
         ("gptj", "GPTJForCausalLM"),
         ("granite", "GraniteForCausalLM"),
         ("granitemoe", "GraniteMoeForCausalLM"),
+        ("granitemoehybrid", "GraniteMoeHybridForCausalLM"),
         ("granitemoeshared", "GraniteMoeSharedForCausalLM"),
         ("helium", "HeliumForCausalLM"),
         ("jamba", "JambaForCausalLM"),

src/transformers/models/bamba/modeling_bamba.py

@@ -854,6 +854,7 @@ def torch_forward(
             # Init cache
             if ssm_state is not None and cache_params is not None:
                 cache_params.ssm_states[self.layer_idx].copy_(ssm_state)
+                cache_params.has_previous_state = True
 
         scan_output = self.norm(y, gate)
 

src/transformers/models/bamba/modular_bamba.py

@@ -651,6 +651,7 @@ def torch_forward(
             # Init cache
             if ssm_state is not None and cache_params is not None:
                 cache_params.ssm_states[self.layer_idx].copy_(ssm_state)
+                cache_params.has_previous_state = True
 
         scan_output = self.norm(y, gate)
 

src/transformers/models/granitemoe/configuration_granitemoe.py

@@ -166,6 +166,8 @@ def __init__(
         self.use_cache = use_cache
         self.rope_theta = rope_theta
         self.rope_scaling = rope_scaling
+        # this model has rope embedding type, hardcoded for BC
+        self.position_embedding_type = "rope"
 
         self.attention_bias = attention_bias
         self.attention_dropout = attention_dropout