import torch
from transformers import AutoModelForCausalLM, AutoProcessor

model_id = "yujiepan/gemma-4-moe-tiny-random"
processor = AutoProcessor.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, dtype=torch.bfloat16, device_map="auto"
)
messages = [
    {
        "role": "user",
        "content": [
            {
                "type": "image",
                "url": "https://upload.wikimedia.org/wikipedia/commons/thumb/8/8d/Lightmatter_panda.jpg/960px-Lightmatter_panda.jpg",
            },
            {"type": "text", "text": "What is shown in this image?"},
        ],
    },
    {
        "role": "assistant",
        "content": [{"type": "text", "text": "Dummy response for image"}],
    },
    {
        "role": "user",
        "content": [
            {
                "type": "video",
                "video": "https://github.com/bebechien/gemma/raw/refs/heads/main/videos/ForBiggerBlazes.mp4",
            },
            {"type": "text", "text": "Describe this video."},
        ],
    },
]
inputs = processor.apply_chat_template(
    messages,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
    add_generation_prompt=True,
).to(model.device)
input_len = inputs["input_ids"].shape[-1]
print("input_len:", input_len)
outputs = model.generate(**inputs, max_new_tokens=32)
response = processor.decode(outputs[0], skip_special_tokens=False)
response = response.replace("<|image|>", "I")
response = response.replace("<|video|>", "V")
print(response)

Codes to create this repo:

Click to expand

python

import json
from pathlib import Path

import torch
from huggingface_hub import file_exists, hf_hub_download

from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoProcessor,
    AutoTokenizer,
    Gemma4ForConditionalGeneration,
    GenerationConfig,
    set_seed,
)

source_model_id = "google/gemma-4-26B-A4B-it"
save_folder = "/tmp/yujiepan/gemma-4-moe-tiny-random"

processor = AutoProcessor.from_pretrained(source_model_id)
processor.save_pretrained(save_folder)

with open(
    hf_hub_download(source_model_id, filename="config.json", repo_type="model"),
    "r",
    encoding="utf-8",
) as f:
    config_json = json.load(f)

config_json["text_config"].update(
    {
        "global_head_dim": 64,
        "head_dim": 32,
        "hidden_size": 8,
        # "hidden_size_per_layer_input": 0, # only "E" variants have this
        "intermediate_size": 64,
        "layer_types": [
            "sliding_attention",
            "full_attention",
            "sliding_attention",
            "full_attention",
        ],
        'moe_intermediate_size': 32,
        "num_attention_heads": 8,
        "num_hidden_layers": 4,
        "num_key_value_heads": 4,
        # "num_kv_shared_layers": 0, # only "E" variants have this
    }
)
config_json["vision_config"].update(
    {
        "num_hidden_layers": 2,
        "hidden_size": 8,
        "intermediate_size": 64,
        "head_dim": 32,
        "global_head_dim": 32,
        "num_attention_heads": 4,
        "num_key_value_heads": 4,
    }
)

with open(f"{save_folder}/config.json", "w", encoding="utf-8") as f:
    json.dump(config_json, f, indent=2)

config = AutoConfig.from_pretrained(
    save_folder,
    trust_remote_code=True,
)
print(config)

torch.set_default_dtype(torch.bfloat16)
model = Gemma4ForConditionalGeneration(config)
torch.set_default_dtype(torch.float32)
if file_exists(
    filename="generation_config.json", repo_id=source_model_id, repo_type="model"
):
    model.generation_config = GenerationConfig.from_pretrained(
        source_model_id,
        trust_remote_code=True,
    )
set_seed(42)
model = model.cpu()
all_numels = 0
for name, p in sorted(model.named_parameters()):
    all_numels += p.numel()
with torch.no_grad():
    for name, p in sorted(model.named_parameters()):
        torch.nn.init.normal_(p, 0, 0.2)
        print(name, p.shape, f"{p.numel() / all_numels * 100: .4f}%")
model.save_pretrained(save_folder)

Printing the model:

Click to expand

text

Gemma4ForConditionalGeneration(
  (model): Gemma4Model(
    (vision_tower): Gemma4VisionModel(
      (patch_embedder): Gemma4VisionPatchEmbedder(
        (input_proj): Linear(in_features=768, out_features=8, bias=False)
      )
      (encoder): Gemma4VisionEncoder(
        (rotary_emb): Gemma4VisionRotaryEmbedding()
        (layers): ModuleList(
          (0-1): 2 x Gemma4VisionEncoderLayer(
            (self_attn): Gemma4VisionAttention(
              (q_proj): Gemma4ClippableLinear(
                (linear): Linear(in_features=8, out_features=128, bias=False)
              )
              (k_proj): Gemma4ClippableLinear(
                (linear): Linear(in_features=8, out_features=128, bias=False)
              )
              (v_proj): Gemma4ClippableLinear(
                (linear): Linear(in_features=8, out_features=128, bias=False)
              )
              (o_proj): Gemma4ClippableLinear(
                (linear): Linear(in_features=128, out_features=8, bias=False)
              )
              (q_norm): Gemma4RMSNorm()
              (k_norm): Gemma4RMSNorm()
              (v_norm): Gemma4RMSNorm()
            )
            (mlp): Gemma4VisionMLP(
              (gate_proj): Gemma4ClippableLinear(
                (linear): Linear(in_features=8, out_features=64, bias=False)
              )
              (up_proj): Gemma4ClippableLinear(
                (linear): Linear(in_features=8, out_features=64, bias=False)
              )
              (down_proj): Gemma4ClippableLinear(
                (linear): Linear(in_features=64, out_features=8, bias=False)
              )
              (act_fn): GELUTanh()
            )
            (input_layernorm): Gemma4RMSNorm()
            (post_attention_layernorm): Gemma4RMSNorm()
            (pre_feedforward_layernorm): Gemma4RMSNorm()
            (post_feedforward_layernorm): Gemma4RMSNorm()
          )
        )
      )
      (pooler): Gemma4VisionPooler()
    )
    (language_model): Gemma4TextModel(
      (embed_tokens): Gemma4TextScaledWordEmbedding(262144, 8, padding_idx=0)
      (layers): ModuleList(
        (0): Gemma4TextDecoderLayer(
          (self_attn): Gemma4TextAttention(
            (q_proj): Linear(in_features=8, out_features=256, bias=False)
            (q_norm): Gemma4RMSNorm()
            (k_norm): Gemma4RMSNorm()
            (v_norm): Gemma4RMSNorm()
            (k_proj): Linear(in_features=8, out_features=128, bias=False)
            (v_proj): Linear(in_features=8, out_features=128, bias=False)
            (o_proj): Linear(in_features=256, out_features=8, bias=False)
          )
          (mlp): Gemma4TextMLP(
            (gate_proj): Linear(in_features=8, out_features=64, bias=False)
            (up_proj): Linear(in_features=8, out_features=64, bias=False)
            (down_proj): Linear(in_features=64, out_features=8, bias=False)
            (act_fn): GELUTanh()
          )
          (input_layernorm): Gemma4RMSNorm()
          (post_attention_layernorm): Gemma4RMSNorm()
          (pre_feedforward_layernorm): Gemma4RMSNorm()
          (post_feedforward_layernorm): Gemma4RMSNorm()
          (router): Gemma4TextRouter(
            (norm): Gemma4RMSNorm()
            (proj): Linear(in_features=8, out_features=128, bias=False)
          )
          (experts): Gemma4TextExperts(
            (act_fn): GELUTanh()
          )
          (post_feedforward_layernorm_1): Gemma4RMSNorm()
          (post_feedforward_layernorm_2): Gemma4RMSNorm()
          (pre_feedforward_layernorm_2): Gemma4RMSNorm()
        )
        (1): Gemma4TextDecoderLayer(
          (self_attn): Gemma4TextAttention(
            (q_proj): Linear(in_features=8, out_features=512, bias=False)
            (q_norm): Gemma4RMSNorm()
            (k_norm): Gemma4RMSNorm()
            (v_norm): Gemma4RMSNorm()
            (k_proj): Linear(in_features=8, out_features=128, bias=False)
            (o_proj): Linear(in_features=512, out_features=8, bias=False)
          )
          (mlp): Gemma4TextMLP(
            (gate_proj): Linear(in_features=8, out_features=64, bias=False)
            (up_proj): Linear(in_features=8, out_features=64, bias=False)
            (down_proj): Linear(in_features=64, out_features=8, bias=False)
            (act_fn): GELUTanh()
          )
          (input_layernorm): Gemma4RMSNorm()
          (post_attention_layernorm): Gemma4RMSNorm()
          (pre_feedforward_layernorm): Gemma4RMSNorm()
          (post_feedforward_layernorm): Gemma4RMSNorm()
          (router): Gemma4TextRouter(
            (norm): Gemma4RMSNorm()
            (proj): Linear(in_features=8, out_features=128, bias=False)
          )
          (experts): Gemma4TextExperts(
            (act_fn): GELUTanh()
          )
          (post_feedforward_layernorm_1): Gemma4RMSNorm()
          (post_feedforward_layernorm_2): Gemma4RMSNorm()
          (pre_feedforward_layernorm_2): Gemma4RMSNorm()
        )
        (2): Gemma4TextDecoderLayer(
          (self_attn): Gemma4TextAttention(
            (q_proj): Linear(in_features=8, out_features=256, bias=False)
            (q_norm): Gemma4RMSNorm()
            (k_norm): Gemma4RMSNorm()
            (v_norm): Gemma4RMSNorm()
            (k_proj): Linear(in_features=8, out_features=128, bias=False)
            (v_proj): Linear(in_features=8, out_features=128, bias=False)
            (o_proj): Linear(in_features=256, out_features=8, bias=False)
          )
          (mlp): Gemma4TextMLP(
            (gate_proj): Linear(in_features=8, out_features=64, bias=False)
            (up_proj): Linear(in_features=8, out_features=64, bias=False)
            (down_proj): Linear(in_features=64, out_features=8, bias=False)
            (act_fn): GELUTanh()
          )
          (input_layernorm): Gemma4RMSNorm()
          (post_attention_layernorm): Gemma4RMSNorm()
          (pre_feedforward_layernorm): Gemma4RMSNorm()
          (post_feedforward_layernorm): Gemma4RMSNorm()
          (router): Gemma4TextRouter(
            (norm): Gemma4RMSNorm()
            (proj): Linear(in_features=8, out_features=128, bias=False)
          )
          (experts): Gemma4TextExperts(
            (act_fn): GELUTanh()
          )
          (post_feedforward_layernorm_1): Gemma4RMSNorm()
          (post_feedforward_layernorm_2): Gemma4RMSNorm()
          (pre_feedforward_layernorm_2): Gemma4RMSNorm()
        )
        (3): Gemma4TextDecoderLayer(
          (self_attn): Gemma4TextAttention(
            (q_proj): Linear(in_features=8, out_features=512, bias=False)
            (q_norm): Gemma4RMSNorm()
            (k_norm): Gemma4RMSNorm()
            (v_norm): Gemma4RMSNorm()
            (k_proj): Linear(in_features=8, out_features=128, bias=False)
            (o_proj): Linear(in_features=512, out_features=8, bias=False)
          )
          (mlp): Gemma4TextMLP(
            (gate_proj): Linear(in_features=8, out_features=64, bias=False)
            (up_proj): Linear(in_features=8, out_features=64, bias=False)
            (down_proj): Linear(in_features=64, out_features=8, bias=False)
            (act_fn): GELUTanh()
          )
          (input_layernorm): Gemma4RMSNorm()
          (post_attention_layernorm): Gemma4RMSNorm()
          (pre_feedforward_layernorm): Gemma4RMSNorm()
          (post_feedforward_layernorm): Gemma4RMSNorm()
          (router): Gemma4TextRouter(
            (norm): Gemma4RMSNorm()
            (proj): Linear(in_features=8, out_features=128, bias=False)
          )
          (experts): Gemma4TextExperts(
            (act_fn): GELUTanh()
          )
          (post_feedforward_layernorm_1): Gemma4RMSNorm()
          (post_feedforward_layernorm_2): Gemma4RMSNorm()
          (pre_feedforward_layernorm_2): Gemma4RMSNorm()
        )
      )
      (norm): Gemma4RMSNorm()
      (rotary_emb): Gemma4TextRotaryEmbedding()
    )
    (embed_vision): Gemma4MultimodalEmbedder(
      (embedding_projection): Linear(in_features=8, out_features=8, bias=False)
      (embedding_pre_projection_norm): Gemma4RMSNorm()
    )
  )
  (lm_head): Linear(in_features=8, out_features=262144, bias=False)
)

Test environment:

torch: 2.10.0+cu130
transformers: 5.9.0