tiny-random-gemma4-unified-it

import torch
from transformers import AutoProcessor
from transformers.models.gemma4_unified.configuration_gemma4_unified import (
    Gemma4UnifiedConfig,
    Gemma4UnifiedTextConfig,
    Gemma4UnifiedVisionConfig,
)
from transformers.models.gemma4_unified.modeling_gemma4_unified import Gemma4UnifiedForConditionalGeneration


save_dir = "./tiny-random-gemma4-unified-it"

# Tiny text config mirroring the 12B gemma4_unified architecture:
#  - PLE disabled (no hidden_size_per_layer_input field exists on unified text config)
#  - attention_k_eq_v=True so full-attention layers fuse v_proj into k_proj
#  - num_global_key_value_heads=1, global_head_dim larger than head_dim
#  - num_kv_shared_layers=0 (12B does not share KV across layers)
#  - use_bidirectional_attention="vision"
text_config = Gemma4UnifiedTextConfig(
    hidden_size=32,
    intermediate_size=64,
    num_hidden_layers=4,
    num_attention_heads=4,
    num_key_value_heads=2,
    head_dim=16,
    global_head_dim=32,
    num_global_key_value_heads=1,
    vocab_size=262144,
    max_position_embeddings=512,
    rms_norm_eps=1e-6,
    hidden_activation="gelu_pytorch_tanh",
    sliding_window=64,
    layer_types=["sliding_attention", "sliding_attention", "sliding_attention", "full_attention"],
    num_kv_shared_layers=0,
    attention_k_eq_v=True,
    use_double_wide_mlp=False,
    use_bidirectional_attention="vision",
    final_logit_softcapping=30.0,
    tie_word_embeddings=True,
)

# Vision is an encoder-free embedder: model_patch_size = patch_size * pooling_kernel_size.
# mm_embed_dim / output_proj_dims must match the text hidden_size.
vision_config = Gemma4UnifiedVisionConfig(
    patch_size=16,
    pooling_kernel_size=3,
    mm_embed_dim=32,
    output_proj_dims=32,
    mm_posemb_size=128,
    rms_norm_eps=1e-6,
)

config = Gemma4UnifiedConfig(
    text_config=text_config.to_dict(),
    vision_config=vision_config.to_dict(),
    audio_config=None,
    boi_token_id=255999,
    eoi_token_id=258882,
    image_token_id=258880,
    video_token_id=258884,
    boa_token_id=256000,
    eoa_token_index=258883,
    audio_token_id=258881,
    tie_word_embeddings=True,
)

# Seed before init so the random weights are reproducible. This seed produces a fixture
# whose greedy generation has no near-ties, so OV-vs-transformers token equality is stable
# under the small (~1e-4) numerical differences of OpenVINO inference.
torch.manual_seed(42)
model = Gemma4UnifiedForConditionalGeneration(config)
model = model.to(dtype=torch.float32)
model.eval()

print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")

model.save_pretrained(save_dir)

# Reuse the reference processor but shrink the soft-token budget so the tiny
# position embedding table (mm_posemb_size=64) is large enough.
processor = AutoProcessor.from_pretrained("google/gemma-4-12b-it")
processor.image_processor.max_soft_tokens = 70
processor.image_processor.image_seq_length = 70
processor.save_pretrained(save_dir)

print(f"Tiny Gemma4Unified model saved to {save_dir}")

# Sanity forward pass
input_ids = torch.randint(0, 262144, (1, 10))
with torch.no_grad():
    out = model(input_ids=input_ids)
print("logits shape:", out.logits.shape)
print("Forward pass OK!")