--- language: en license: apache-2.0 library_name: pytorch tags: - tiny-model - education - instruction-following - custom-architecture - pytorch - text-generation datasets: - HuggingFaceFW/fineweb-edu - teknium/OpenHermes-2.5 - m-a-p/CodeFeedback-Filtered-Instruction base_model: Arko007/tiny-edu-50m pipeline_tag: text-generation inference: false --- Arko007/Zenyx_114M-Tiny-Edu-Instruct 🎓 Note: This is an experimental research artifact. While the model achieved excellent training convergence (Validation Loss: ~1.04), inference on such a small scale (50M parameters) is highly sensitive to sampling parameters and may result in repetition loops. TinyEdu-50M-Instruct is a 50-million parameter language model designed to test the limits of instruction tuning on extremely small architectures. It is based on the TinyEdu-50M foundation model, which was pre-trained on the FineWeb-Edu dataset (100B+ tokens), and subsequently fine-tuned on a high-quality mix of instruction and code data. 📊 Model Details Architecture: Custom GPT-2 style Transformer with Rotary Positional Embeddings (RoPE) and Grouped Query Attention (GQA). Parameters: ~114 Million Context Length: 4096 tokens Vocabulary Size: 50,257 (GPT-2 Tokenizer) Precision: Float16 (Training), Float32 (Recommended for Inference) 📉 Training Data & Metrics Datasets used for SFT (Supervised Fine-Tuning): OpenHermes-2.5: 60% of the mix. High-quality general instruction following. CodeFeedback-Filtered: 40% of the mix. Code generation and logic. Training Dynamics Hardware: Single NVIDIA Tesla P100 (16GB VRAM) Optimization: DeepSpeed ZeRO-2, Gradient Checkpointing Batch Size: 1 (Effective Batch Size: 64 via Gradient Accumulation) Start Loss: ~11.0 Final Validation Loss: 1.0391 (Converged) ⚠️ Known Limitations Due to the extremely small parameter count (114M), this model exhibits specific behaviors: Repetition Loops: Without a high repetition_penalty (>= 1.2), the model tends to get stuck repeating tokens (e.g., )))) or ......). Prompt Sensitivity: It requires strict Alpaca formatting (### Instruction: ... ### Response:). Hallucination: As with all tiny models, factuality is low. It captures the structure of language and code well but may invent facts. 🚀 How to Use Since this uses a custom architecture implementation in raw PyTorch, you cannot use AutoModelForCausalLM directly. You must define the class structure first. 1. Install Dependencies pip install torch transformers huggingface_hub 2. Run Inference import torch import torch.nn as nn import torch.nn.functional as F from transformers import AutoTokenizer from huggingface_hub import hf_hub_download # --- 1. Define Model Architecture --- class Config: vocab_size = 50257 n_layer = 12 n_head = 12 n_embd = 768 n_kv_heads = 4 block_size = 4096 dropout = 0.0 bias = False rope_base = 10000 class RotaryEmbedding(nn.Module): def __init__(self, dim, max_position_embeddings=2048, base=10000): super().__init__() inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) self.max_seq_len_cached = max_position_embeddings t = torch.arange(self.max_seq_len_cached, dtype=self.inv_freq.dtype) freqs = torch.einsum("i,j->ij", t, self.inv_freq) emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False) self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False) def forward(self, x, seq_len=None): if seq_len > self.max_seq_len_cached: self.max_seq_len_cached = seq_len t = torch.arange(self.max_seq_len_cached, dtype=self.inv_freq.dtype, device=x.device) freqs = torch.einsum("i,j->ij", t, self.inv_freq) emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False) self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False) return self.cos_cached[:, :, :seq_len, ...], self.sin_cached[:, :, :seq_len, ...] def rotate_half(x): x1, x2 = x.chunk(2, dim=-1) return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb(q, k, cos, sin): cos = cos.to(q.device) sin = sin.to(q.device) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed class GroupedQueryAttention(nn.Module): def __init__(self, config): super().__init__() self.n_head = config.n_head self.n_kv_heads = config.n_kv_heads self.head_dim = config.n_embd // config.n_head self.n_rep = self.n_head // self.n_kv_heads self.q_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias) self.k_proj = nn.Linear(config.n_embd, self.n_kv_heads * self.head_dim, bias=config.bias) self.v_proj = nn.Linear(config.n_embd, self.n_kv_heads * self.head_dim, bias=config.bias) self.out_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias) self.rotary_emb = RotaryEmbedding(self.head_dim, max_position_embeddings=config.block_size, base=config.rope_base) def forward(self, x): B, T, C = x.size() q = self.q_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2) k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2) v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2) cos, sin = self.rotary_emb(v, seq_len=T) q, k = apply_rotary_pos_emb(q, k, cos, sin) k = k.repeat_interleave(self.n_rep, dim=1) v = v.repeat_interleave(self.n_rep, dim=1) y = F.scaled_dot_product_attention(q, k, v, is_causal=True) y = y.transpose(1, 2).contiguous().view(B, T, C) return self.out_proj(y) class MLP(nn.Module): def __init__(self, config): super().__init__() self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias) self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias) def forward(self, x): return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh')) class Block(nn.Module): def __init__(self, config): super().__init__() self.ln_1 = nn.LayerNorm(config.n_embd) self.attn = GroupedQueryAttention(config) self.ln_2 = nn.LayerNorm(config.n_embd) self.mlp = MLP(config) def forward(self, x): x = x + self.attn(self.ln_1(x)) return x + self.mlp(self.ln_2(x)) class TinyEduGPT(nn.Module): def __init__(self, config): super().__init__() self.config = config self.transformer = nn.ModuleDict(dict( wte = nn.Embedding(config.vocab_size, config.n_embd), drop = nn.Dropout(config.dropout), h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]), ln_f = nn.LayerNorm(config.n_embd), )) self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) self.transformer.wte.weight = self.lm_head.weight def forward(self, idx): x = self.transformer.wte(idx) for block in self.transformer.h: x = block(x) return self.lm_head(self.transformer.ln_f(x)) # --- 2. Load Weights --- config = Config() device = "cuda" if torch.cuda.is_available() else "cpu" model = TinyEduGPT(config).to(device) model_path = hf_hub_download("Arko007/tiny-edu-50m-instruct", "pytorch_model.bin") state_dict = torch.load(model_path, map_location=device) # Key Mapping Fix clean_sd = {} for k, v in state_dict.items(): new_k = k.replace("_orig_mod.", "").replace("module.", "") clean_sd[new_k] = v model.load_state_dict(clean_sd, strict=False) model.eval() # --- 3. Generate --- tokenizer = AutoTokenizer.from_pretrained("gpt2") prompt = "### Instruction:\nWrite a python function to add two numbers.\n\n### Response:\n" inputs = tokenizer.encode(prompt, return_tensors="pt").to(device) # Greedy decoding with penalty is recommended for this model size with torch.no_grad(): outputs = model(inputs) # Note: Full generation loop required (see generation scripts in discussions) 🛠️ Intended Uses Educational analysis of Transformer internals. Research into training dynamics of small language models. Low-latency inference testing on edge devices (requires conversion to ONNX/GGUF). ⚠️ Licensing This model is released under the Apache 2.0 License.