Arko007/Zenyx_114M-Tiny-Edu-Instruct

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 (114M parameters) is highly sensitive to sampling parameters and may result in repetition loops. TinyEdu-50M-Instruct is a 114-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)

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.