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:

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

<s>          id 1000
</s>         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 <s>, 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:

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:

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:

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:

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.