--- license: mit language: - en tags: - llama4 - llama4-text - llama4-moe - moe - mixture-of-experts - causal-lm - tinystories - tiny-model - validation - debug-model - transformers --- # Tiny Llama4 MoE Text This repository contains a tiny Llama4 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 Llama4 MoE text-model code paths in Hugging Face Transformers. This checkpoint is useful for implementation testing because it includes chunked attention, full attention, grouped-query attention, QK norm, and MoE routing with multiple experts. ## Model purpose This model is designed for: * testing `Llama4ForCausalLM` * validating `Llama4TextConfig` * testing Llama4 text-only MoE model loading * checking model save/load behavior * checking tokenizer save/load behavior * exercising chunked attention layers * exercising full attention layers * exercising grouped-query attention * exercising QK norm paths * exercising MoE expert parameters * exercising top-k expert routing * providing a compact Llama4 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 `Llama4ForCausalLM` with a small Llama4 text MoE configuration. Representative configuration: ```text model_type: llama4_text vocab_size: 1024 hidden_size: 96 intermediate_size: 192 intermediate_size_mlp: 384 num_hidden_layers: 5 num_attention_heads: 4 num_key_value_heads: 1 head_dim: 24 max_position_embeddings: 1024 attention_chunk_size: 128 layer_types: - chunked_attention - full_attention - chunked_attention - full_attention - chunked_attention hidden_act: silu tie_word_embeddings: false attention_bias: false rms_norm_eps: 1e-05 rope_theta: 500000.0 use_qk_norm: true attn_temperature_tuning: false floor_scale: 8192 attn_scale: 0.1 num_local_experts: 4 num_experts_per_tok: 2 moe_layers: - 0 - 1 - 2 - 3 - 4 interleave_moe_layer_step: 1 router_aux_loss_coef: 0.001 router_jitter_noise: 0.0 output_router_logits: false pad_token_id: 2 bos_token_id: 0 eos_token_id: 1 ``` The attention pattern is: ```text CFCFC ``` where `C` means `chunked_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 chunked attention path. ## MoE configuration This model enables Llama4 MoE blocks. ```text num_local_experts: 4 num_experts_per_tok: 2 moe_layers: all layers interleave_moe_layer_step: 1 router_aux_loss_coef: 0.001 ``` The `num_local_experts=4` and `num_experts_per_tok=2` setting is intentional. A smaller configuration such as `num_local_experts=2, num_experts_per_tok=1` would exercise only a much simpler routing path. This checkpoint is intended to cover: ```text router parameters multiple experts top-2 expert selection weighted expert combination MoE FFN parameters chunked and full attention interaction MoE layer execution in a small model ``` ## Parameter count The exact parameter count depends on the saved checkpoint configuration. The default training script prints the parameter count at startup. For the default configuration, check the run log for: ```text Parameter count: ... MoE/router/expert parameter names found: ... Prefix breakdown: ... ``` After training, the metadata file also records the count: ```text artifact_metadata.json ``` ## Training data The model was trained on TinyStories-style English story text. The tokenizer is a small byte-level BPE tokenizer with a vocabulary size of 1024. The small vocabulary keeps the checkpoint compact, but it also limits text generation quality. ## Training setup Representative training settings: ```text num_epochs: 1 learning_rate: 3e-4 batch_size: 32 block_size: 256 vocab_size: 1024 hidden_size: 96 intermediate_size: 192 intermediate_size_mlp: 384 num_hidden_layers: 5 num_attention_heads: 4 num_key_value_heads: 1 head_dim: 24 layer_pattern: CFCFC attention_chunk_size: 128 num_local_experts: 4 num_experts_per_tok: 2 moe_layers: all interleave_moe_layer_step: 1 router_aux_loss_coef: 0.001 ``` The reference run can be reproduced with: ```bash python train_llama4_moe_text_tinystories_epoch.py \ --output-dir tinyllama4moe2m \ --num-epochs 1 \ --max-steps 0 ``` A short smoke test can be run with: ```bash python train_llama4_moe_text_tinystories_epoch.py \ --output-dir tinyllama4moe_smoke \ --max-rows 2000 \ --max-steps 100 \ --log-steps 10 \ --eval-steps 50 ``` ## Reference result The reference run completed one epoch of TinyStories-style training. ```text Final loss: 2.458171248435974 ``` Example generations from the reference checkpoint: ```text Prompt: Once upon Once upon a time, there was a little girl named Lily. She loved to play with her toys and her friends. One day, she saw a big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, big, ``` ```text Prompt: There was a little There was a little girl named Lily. She loved to play with her toys and her friends. One day, she saw a big, big, big carrot in the park. She wanted to play with it, but she was very sad. Lily saw a big dog who was playing with a ballrot. She wanted to play with it, but she was too big to play with. She asked her mom to help her. She said, "I can't go to the park, but I ``` ```text Prompt: One day One day, a little girl named Lily went to the park with her mom. They saw a big tree with a big bowl of cookies. Lily wanted to play with the big, but her mom said they could go to the park. Lily was very happy. She saw a big tree with a big bowl of cookies. She wanted to play with it. She asked her mom, "Why are you sad?" Her mom said, "I'm sorry ``` The model can generate TinyStories-like fragments, but repetition, template collapse, and occasional invented words are expected. This is normal for this checkpoint and is not considered a failure for its 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, Llama4ForCausalLM repo = "shibatch/tinyllama4moe2m" tokenizer = PreTrainedTokenizerFast.from_pretrained(repo, subfolder="hf") model = Llama4ForCausalLM.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 Llama4 and Llama4 MoE. The following imports should work: ```python from transformers import Llama4ForCausalLM, Llama4TextConfig ``` If these imports fail, update Transformers to a version with Llama4 support. ## Tokenizer 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 Llama4TextConfig Llama4ForCausalLM chunked_attention layers full_attention layers GQA with num_key_value_heads = 1 QK norm Llama4 RMSNorm behavior Llama4 MLP activation: silu Llama4 MoE expert parameters num_local_experts = 4 num_experts_per_tok = 2 moe_layers expert routing expert output combination generate() save_pretrained() from_pretraine() ``` ## Limitations This is a tiny debug model. It should not be used as a general-purpose language model. Known limitations: * frequent phrase repetition may occur * TinyStories template collapse may occur * 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 Llama4 MoE text-model code paths easy to test without downloading a large model. ## Why include MoE? A dense tiny Llama4 model is simpler to train, but it does not cover MoE-specific implementation paths. This checkpoint intentionally includes: ```text num_local_experts = 4 num_experts_per_tok = 2 ``` to exercise a more realistic MoE routing path than a minimal top-1 routing configuration. ## Why include chunked attention? A full-attention-only tiny model may train more cleanly, but it would not cover Llama4 chunked attention behavior. This checkpoint uses: ```text chunked_attention full_attention chunked_attention full_attention chunked_attention ``` to cover both attention implementations. ## 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 Llama4 OCR or Llama4 multimodal MoE validation model would be a separate project. It would require a tiny multimodal Llama4 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/tinyllama4moe2m ``` Alternative names: ```text shibatch/tinyllama4moetext2m shibatch/tinyllama4-moe-text ``` ## Citation This is a synthetic tiny validation checkpoint derived from Llama4-compatible MoE text architecture settings. It is intended for debugging and implementation testing.