Welmia/welmia-1.0-81m

Welmia-1.0 (81M)

A lightweight 81M-parameter causal language model trained completely from scratch using a TinyLLaMA-inspired architecture with RoPE, RMSNorm, SwiGLU, and KV-cache streaming support. Designed for fast local inference, edge deployment, and research experimentation.

🏗️ Architecture Details

Parameter	Value
Parameters	81M
Layers	6
Attention Heads	12
Embedding Dim	768
Context Length	1024
Vocab Size	50,257
Normalization	RMSNorm
Activation	SwiGLU
Positional Enc.	RoPE
Tokenizer	GPT-2 (tiktoken)
Weight Tying	✅ Yes

🚀 Quick Start

Installation

pip install torch safetensors tiktoken transformers>=4.40.0

### Load with Transformers (Recommended)
```python
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "YOUR_HF_USERNAME/welmia-1.0-81m",
    trust_remote_code=True,
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("YOUR_HF_USERNAME/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, temperature=0.7, top_k=50)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Streaming Inference with KV Cache

This model supports token-by-token streaming generation with persistent KV cache for efficient long-context inference:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "YOUR_HF_USERNAME/welmia-1.0-81m", trust_remote_code=True
).eval()
tokenizer = AutoTokenizer.from_pretrained("YOUR_HF_USERNAME/welmia-1.0-81m")

@torch.inference_mode()
def stream_generate(prompt, max_new_tokens=300, temperature=0.7, top_k=50):
    formatted = f"### Instruction:\n{prompt}\n\n### Response:\n"
    ids = tokenizer.encode(formatted)
    x = torch.tensor([ids], dtype=torch.long, device=model.device)
    
    logits, caches = model(x, use_cache=True)
    generated = []
    
    for _ in range(max_new_tokens):
        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() == tokenizer.eos_token_id:
            break
            
        generated.append(next_token.item())
        print(tokenizer.decode(generated[-1]), end="", flush=True)
        logits, caches = model(next_token, kv_caches=caches, 
                               start_pos=len(ids)+len(generated)-1, use_cache=True)
    print()

stream_generate("Explain why the sky is blue")

⚠️ Important Notes

• Custom Architecture: This model uses a non-standard architecture. You must pass trust_remote_code=True when loading.
• Tokenizer: Uses the standard GPT-2 BPE tokenizer via tiktoken. The included HF tokenizer files are compatible wrappers.
• Training: Trained completely from scratch (not fine-tuned from an existing checkpoint). Training details and dataset information will be added in future updates.
• Context Limit: Maximum context length is 1024 tokens. Inputs exceeding this will be truncated from the left.

📊 Intended Use & Limitations

✅ Good for:

• Edge/local deployment on CPU or low-VRAM GPUs
• Research into small-scale LM training dynamics
• Fast prototyping and instruction-following experiments
• Educational purposes and architecture exploration

❌ Not suitable for:

• Production applications requiring high accuracy
• Tasks requiring deep world knowledge or reasoning
• Multilingual generation (English only)
• Long-context tasks beyond 1024 tokens

📄 License

Apache License 2.0

👤 Author

Trained and released by [Your Name/Org]

🙏 Acknowledgements

Architecture inspired by TinyLLaMA. Built with PyTorch, safetensors, and Hugging Face Transformers.

---

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