Exporting google/gemma-3n-e4b-it language_model (decoder) into ONNX format

[arkadaz]

System Info

import os

import numpy as np
import torch
import torch.nn as nn
from transformers import AutoConfig, AutoProcessor, Gemma3nForConditionalGeneration


# Wrapper for the Vision Encoder
class VisionEncoderWrapper(nn.Module):
    def __init__(self, vision_encoder, embed_vision, config):
        super().__init__()
        self.vision_encoder = vision_encoder
        self.embed_vision = embed_vision
        self.config = config

    def forward(self, inputs):
        image_features = self.vision_encoder.timm_model.forward_features(inputs)
        image_features_reshaped = image_features.reshape(
            image_features.shape[0],
            self.config.vision_config.hidden_size,
            self.config.vision_soft_tokens_per_image,
        ).permute(0, 2, 1)
        return self.embed_vision(inputs_embeds=image_features_reshaped)


# Wrapper for the Audio Encoder
class AudioEncoderWrapper(nn.Module):
    def __init__(self, audio_encoder, embed_audio, config):
        super().__init__()
        self.audio_encoder = audio_encoder
        self.embed_audio = embed_audio
        self.config = config

    def forward(self, inputs_features, inputs_features_mask):
        audio_features, audio_mask = model.model.audio_tower(
            inputs_features.to(torch.float32), ~inputs_features_mask
        )
        audio_embeds = self.embed_audio(inputs_embeds=audio_features)
        padding_token_id = 262400 - 1
        audio_padding_toks = torch.tensor([[padding_token_id]], dtype=torch.long)
        audio_padding_embs = self.embed_audio(input_ids=audio_padding_toks)
        audio_features = torch.where(
            audio_mask.unsqueeze(-1), audio_padding_embs, audio_embeds
        )

        audio_batch_size, audio_seq_len, audio_embed_dim = audio_features.shape
        extra_padding_tokens = self.config.audio_soft_tokens_per_image - audio_seq_len
        extra_padding_features = audio_padding_embs.expand(
            audio_batch_size, extra_padding_tokens, audio_embed_dim
        )
        return torch.cat((audio_features, extra_padding_features), dim=1)


# Wrapper for the Token Embedder
class TokenEmbedderWrapper(nn.Module):
    def __init__(self, text_embedder, per_layer_inputs_model, config):
        super().__init__()
        self.text_embedder = text_embedder
        self.per_layer_inputs_model = per_layer_inputs_model
        self.config = config

    def forward(self, input_ids):
        # This layer converts token IDs to dense vectors
        per_layer_inputs_mask = torch.logical_and(
            input_ids >= 0,
            input_ids < self.config.text_config.vocab_size_per_layer_input,
        )
        per_layer_inputs_tokens = torch.where(
            per_layer_inputs_mask,
            input_ids,
            torch.zeros_like(input_ids),
        )
        per_layer_inputs = self.per_layer_inputs_model(per_layer_inputs_tokens)
        return self.text_embedder(input_ids), per_layer_inputs

class LanguageWrapper(nn.Module):
    def __init__(self, lauguage_model):
        super().__init__()
        self.lauguage_model = lauguage_model

    def forward(self, inputs_embeds, position_ids, per_layer_inputs, **kwargs):
        return  self.lauguage_model(inputs_embeds, position_ids, per_layer_inputs, **kwargs)

if __name__ == "__main__":

    model_id = "google/gemma-3n-e4b-it"
    output_dir = "onnx_exported_model_float32"
    os.makedirs(output_dir, exist_ok=True)

    print("🔄 Loading processor and model...")
    # Using float32 for higher precision
    model = Gemma3nForConditionalGeneration.from_pretrained(
        model_id, torch_dtype=torch.float32
    )
    processor = AutoProcessor.from_pretrained(model_id)
    config = AutoConfig.from_pretrained(model_id)

    # Move model to CPU for export
    model.to("cpu")

    # === Dummy input ===
    # Using the same multimodal inputs as before
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg",
                },
                {"type": "text", "text": "Describe this image."},
                {
                    "type": "audio",
                    "audio": "https://raw.githubusercontent.com/google-gemini/gemma-cookbook/refs/heads/main/Demos/sample-data/shopping3.wav",
                },
            ],
        },
    ]

    # Process inputs to get all necessary tensors
    inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    )
    inputs = {k: v.to("cpu") for k, v in inputs.items()}

    # Ensure pixel_values are float32 to match the model
    if "pixel_values" in inputs:
        inputs["pixel_values"] = inputs["pixel_values"].to(torch.float32)

    # === 1. Export Vision Encoder ===
    print("\n📦 [1/4] Exporting Vision Encoder...")
    vision_wrapper = VisionEncoderWrapper(
        model.model.vision_tower, model.model.embed_vision, config
    ).eval()
    torch.onnx.export(
        model=vision_wrapper,
        args=(inputs["pixel_values"],),
        f=os.path.join(output_dir, "vision_encoder.onnx"),
        input_names=["pixel_values"],
        output_names=["image_features"],
        dynamic_axes={
            "pixel_values": {0: "batch_size"},
            "image_features": {0: "batch_size"},
        },
        opset_version=21,
        dynamo=True,
    )
    print("✅ Vision encoder exported.")

    # === 2. Export Audio Encoder ===
    print("\n🔊 [2/4] Exporting Audio Encoder...")
    audio_wrapper = AudioEncoderWrapper(
        model.model.audio_tower, model.model.embed_audio, config
    ).eval()
    torch.onnx.export(
        model=audio_wrapper,
        args=(
            inputs["input_features"].to(torch.float32),
            inputs["input_features_mask"],
        ),
        f=os.path.join(output_dir, "audio_encoder.onnx"),
        input_names=["input_features", "input_features_mask"],
        output_names=["audio_features"],
        dynamic_axes={
            "input_features": {0: "batch_size", 1: "sequence_length"},
            "input_features_mask": {0: "batch_size", 1: "sequence_length"},
            "audio_features": {0: "batch_size"},
        },
        opset_version=21,
        dynamo=True,
    )
    print(f"   - Audio features shape: {inputs['input_features'].shape}")
    print("✅ Audio encoder exported.")

    # === 3. Export Token Embedder ===
    print("\n📝 [3/4] Exporting Token Embedder...")
    token_embedder_wrapper = TokenEmbedderWrapper(
        model.model.language_model.embed_tokens,
        model.language_model.get_per_layer_inputs,
        config,
    ).eval()
    torch.onnx.export(
        model=token_embedder_wrapper,
        args=(inputs["input_ids"],),
        f=os.path.join(output_dir, "token_embedder.onnx"),
        input_names=["input_ids"],
        output_names=["text_embeds"],
        dynamic_axes={
            "input_ids": {0: "batch_size", 1: "sequence_length"},
            "text_embeds": {0: "batch_size", 1: "sequence_length"},
        },
        opset_version=21,
        dynamo=True,
    )

    #################################################################################
    model.language_model.save_pretrained('language-model')`

I am working on exporting the "google/gemma-3n-e4b-it" model to the ONNX format and am encountering issues with the language model (decoder) component.

I have been following the approach outlined in a Colab Notebook on how to export Llava, which was referenced in a related GitHub issue (#38924):

Notebook: https://colab.research.google.com/drive/1mtcxKWHAR7D9LbjzRcTB_b5CURm2jw-z

GitHub Issue: huggingface/transformers#38924

Could you please suggest how to export this model? I have found that a 2B version is available in the ONNX format here: https://huggingface.co/onnx-community/gemma-3n-E2B-it-ONNX/tree/main. I am curious about how to export it myself.

Who can help?

@xenova @zucchini-nlp
need some tutorial about how to export the decoder model 🙏

[zucchini-nlp]

Can you share the error you are encountering?

[arkadaz]

Can you share the error you are encountering?

When I use this code:

model.model.language_model.save_pretrained('language-model')
pipe.tokenizer.save_pretrained('tokenizer')

and this command:

python -m scripts.convert --model_id ./language-model --tokenizer_id ./tokenizer \
--task text-generation-with-past --skip_validation --trust_remote_code

Traceback (most recent call last): File "<frozen runpy>", line 198, in _run_module_as_main File "<frozen runpy>", line 88, in _run_code File "/content/transformers.js/scripts/convert.py", line 456, in <module> main() File "/content/transformers.js/scripts/convert.py", line 341, in main main_export(**export_kwargs) File "/usr/local/lib/python3.11/dist-packages/optimum/exporters/onnx/__main__.py", line 414, in main_export onnx_export_from_model( File "/usr/local/lib/python3.11/dist-packages/optimum/exporters/onnx/convert.py", line 1039, in onnx_export_from_model raise ValueError( ValueError: Trying to export a gemma3n-text model, that is a custom or unsupported architecture, but no custom onnx configuration was passed as custom_onnx_configs. Please refer to https://huggingface.co/docs/optimum/main/en/exporters/onnx/usage_guides/export_a_model#custom-export-of-transformers-models for an example on how to export custom models. Please open an issue at https://github.com/huggingface/optimum/issues if you would like the model type gemma3n-text to be supported natively in the ONNX export.

It is not supported for now, but this model from https://huggingface.co/onnx-community/gemma-3n-E2B-it-ONNX/tree/main can be exported. I just don't know how. Please help.

[xenova]

Hi there 👋 I exported that model with a custom script, which is why it's not working with Optimum (which the transformers.js conversion script uses behind the scenes).

That said, efforts are being made to integrate this model into onnxruntime-genai's model builder, which will ensure a far more optimized conversion (over the torch-exported version).

Since this isn't a transformers issue (it is a feature request for https://github.com/huggingface/optimum/), I'll close the issue here, but you can re-open in the optimum repo, as the errors suggests:

Please open an issue at https://github.com/huggingface/optimum/issues if you would like the model type gemma3n-text to be supported natively in the ONNX export.

[arkadaz]

Hi there 👋 I exported that model with a custom script, which is why it's not working with Optimum (which the transformers.js conversion script uses behind the scenes).

That said, efforts are being made to integrate this model into onnxruntime-genai's model builder, which will ensure a far more optimized conversion (over the torch-exported version).

Since this isn't a transformers issue (it is a feature request for https://github.com/huggingface/optimum/), I'll close the issue here, but you can re-open in the optimum repo, as the errors suggests:

Please open an issue at https://github.com/huggingface/optimum/issues if you would like the model type gemma3n-text to be supported natively in the ONNX export.

Could you please share a custom script? It would be great. I've seen a lot of people trying to do this, but there are no guides for it.

[bil-ash]

Hi there 👋 I exported that model with a custom script, which is why it's not working with Optimum (which the transformers.js conversion script uses behind the scenes).
That said, efforts are being made to integrate this model into onnxruntime-genai's model builder, which will ensure a far more optimized conversion (over the torch-exported version).
Since this isn't a transformers issue (it is a feature request for https://github.com/huggingface/optimum/), I'll close the issue here, but you can re-open in the optimum repo, as the errors suggests:

Please open an issue at https://github.com/huggingface/optimum/issues if you would like the model type gemma3n-text to be supported natively in the ONNX export.

Could you please share a custom script? It would be great. I've seen a lot of people trying to do this, but there are no guides for it.

@xenova I guess that it will take some time for gemma3n support to come in optimum or onnxruntime-genai. Would be nice if you could provide your conversion scripts in the meantime.

[william-lijunxuan]

System Info

import os

import numpy as np
import torch
import torch.nn as nn
from transformers import AutoConfig, AutoProcessor, Gemma3nForConditionalGeneration


# Wrapper for the Vision Encoder
class VisionEncoderWrapper(nn.Module):
    def __init__(self, vision_encoder, embed_vision, config):
        super().__init__()
        self.vision_encoder = vision_encoder
        self.embed_vision = embed_vision
        self.config = config

    def forward(self, inputs):
        image_features = self.vision_encoder.timm_model.forward_features(inputs)
        image_features_reshaped = image_features.reshape(
            image_features.shape[0],
            self.config.vision_config.hidden_size,
            self.config.vision_soft_tokens_per_image,
        ).permute(0, 2, 1)
        return self.embed_vision(inputs_embeds=image_features_reshaped)


# Wrapper for the Audio Encoder
class AudioEncoderWrapper(nn.Module):
    def __init__(self, audio_encoder, embed_audio, config):
        super().__init__()
        self.audio_encoder = audio_encoder
        self.embed_audio = embed_audio
        self.config = config

    def forward(self, inputs_features, inputs_features_mask):
        audio_features, audio_mask = model.model.audio_tower(
            inputs_features.to(torch.float32), ~inputs_features_mask
        )
        audio_embeds = self.embed_audio(inputs_embeds=audio_features)
        padding_token_id = 262400 - 1
        audio_padding_toks = torch.tensor([[padding_token_id]], dtype=torch.long)
        audio_padding_embs = self.embed_audio(input_ids=audio_padding_toks)
        audio_features = torch.where(
            audio_mask.unsqueeze(-1), audio_padding_embs, audio_embeds
        )

        audio_batch_size, audio_seq_len, audio_embed_dim = audio_features.shape
        extra_padding_tokens = self.config.audio_soft_tokens_per_image - audio_seq_len
        extra_padding_features = audio_padding_embs.expand(
            audio_batch_size, extra_padding_tokens, audio_embed_dim
        )
        return torch.cat((audio_features, extra_padding_features), dim=1)


# Wrapper for the Token Embedder
class TokenEmbedderWrapper(nn.Module):
    def __init__(self, text_embedder, per_layer_inputs_model, config):
        super().__init__()
        self.text_embedder = text_embedder
        self.per_layer_inputs_model = per_layer_inputs_model
        self.config = config

    def forward(self, input_ids):
        # This layer converts token IDs to dense vectors
        per_layer_inputs_mask = torch.logical_and(
            input_ids >= 0,
            input_ids < self.config.text_config.vocab_size_per_layer_input,
        )
        per_layer_inputs_tokens = torch.where(
            per_layer_inputs_mask,
            input_ids,
            torch.zeros_like(input_ids),
        )
        per_layer_inputs = self.per_layer_inputs_model(per_layer_inputs_tokens)
        return self.text_embedder(input_ids), per_layer_inputs

class LanguageWrapper(nn.Module):
    def __init__(self, lauguage_model):
        super().__init__()
        self.lauguage_model = lauguage_model

    def forward(self, inputs_embeds, position_ids, per_layer_inputs, **kwargs):
        return  self.lauguage_model(inputs_embeds, position_ids, per_layer_inputs, **kwargs)

if __name__ == "__main__":

    model_id = "google/gemma-3n-e4b-it"
    output_dir = "onnx_exported_model_float32"
    os.makedirs(output_dir, exist_ok=True)

    print("🔄 Loading processor and model...")
    # Using float32 for higher precision
    model = Gemma3nForConditionalGeneration.from_pretrained(
        model_id, torch_dtype=torch.float32
    )
    processor = AutoProcessor.from_pretrained(model_id)
    config = AutoConfig.from_pretrained(model_id)

    # Move model to CPU for export
    model.to("cpu")

    # === Dummy input ===
    # Using the same multimodal inputs as before
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg",
                },
                {"type": "text", "text": "Describe this image."},
                {
                    "type": "audio",
                    "audio": "https://raw.githubusercontent.com/google-gemini/gemma-cookbook/refs/heads/main/Demos/sample-data/shopping3.wav",
                },
            ],
        },
    ]

    # Process inputs to get all necessary tensors
    inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    )
    inputs = {k: v.to("cpu") for k, v in inputs.items()}

    # Ensure pixel_values are float32 to match the model
    if "pixel_values" in inputs:
        inputs["pixel_values"] = inputs["pixel_values"].to(torch.float32)

    # === 1. Export Vision Encoder ===
    print("\n📦 [1/4] Exporting Vision Encoder...")
    vision_wrapper = VisionEncoderWrapper(
        model.model.vision_tower, model.model.embed_vision, config
    ).eval()
    torch.onnx.export(
        model=vision_wrapper,
        args=(inputs["pixel_values"],),
        f=os.path.join(output_dir, "vision_encoder.onnx"),
        input_names=["pixel_values"],
        output_names=["image_features"],
        dynamic_axes={
            "pixel_values": {0: "batch_size"},
            "image_features": {0: "batch_size"},
        },
        opset_version=21,
        dynamo=True,
    )
    print("✅ Vision encoder exported.")

    # === 2. Export Audio Encoder ===
    print("\n🔊 [2/4] Exporting Audio Encoder...")
    audio_wrapper = AudioEncoderWrapper(
        model.model.audio_tower, model.model.embed_audio, config
    ).eval()
    torch.onnx.export(
        model=audio_wrapper,
        args=(
            inputs["input_features"].to(torch.float32),
            inputs["input_features_mask"],
        ),
        f=os.path.join(output_dir, "audio_encoder.onnx"),
        input_names=["input_features", "input_features_mask"],
        output_names=["audio_features"],
        dynamic_axes={
            "input_features": {0: "batch_size", 1: "sequence_length"},
            "input_features_mask": {0: "batch_size", 1: "sequence_length"},
            "audio_features": {0: "batch_size"},
        },
        opset_version=21,
        dynamo=True,
    )
    print(f"   - Audio features shape: {inputs['input_features'].shape}")
    print("✅ Audio encoder exported.")

    # === 3. Export Token Embedder ===
    print("\n📝 [3/4] Exporting Token Embedder...")
    token_embedder_wrapper = TokenEmbedderWrapper(
        model.model.language_model.embed_tokens,
        model.language_model.get_per_layer_inputs,
        config,
    ).eval()
    torch.onnx.export(
        model=token_embedder_wrapper,
        args=(inputs["input_ids"],),
        f=os.path.join(output_dir, "token_embedder.onnx"),
        input_names=["input_ids"],
        output_names=["text_embeds"],
        dynamic_axes={
            "input_ids": {0: "batch_size", 1: "sequence_length"},
            "text_embeds": {0: "batch_size", 1: "sequence_length"},
        },
        opset_version=21,
        dynamo=True,
    )

    #################################################################################
    model.language_model.save_pretrained('language-model')`

I am working on exporting the "google/gemma-3n-e4b-it" model to the ONNX format and am encountering issues with the language model (decoder) component.

I have been following the approach outlined in a Colab Notebook on how to export Llava, which was referenced in a related GitHub issue (#38924):

Notebook: https://colab.research.google.com/drive/1mtcxKWHAR7D9LbjzRcTB_b5CURm2jw-z

GitHub Issue: #38924

Could you please suggest how to export this model? I have found that a 2B version is available in the ONNX format here: https://huggingface.co/onnx-community/gemma-3n-E2B-it-ONNX/tree/main. I am curious about how to export it myself.

Who can help?

@xenova @zucchini-nlp need some tutorial about how to export the decoder model 🙏

Hello,have you successfully convert safetensors to onnx? if you succeed, please let me know! thanks!

[william-lijunxuan]

Hi there 👋 I exported that model with a custom script, which is why it's not working with Optimum (which the transformers.js conversion script uses behind the scenes).

That said, efforts are being made to integrate this model into onnxruntime-genai's model builder, which will ensure a far more optimized conversion (over the torch-exported version).

Since this isn't a transformers issue (it is a feature request for https://github.com/huggingface/optimum/), I'll close the issue here, but you can re-open in the optimum repo, as the errors suggests:

Please open an issue at https://github.com/huggingface/optimum/issues if you would like the model type gemma3n-text to be supported natively in the ONNX export.

Hello, Can you share how to convert safetensors to onnx script to me ? i'm very eager know how! thanks !

[shahizat]

Hello, has someone managed to convert from the .safetensors format of gemma3n to the .onnx format? If so, please share the conversion script.

[hariprasad-ms]

Any updates on gemma-3n models to onnx?

[zucchini-nlp]

To anyone interested, I suggest to open a feature request in optimum repo asking to add support for gemma3 models. The issue is not related to transformers and I cannot give any updates on it

[xenova]

As mentioned in my previous comment, this issue is best suited for the Optimum repo

[arkadaz]

As mentioned in my previous comment, this issue is best suited for the Optimum repo

Hi xenova, could you please share the custom script you mentioned for exporting to ONNX? Thank you!