--- license: mit language: - en tags: - gemma4 - gemma4-text - gemma4-moe - moe - mixture-of-experts - causal-lm - tinystories - tiny-model - validation - debug-model - transformers - legacy-tokenizer pipeline_tag: text-generation --- # Tiny Gemma4 MoE Text 3M This repository contains an approximately 3M-parameter tiny Gemma4 text-only Mixture-of-Experts causal language model for validation and debugging. The model is intentionally small. It is not intended to be a high-quality text generation model. Its main purpose is to provide a compact checkpoint that exercises Gemma4 MoE text-model code paths in Hugging Face Transformers and in independent inference engines. This checkpoint is useful for implementation testing because it includes sliding attention layers, full attention layers, grouped-query attention, per-layer input embeddings, and MoE routing with multiple experts. Compared with the previous larger tiny checkpoint, this 3M version is intended to be the recommended compact validation checkpoint. ## Model purpose This model is designed for: * testing `Gemma4ForCausalLM` * validating `Gemma4TextConfig` * testing Gemma4 text-only MoE model loading * checking model save/load behavior * checking tokenizer save/load behavior * exercising sliding attention layers * exercising full attention layers * exercising grouped-query attention * exercising Gemma4 per-layer input embedding paths * exercising MoE expert parameters * exercising top-k expert routing * providing a compact Gemma4 MoE checkpoint for inference-engine validation It is not designed for: * high-quality story generation * instruction following * chat use * OCR * multimodal inference * benchmark comparison against production language models * production deployment ## Model architecture The model uses `Gemma4ForCausalLM` with a small Gemma4 text MoE configuration. Representative configuration: ```text model_type: gemma4_text vocab_size: 1024 vocab_size_per_layer_input: 1024 hidden_size: 128 hidden_size_per_layer_input: 16 intermediate_size: 384 intermediate_dim: 192 moe_intermediate_size: 192 num_hidden_layers: 6 num_attention_heads: 4 num_key_value_heads: 1 num_global_key_value_heads: 1 head_dim: 32 global_head_dim: 32 sliding_window: 128 max_position_embeddings: 1024 layer_types: - sliding_attention - sliding_attention - full_attention - sliding_attention - sliding_attention - full_attention hidden_activation: gelu_pytorch_tanh tie_word_embeddings: true attention_bias: false attention_dropout: 0.0 rms_norm_eps: 1e-06 initializer_range: 0.02 use_cache: true final_logit_softcapping: null use_bidirectional_attention: null attention_k_eq_v: false num_kv_shared_layers: 0 use_double_wide_mlp: false enable_moe_block: true num_experts: 4 top_k_experts: 2 expert_interval: 2 router_aux_loss_coef: 0.0 pad_token_id: 1000 bos_token_id: 1000 eos_token_id: 1001 ``` The attention pattern is: ```text ssFssF ``` where `s` means `sliding_attention` and `F` means `full_attention`. This pattern was chosen for validation coverage. A full-attention-only model may be easier to train, but it would not exercise the sliding attention path. ## MoE configuration This model enables Gemma4 MoE blocks. ```text enable_moe_block: true num_experts: 4 top_k_experts: 2 expert_interval: 2 moe_intermediate_size: 192 intermediate_dim: 192 ``` The `num_experts=4` and `top_k_experts=2` setting is intentional. A smaller configuration such as `num_experts=2, top_k=1` would exercise only a much simpler routing path. This checkpoint is intended to cover: ```text router / gate parameters multiple experts top-2 expert selection weighted expert combination MoE FFN parameters dense and MoE layer interaction ``` ## Tokenizer The model uses a small legacy-style byte-level BPE tokenizer. The tokenizer is trained with: ```text RawTokenizer(BPE()) ByteLevel(add_prefix_space=False) ByteLevelDecoder() BpeTrainer( vocab_size=1000, min_frequency=2, special_tokens=[], initial_alphabet=ByteLevel.alphabet(), ) ``` After BPE training, the following special tokens are added: ```text id 1000 id 1001 <|im_start|> id 1002 ``` The model config keeps `vocab_size=1024`, leaving a small reserved range above the actual tokenizer size. The pad token is set to ``, so `pad_token_id` and `bos_token_id` are both 1000. This tokenizer setup was chosen because the tiny model trained substantially better with this legacy byte-level configuration than with a standard `ByteLevelBPETokenizer` setup that trains special tokens directly into the vocabulary. ## Training data The model was trained on TinyStories-style English story text. The small vocabulary keeps the checkpoint compact, but it also limits text generation quality. The model often learns common TinyStories templates, especially stories about Lily, her mom, parks, apples, slides, toys, and simple moral situations. ## Training setup Representative training settings: ```text num_epochs: 1 learning_rate: 2e-4 batch_size: 32 block_size: 256 max_steps: derived from one epoch device: auto dtype: float32 by default grad_clip: 1.0 error_on_nonfinite_gradients: true weight_decay: 0.0 seed: 1234 vocab_size: 1024 base_vocab_size: 1000 legacy_tokenizer: true legacy_special_token_ids: true hidden_size: 128 intermediate_size: 384 moe_intermediate_size: 192 num_hidden_layers: 6 num_attention_heads: 4 num_key_value_heads: 1 head_dim: 32 hidden_size_per_layer_input: 16 layer_pattern: ssFssF sliding_window: 128 max_position_embeddings: 1024 enable_moe_block: true num_experts: 4 top_k_experts: 2 expert_interval: 2 router_aux_loss_coef: 0.0 ``` The final evaluation loss in the reference run was approximately: ```text Final loss: 1.5030 ``` This loss should not be interpreted as a benchmark against production models. The model is very small and includes Gemma4 MoE-specific architectural paths primarily for validation coverage. The result is useful because the model trains cleanly and generates coherent TinyStories-like text fragments while remaining compact. ## Example generation Example output from the reference checkpoint: ```text Prompt: Once upon Once upon a time, there was a little girl named Lily. She loved to play outside in the sunshine. One day, she saw a big, red apple on the ground. She picked it up and showed it to her mom. "Mommy, look what I found!" Lily said. "That's a big apple, Lily. It's a special apple. It's very special," her mom replied. ``` ```text Prompt: There was a little There was a little girl named Lily. She loved to play outside in the sunshine. One day, she saw a big, red apple on the ground. She picked it up and showed it to her mom. "Mommy, look at the apple!" Lily said. "That's a nice apple, Lily. It's very pretty," her mom replied. Lily was happy to have a new apple and wanted to ``` ```text Prompt: One day One day, a little girl named Lily went to the park with her mom. She saw a big slide and wanted to try it. But her mom said, "No, Lily. You have to wait. It's not safe." Lily was sad. She wanted to go on the slide. She asked her mom, "Can I go on the slide?" Her mom said, "No, Lily. You have to wait until the slide is safe." ``` The model can generate coherent TinyStories-like text fragments. Repetition, template convergence, weak long-form coherence, and repeated character names are expected and are not considered failures for the intended validation purpose. ## Usage If the model files are stored under an `hf/` subdirectory, use the following example. ```python import torch from transformers import PreTrainedTokenizerFast, Gemma4ForCausalLM repo = "shibatch/tinygemma4moe3m" tokenizer = PreTrainedTokenizerFast.from_pretrained(repo, subfolder="hf") model = Gemma4ForCausalLM.from_pretrained( repo, subfolder="hf", torch_dtype=torch.float32, ) model.eval() prompt = "Once upon" inputs = tokenizer(prompt, return_tensors="pt") with torch.no_grad(): output_ids = model.generate( **inputs, max_new_tokens=100, do_sample=False, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id, ) print(tokenizer.decode(output_ids[0], skip_special_tokens=True)) ``` ## Loading requirements This checkpoint requires a Transformers version that supports Gemma4 and Gemma4 MoE. The following imports should work: ```python from transformers import Gemma4ForCausalLM, Gemma4TextConfig ``` If these imports fail, update Transformers to a version with Gemma4 support. ## Tokenizer loading note This repository uses a custom byte-level BPE tokenizer saved as a `PreTrainedTokenizerFast`. For this reason, examples use: ```python from transformers import PreTrainedTokenizerFast ``` instead of `AutoTokenizer`. Using `AutoTokenizer` may fail in some environments if the tokenizer backend cannot be inferred automatically. The expected tokenizer files include: ```text tokenizer.json tokenizer_config.json special_tokens_map.json ``` ## Intended validation coverage This checkpoint is intended to validate support for: ```text Gemma4TextConfig Gemma4ForCausalLM sliding_attention layers full_attention layers GQA with num_key_value_heads = 1 global key/value head configuration per-layer input embeddings tied word embeddings Gemma4 RMSNorm behavior Gemma4 MLP activation: gelu_pytorch_tanh Gemma4 MoE expert parameters num_experts = 4 top_k_experts = 2 expert_interval = 2 MoE expert dispatch MoE expert output combination legacy byte-level BPE tokenizer loading generate() save_pretrained() from_pretrained() ``` ## Limitations This is a tiny debug model. It should not be used as a general-purpose language model. Known limitations: * TinyStories template convergence * repeated simple story patterns * weak long-form coherence * small vocabulary * weak semantic consistency * no instruction tuning * no chat formatting * no multimodal capability * no OCR capability * no production use The checkpoint is primarily intended to make Gemma4 MoE text-model code paths easy to test without downloading a large model. ## Why include MoE? A dense tiny Gemma4 model is simpler to train, but it does not cover MoE-specific implementation paths. This checkpoint intentionally includes: ```text num_experts = 4 top_k_experts = 2 ``` to exercise a more realistic MoE routing path than a minimal `top_k=1` configuration. ## Why not full attention only? A full-attention-only tiny model may train more cleanly, but it would not cover Gemma4 sliding attention behavior. This checkpoint uses: ```text sliding_attention sliding_attention full_attention sliding_attention sliding_attention full_attention ``` to cover both attention implementations. ## Why a legacy tokenizer? Earlier tiny Gemma4 MoE training attempts were sensitive to tokenizer details. The legacy byte-level BPE setup used here produced a substantially better tiny-model result than the earlier tokenizer setup. The tokenizer intentionally keeps byte-level behavior explicit: ```text ByteLevel(add_prefix_space=False) initial_alphabet=ByteLevel.alphabet() special tokens added after BPE training ``` This is useful for validation because the tokenizer can be saved and loaded through `PreTrainedTokenizerFast` without requiring SentencePiece or tiktoken inference. ## Notes on OCR and multimodal use This repository is text-only. It does not include a vision tower, image projector, image-token alignment, or OCR training. A Gemma4 OCR or Gemma4 MoE OCR validation model would be a separate project. It would require a tiny multimodal Gemma4 configuration, a synthetic OCR dataset, image-token handling, vision/text alignment, OCR fine-tuning, and additional validation scripts. ## Suggested repository name Suggested Hugging Face repository name: ```text shibatch/tinygemma4moe3m ``` ## Citation This is a synthetic tiny validation checkpoint derived from Gemma4-compatible MoE text architecture settings. It is intended for debugging and implementation testing.