Welmia 1.0 (81M) — Instruction-Tuned Small Language Model

Welmia 1.0 81M is a lightweight, 81-million-parameter causal language model, trained completely from scratch — original architecture implementation, from-random-initialization pretraining, no borrowed backbone — using a TinyLLaMA-inspired design: RoPE positional encoding, RMSNorm, SwiGLU feed-forward layers, and full KV-cache streaming for fast token-by-token generation. It's instruction-tuned to follow Alpaca-style prompts, and small enough to run comfortably on CPU.

🌐 Official Website

welmia.kesug.com · GitHub

🆕 Fixed: Transformers loading now works correctly

Earlier versions of this repository had a bug in modeling_welmia.py: the RoPE (rope_cos/rope_sin) position-encoding buffers were registered as non-persistent. On recent versions of transformers, which use a meta-device fast-loading path by default, this caused those buffers to load as uninitialized garbage instead of real values — producing NaN output when loaded via AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True). The repository's KV-cache plumbing was also incompletely wired through the causal-LM wrapper.

Both are now fixed: RoPE buffers are baked directly into model.safetensors, attention uses a manual (non-SDPA) implementation for numerical stability across backends, and KV-cache is correctly threaded through generate(). If you downloaded this model before this fix, please re-download model.safetensors and modeling_welmia.py.

🏗️ Architecture Details

🚀 Quick Start

Option 1: Standalone (Recommended — simplest, no Transformers needed)

Runs with just 3 dependencies — no transformers, no custom architecture registration, no version conflicts.

Install:

pip install torch safetensors tiktoken huggingface-hub

Run:

# inference.py — Welmia-1.0 standalone inference with KV-cache streaming
import os, torch, torch.nn as nn, torch.nn.functional as F
from dataclasses import dataclass
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download
import tiktoken

torch.set_num_threads(min(os.cpu_count(), 4))  # see note below on thread count
device = "cuda" if torch.cuda.is_available() else "cpu"

# ── Model Architecture ──────────────────────────────────────────
@dataclass
class Config:
    vocab_size: int = 50257
    block_size: int = 1024
    n_layer: int = 6
    n_head: int = 12
    n_embd: int = 768
    dropout: float = 0.0

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
    def forward(self, x):
        norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return norm * self.weight

def build_rope_cache(seq_len, dim, base=10000, device="cpu"):
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=device).float() / dim))
    t = torch.arange(seq_len, device=device).float()
    freqs = torch.outer(t, inv_freq)
    return torch.cos(freqs), torch.sin(freqs)

def apply_rope(x, cos, sin):
    x1, x2 = x[..., ::2], x[..., 1::2]
    cos = cos[None, None, :x.size(2), :]
    sin = sin[None, None, :x.size(2), :]
    rx1 = x1 * cos - x2 * sin
    rx2 = x1 * sin + x2 * cos
    return torch.stack([rx1, rx2], dim=-1).flatten(-2)

class Attention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.n_head = cfg.n_head
        self.head_dim = cfg.n_embd // cfg.n_head
        self.qkv = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=False)
        self.proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=False)

    def forward(self, x, cos_full, sin_full, start_pos, cache, use_cache):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        cos = cos_full[start_pos:start_pos + T]
        sin = sin_full[start_pos:start_pos + T]
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)

        if cache is not None:
            past_k, past_v = cache
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)
        new_cache = (k, v) if use_cache else None

        is_causal = cache is None
        y = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal, dropout_p=0.0)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(y), new_cache

class SwiGLU(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        hidden = int(8 * cfg.n_embd / 3)
        hidden = ((hidden + 63) // 64) * 64
        self.w1 = nn.Linear(cfg.n_embd, hidden, bias=False)
        self.w2 = nn.Linear(hidden, cfg.n_embd, bias=False)
        self.w3 = nn.Linear(cfg.n_embd, hidden, bias=False)
    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.norm1 = RMSNorm(cfg.n_embd)
        self.attn = Attention(cfg)
        self.norm2 = RMSNorm(cfg.n_embd)
        self.mlp = SwiGLU(cfg)
    def forward(self, x, cos, sin, start_pos, cache, use_cache):
        attn_out, new_cache = self.attn(self.norm1(x), cos, sin, start_pos, cache, use_cache)
        x = x + attn_out
        x = x + self.mlp(self.norm2(x))
        return x, new_cache

class GPT(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.wte = nn.Embedding(cfg.vocab_size, cfg.n_embd)
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)])
        self.norm_f = RMSNorm(cfg.n_embd)
        self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
        self.lm_head.weight = self.wte.weight
        head_dim = cfg.n_embd // cfg.n_head
        cos, sin = build_rope_cache(cfg.block_size, head_dim)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)

    def forward(self, idx, kv_caches=None, start_pos=0, use_cache=False):
        x = self.wte(idx)
        new_caches = [] if use_cache else None
        for i, blk in enumerate(self.blocks):
            cache_i = kv_caches[i] if kv_caches is not None else None
            x, cache_out = blk(x, self.rope_cos, self.rope_sin, start_pos, cache_i, use_cache)
            if use_cache:
                new_caches.append(cache_out)
        x = self.norm_f(x)
        logits = self.lm_head(x)
        if use_cache:
            return logits, new_caches
        return logits

# ── Download & Load ──────────────────────────────────────────────
print("⬇️  Downloading model from Hugging Face...")
cfg = Config(vocab_size=50257, block_size=1024, n_layer=6, n_head=12, n_embd=768)
model = GPT(cfg).to(device)

safetensors_path = hf_hub_download(repo_id="Welmia/welmia-1.0-81m", filename="model.safetensors")
state_dict = load_file(safetensors_path, device=str(device))

# This class computes rope_cos/rope_sin itself in __init__ (always correct,
# no meta-device path here) — drop the checkpoint's copies rather than load them.
state_dict.pop("rope_cos", None)
state_dict.pop("rope_sin", None)
# Restore weight tying (lm_head.weight is not stored separately, to save space)
state_dict["lm_head.weight"] = state_dict["wte.weight"]

model.load_state_dict(state_dict, strict=True)
model.eval()

params = sum(p.numel() for p in model.parameters()) / 1e6
print(f"✅ Loaded ({params:.1f}M params) on {device}")

enc = tiktoken.get_encoding("gpt2")

# ── Streaming Generation with KV-Cache ───────────────────────────
@torch.inference_mode()
def stream_generate(prompt, max_new_tokens=300, temperature=0.7, top_k=50):
    temperature = max(temperature, 1e-6)
    formatted = f"### Instruction:\n{prompt}\n\n### Response:\n"
    ids = enc.encode_ordinary(formatted)
    if len(ids) > cfg.block_size:
        ids = ids[-cfg.block_size:]
    x = torch.tensor([ids], dtype=torch.long, device=device)
    pos = x.size(1)

    print("💡 ", end="", flush=True)
    generated_tokens = []

    logits, caches = model(x, kv_caches=None, start_pos=0, use_cache=True)

    for _ in range(max_new_tokens):
        if pos >= cfg.block_size:
            break
        next_logits = logits[:, -1, :] / temperature
        if top_k > 0:
            v, _ = torch.topk(next_logits, min(top_k, next_logits.size(-1)))
            next_logits[next_logits < v[:, [-1]]] = float("-inf")
        probs = torch.softmax(next_logits, dim=-1)
        next_token = torch.multinomial(probs, 1)

        if next_token.item() == enc.eot_token:
            break

        generated_tokens.append(next_token.item())
        token_text = enc.decode(generated_tokens)
        prev_text = enc.decode(generated_tokens[:-1]) if len(generated_tokens) > 1 else ""
        print(token_text[len(prev_text):], end="", flush=True)

        logits, caches = model(next_token, kv_caches=caches, start_pos=pos, use_cache=True)
        pos += 1

    print("\n", flush=True)

# ── Interactive Loop ──────────────────────────────────────────────
print("\n🎯 TOKEN-BY-TOKEN STREAMING (type 'quit' to exit)\n")
while True:
    user_input = input("📝 You: ").strip()
    if user_input.lower() in ("quit", "exit", "q"):
        print("Bye! 👋")
        break
    if not user_input:
        continue
    stream_generate(user_input)
    print()

Thread count note: on machines with many CPU cores (e.g. Colab/Kaggle), too many threads can cause small numerical differences in floating-point reduction that noticeably affect a model this small. torch.set_num_threads(4) is a safe default; try 2–8 if you want to tune it.

Option 2: Hugging Face Transformers (now fully working, with generate() + KV-cache)

Requires transformers >= 4.40.0 for the custom welmia architecture. Confirmed working on modern transformers releases (5.x) after the fix above.

Install:

pip install torch safetensors "transformers>=4.40.0"

Run:

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "Welmia/welmia-1.0-81m",
    trust_remote_code=True,
    torch_dtype="auto",
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("Welmia/welmia-1.0-81m")

prompt = "### Instruction:\nWhat is machine learning?\n\n### Response:\n"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, use_cache=True)
print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))

trust_remote_code=True is required — custom architecture defined in modeling_welmia.py in this repository. use_cache=True now works correctly and gives real streaming-style speedups over recomputing the full sequence each step.

⚠️ Important Notes

• Custom architecture (model_type: welmia) — trust_remote_code=True is required with Transformers.
• Tokenizer: standard GPT-2 BPE (tiktoken-compatible, vocab size 50,257).
• Context limit: 1024 tokens. Longer inputs are truncated.
• Weight tying: lm_head is tied to the token embeddings and not stored separately in model.safetensors. If loading the raw file manually (Option 1), restore it with state_dict["lm_head.weight"] = state_dict["wte.weight"].
• Instruction format: trained on Alpaca-style prompts — always wrap your input as ### Instruction:\n{text}\n\n### Response:\n for coherent output. This is not a raw text-continuation base model.
• Thread count: on many-core CPUs, set torch.set_num_threads(4) to avoid small floating-point inconsistencies that are more noticeable on small models.

📊 Intended Use & Limitations

✅ Good for:

• Edge/local deployment on CPU or low-VRAM GPUs
• Research into small-scale LM training dynamics and KV-cache streaming
• Fast prototyping and instruction-following experiments
• Educational purposes and architecture exploration (RoPE, RMSNorm, SwiGLU from scratch)

❌ Not suitable for:

• Production applications requiring high accuracy
• Tasks requiring deep world knowledge or multi-step reasoning
• Multilingual generation (English only)
• Long-context tasks beyond 1024 tokens
• Safety-critical or unmoderated deployments (no RLHF/safety alignment beyond Alpaca SFT)

🔗 Related Models

• Welmia/welmia-2.0-103m-base — newer 103M-parameter Welmia base model, custom 24K tokenizer, 512-token context
• Welmia/welmia-2.0-103m-instruct — instruction-tuned version of the above

Note: Welmia 2.0 uses a different tokenizer (custom 24K vocab vs. GPT-2's 50,257) and is not weight- or tokenizer-compatible with this 81M model.

📄 License

Apache License 2.0

👤 Author

Trained and released by Welmia 🌐 welmia.kesug.com · 💻 github.com/ShellCrafter

🙏 Acknowledgements

Architecture inspired by TinyLLaMA. Built with PyTorch, safetensors, and tiktoken.

If you find this model useful, please consider leaving a ❤️ like and sharing your experiments!