Skip to content

Commit 09d4c96

Browse files
popovaanrkazantsCopilotecharlaix
authored
[OpenVINO] Support of Gemma3n (#1650)
* [OpenVINO] Support Gemma3n * Add custom PKV generator and inputs * Apply suggestion from @Copilot Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Add patching for language part of Gemma3n * Apply code-formatting * Correct inference for Gemma3n * Add per_layer_inputs, applied vision tower patch. * Add attention layer patching. * Fixed case with no shared layers. * Minor fixes * Minor corrections. * Removed not needed code. * Remove duplicating code. * Remove duplications. * Added test, added comment. * Added tests. * Remove duplications from config. * Remove not needed code. * Add comment. * Code style. * Fix typos. Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> * Minor corrections. * Minor fix. * Update docs/source/openvino/models.mdx Co-authored-by: Ella Charlaix <80481427+echarlaix@users.noreply.github.com> * Update optimum/exporters/openvino/model_configs.py Co-authored-by: Ella Charlaix <80481427+echarlaix@users.noreply.github.com> * Fix for transformers==5.0 * Apply comments. * Minor corrections. * Remove not needed code. * Added link. * Added tests. * Turn off quantization tests for older openvino. * Removed not needed changes. * Returned patching of project_per_layer_inputs. * Minor test fix. * Remove not used input in test. * Changed min transformers version. * Minor alignment. * Test fix. * Test fix. --------- Co-authored-by: Kazantsev, Roman <roman.kazantsev@intel.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Ella Charlaix <80481427+echarlaix@users.noreply.github.com>
1 parent 3517891 commit 09d4c96

11 files changed

Lines changed: 393 additions & 8 deletions

File tree

‎docs/source/openvino/models.mdx‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ Here is the list of the supported architectures :
7373
- Gemma
7474
- Gemma 2
7575
- Gemma 3
76+
- Gemma 3n
7677
- Gemma 4
7778
- Gemma 4 Unified
7879
- GOT-OCR 2.0

‎optimum/exporters/openvino/input_generators.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -430,6 +430,7 @@ def __init__(
430430
self.num_global_key_value_heads = (
431431
getattr(normalized_config.config, "num_global_key_value_heads", None) or self.num_key_value_heads
432432
)
433+
self.model_type = normalized_config.config.model_type
433434

434435
def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
435436
# some layers do not produce their own KV-cache, they use the shared KV-cache
@@ -443,7 +444,7 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int
443444
shape = (
444445
self.batch_size,
445446
self.num_key_value_heads,
446-
self.sliding_window,
447+
self.sequence_length if self.model_type == "gemma3n_text" else self.sliding_window,
447448
self.head_dim,
448449
)
449450
else:

‎optimum/exporters/openvino/model_configs.py‎

Lines changed: 155 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,8 @@
104104
FluxTransformerModelPatcher,
105105
Gemma2ModelPatcher,
106106
Gemma3LMModelPatcher,
107+
Gemma3nImageEmbeddingsModelPatcher,
108+
Gemma3nLMModelPatcher,
107109
Gemma4ImageEmbeddingsModelPatcher,
108110
Gemma4LMModelPatcher,
109111
Gemma4UnifiedImageEmbeddingsModelPatcher,
@@ -1333,6 +1335,21 @@ class Gemma4UnifiedTextOpenVINOConfig(Gemma4TextOpenVINOConfig):
13331335
MAX_TRANSFORMERS_VERSION = "5.10.99"
13341336

13351337

1338+
@register_in_tasks_manager(
1339+
"gemma3n_text",
1340+
*[
1341+
"feature-extraction",
1342+
"feature-extraction-with-past",
1343+
"text-generation",
1344+
"text-generation-with-past",
1345+
"text-classification",
1346+
],
1347+
library_name="transformers",
1348+
)
1349+
class Gemma3nTextOpenVINOConfig(Gemma4TextOpenVINOConfig):
1350+
MIN_TRANSFORMERS_VERSION = "5.0"
1351+
1352+
13361353
@register_in_tasks_manager("deci", *["text-generation", "text-generation-with-past"], library_name="transformers")
13371354
class DeciOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig):
13381355
MAX_TRANSFORMERS_VERSION = "4.57.6"
@@ -3925,7 +3942,7 @@ def __init__(
39253942
behavior=behavior,
39263943
)
39273944
self._behavior = behavior
3928-
if self._behavior == Gemma4ConfigBehavior.VISION_EMBEDDINGS:
3945+
if self._behavior == Gemma4ConfigBehavior.VISION_EMBEDDINGS and config.model_type == "gemma4":
39293946
self.DUMMY_INPUT_GENERATOR_CLASSES = (DummyGemma4VisionInputGenerator,)
39303947
# Attach image_seq_length from preprocessor to normalized config so
39313948
# the dummy input generator can compute the correct number of patches.
@@ -3950,14 +3967,26 @@ def __init__(
39503967
break
39513968
if image_seq_length is not None:
39523969
self._normalized_config.image_seq_length = image_seq_length
3953-
elif self._behavior in (
3954-
Gemma4ConfigBehavior.TEXT_EMBEDDINGS,
3955-
Gemma4ConfigBehavior.TEXT_EMBEDDINGS_PER_LAYER,
3970+
elif (
3971+
self._behavior
3972+
in (
3973+
Gemma4ConfigBehavior.TEXT_EMBEDDINGS,
3974+
Gemma4ConfigBehavior.TEXT_EMBEDDINGS_PER_LAYER,
3975+
)
3976+
and config.model_type == "gemma4"
39563977
):
39573978
self.DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator,)
39583979
self._config = config.text_config
39593980
self._normalized_config = NormalizedTextConfig(self._config)
39603981

3982+
@staticmethod
3983+
def _get_language_model(model):
3984+
if hasattr(model, "language_model"):
3985+
return model.language_model
3986+
if hasattr(model, "model") and hasattr(model.model, "language_model"):
3987+
return model.model.language_model
3988+
raise AttributeError("Gemma3n model does not expose language_model")
3989+
39613990
def with_behavior(self, behavior: Union[str, Gemma4ConfigBehavior]):
39623991
if isinstance(behavior, str) and not isinstance(behavior, Gemma4ConfigBehavior):
39633992
behavior = Gemma4ConfigBehavior(behavior)
@@ -4089,6 +4118,128 @@ def outputs(self) -> Dict[str, Dict[int, str]]:
40894118
return super().outputs
40904119

40914120

4121+
@register_in_tasks_manager("gemma3n", *["image-text-to-text"], library_name="transformers")
4122+
class Gemma3nOpenVINOConfig(Gemma4OpenVINOConfig):
4123+
def __init__(
4124+
self,
4125+
config: "PretrainedConfig",
4126+
task: str = "feature-extraction",
4127+
int_dtype: str = "int64",
4128+
float_dtype: str = "fp32",
4129+
behavior: Gemma4ConfigBehavior = Gemma4ConfigBehavior.VISION_EMBEDDINGS,
4130+
preprocessors: Optional[List[Any]] = None,
4131+
):
4132+
super().__init__(
4133+
config=config,
4134+
task=task,
4135+
int_dtype=int_dtype,
4136+
float_dtype=float_dtype,
4137+
behavior=behavior,
4138+
preprocessors=preprocessors,
4139+
)
4140+
# For VISION_EMBEDDINGS, the vision model contains a reshape that uses
4141+
# vision_soft_tokens_per_image as a fixed constant. The dummy input must
4142+
# produce spatial output matching that value, so we extract image size
4143+
# from the preprocessor to ensure the trace uses the correct dimensions.
4144+
if self._behavior == Gemma4ConfigBehavior.VISION_EMBEDDINGS and config.model_type == "gemma3n":
4145+
image_size = None
4146+
if preprocessors is not None:
4147+
for p in preprocessors:
4148+
ip = getattr(p, "image_processor", p)
4149+
size = getattr(ip, "size", None)
4150+
if size is not None:
4151+
if isinstance(size, dict):
4152+
image_size = size.get("height", size.get("shortest_edge"))
4153+
elif isinstance(size, (int, float)):
4154+
image_size = int(size)
4155+
break
4156+
if image_size is not None:
4157+
self._config.image_size = image_size
4158+
4159+
def with_behavior(self, behavior: Union[str, Gemma4ConfigBehavior]):
4160+
if behavior == Gemma4ConfigBehavior.LANGUAGE:
4161+
return get_vlm_text_generation_config(
4162+
"gemma3n_text",
4163+
self._orig_config.text_config,
4164+
self.int_dtype,
4165+
self.float_dtype,
4166+
model_patcher=Gemma3nLMModelPatcher,
4167+
inputs_update={"per_layer_inputs": {0: "batch_size", 1: "sequence_length", 2: "num_hidden_layers"}},
4168+
)
4169+
return super().with_behavior(behavior)
4170+
4171+
def get_model_for_behavior(self, model, behavior: Union[str, VLMConfigBehavior]):
4172+
if behavior == Gemma4ConfigBehavior.TEXT_EMBEDDINGS_PER_LAYER:
4173+
import torch
4174+
4175+
class PerLayerInputsModule(torch.nn.Module):
4176+
def __init__(self, language_model, vocab_size_per_layer_input: int):
4177+
super().__init__()
4178+
self.language_model = language_model
4179+
self.vocab_size_per_layer_input = vocab_size_per_layer_input
4180+
4181+
# Original code: https://github.com/huggingface/transformers/blob/v5.0.0/src/transformers/models/gemma3n/modeling_gemma3n.py#L2081-L2089
4182+
# Patching is needed to move the preparation of per_layer_inputs outside the language model's forward,
4183+
# as this part requires input_ids to be passed to forward, while for VLM models we pass only input_embeds.
4184+
def forward(self, input_ids: torch.Tensor):
4185+
per_layer_inputs_mask = torch.logical_and(
4186+
input_ids >= 0, input_ids < self.vocab_size_per_layer_input
4187+
)
4188+
per_layer_inputs_tokens = torch.where(
4189+
per_layer_inputs_mask, input_ids, torch.zeros_like(input_ids)
4190+
)
4191+
per_layer_inputs = self.language_model.get_per_layer_inputs(per_layer_inputs_tokens)
4192+
return per_layer_inputs
4193+
4194+
model = PerLayerInputsModule(
4195+
self._get_language_model(model), model.config.text_config.vocab_size_per_layer_input
4196+
)
4197+
return model
4198+
if behavior == VLMConfigBehavior.TEXT_EMBEDDINGS:
4199+
import torch
4200+
4201+
class TextEmbeddingsModule(torch.nn.Module):
4202+
def __init__(self, model):
4203+
super().__init__()
4204+
self.model = model
4205+
4206+
def forward(self, input_ids: torch.Tensor):
4207+
inputs_embeds = self.model.get_input_embeddings()(input_ids)
4208+
vision_mask = torch.logical_and(
4209+
input_ids >= self.model.model.embed_vision.vocab_offset,
4210+
input_ids < self.model.model.embed_audio.vocab_offset,
4211+
)
4212+
dummy_vision_token_id = (
4213+
self.model.model.embed_vision.vocab_offset + self.model.model.embed_vision.vocab_size - 1
4214+
)
4215+
vision_input_ids = torch.where(vision_mask, input_ids, dummy_vision_token_id).to(
4216+
inputs_embeds.device
4217+
)
4218+
vision_embeds = self.model.model.embed_vision(input_ids=vision_input_ids)
4219+
expanded_vision_mask = vision_mask.unsqueeze(-1).expand_as(inputs_embeds)
4220+
inputs_embeds = torch.where(expanded_vision_mask, vision_embeds, inputs_embeds)
4221+
4222+
return inputs_embeds
4223+
4224+
text_embedding = TextEmbeddingsModule(model)
4225+
text_embedding.config = self._get_language_model(model).config
4226+
return text_embedding
4227+
4228+
return super().get_model_for_behavior(model, behavior)
4229+
4230+
def patch_model_for_export(self, model: Union["PreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None):
4231+
model_kwargs = model_kwargs or {}
4232+
if self._behavior == VLMConfigBehavior.VISION_EMBEDDINGS:
4233+
return Gemma3nImageEmbeddingsModelPatcher(self, model, model_kwargs)
4234+
return super().patch_model_for_export(model, model_kwargs)
4235+
4236+
@property
4237+
def inputs(self) -> Dict[str, Dict[int, str]]:
4238+
if self._behavior == Gemma4ConfigBehavior.VISION_EMBEDDINGS:
4239+
return {"pixel_values": {0: "batch_size", 2: "height", 3: "width"}}
4240+
return super().inputs
4241+
4242+
40924243
@register_in_tasks_manager("gemma4_unified", *["image-text-to-text"], library_name="transformers")
40934244
class Gemma4UnifiedOpenVINOConfig(Gemma3OpenVINOConfig):
40944245
# gemma4_unified (e.g. google/gemma-4-12B) reuses the gemma3 VLM scaffolding but has an

0 commit comments

Comments
 (0)