Omyga-150M Base
A 149M-parameter decoder-only language model pretrained from scratch on ~3B tokens on a single GPU.
This is the base (pretrained) model: it continues text, it does not follow instructions or chat. For the chat versions, see Related models.
| Parameters | 149,341,696 |
| Training tokens | 2,999,975,936 (~3B) |
| Context length | 4,096 tokens |
| Final validation loss | 3.0848 (perplexity 21.86) |
| Throughput | ~68.4K tokens/sec (single GPU, BF16) |
Architecture
| Type | Decoder-only Transformer, pre-norm |
| Layers | 22 |
| Hidden size | 640 |
| Attention heads | 10 (head dim 64), 10 KV heads |
| MLP | SwiGLU, 640 → 1,920 → 640 |
| Normalization | RMSNorm (eps 1e-5), plus QK-norm on queries and keys |
| Positions | RoPE (theta 10,000) |
| Embeddings | Tied input/output embeddings (50,257 × 640) |
| Bias / dropout | None / 0.0 |
| Tokenizer | GPT-2 byte-level BPE, vocabulary 50,257 |
Training data
| Source | Tokens | Share |
|---|---|---|
| FineWeb-Edu | 1.35B | 45.0% |
| DCLM Baseline 1.0 | 1.00B | 33.3% |
| FineMath (4+) | 300M | 10.0% |
| English Wikipedia | 225M | 7.5% |
| OpenWebMath | 125M | 4.2% |
| Total | 3.0B | 100% |
Text was tokenized with the GPT-2 tokenizer into 60 training shards of 50M tokens each (plus a 3M-token validation shard), with deterministic document-level splits. Training samples are random 4,096-token windows.
Training setup
| Steps | 5,722 optimizer steps |
| Batch | 4 sequences × 32 gradient-accumulation steps × 4,096 tokens = 524,288 tokens per step |
| Learning rate | linear warmup to 6e-4 over 200 steps, then cosine decay |
| Precision | BF16 mixed precision |
| Peak GPU memory | 27.0 GB |
| Training compute | ~12.2 hours (3B tokens at ~68.4K tokens/sec) |
Validation loss during training:
| Step | 250 | 500 | 1,000 | 2,000 | 2,750 | 5,722 (final) |
|---|---|---|---|---|---|---|
| Val loss | 5.231 | 4.218 | 3.644 | 3.358 | 3.254 | 3.085 |
The full per-step log (loss, learning rate, gradient norm, throughput) is in training/metrics.jsonl.
How it works: data → training → inference
Evaluation
0-shot, full test/validation splits. Knowledge tasks are scored by the log-likelihood of each answer choice; acc_norm is the mean log-probability per token.
| Benchmark | Metric | Chance | Omyga-150M Base |
|---|---|---|---|
| SciQ | acc_norm | 25% | 65.30 |
| PIQA | acc_norm | 50% | 60.94 |
| ARC-Easy | acc_norm | 25% | 42.55 |
| OpenBookQA | acc_norm | 25% | 28.60 |
| HellaSwag | acc_norm | 25% | 28.26 |
| ARC-Challenge | acc_norm | 25% | 24.40 |
| WinoGrande | acc_norm | 50% | 51.30 |
| Arithmetic (GPT-3 set: 2/3-digit +, −, ×) | greedy exact match | – | 0.85 |
| GSM8K | 5-shot, greedy | – | 1.36 |
SciQ, PIQA and ARC-Easy are clearly above chance. HellaSwag, ARC-Challenge, WinoGrande and OpenBookQA are near chance, which is expected at this model size and token budget.
Usage
from transformers import AutoTokenizer, AutoModelForCausalLM
repo = "ChakradharS/Omyga-150m-base"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
inputs = tok("The water cycle is", return_tensors="pt")
out = model.generate(**inputs, max_new_tokens=60, do_sample=True, temperature=0.8, top_p=0.95,
pad_token_id=tok.eos_token_id)
print(tok.decode(out[0], skip_special_tokens=True))
This repository contains custom model code, so trust_remote_code=True is required. Review the code and pin a revision for reproducible use. This version generates without a KV cache (correct, but slower); the KV cache was added in the post-training models.
Limitations
- Base model only: it completes text and does not follow instructions.
- Small model trained on 3B tokens: it often produces incorrect or made-up facts and is weak at reasoning and arithmetic.
- English only. Not intended for production use.
Related models
| Model | What it is |
|---|---|
| Omyga-150m-base-chatml | This model + ChatML chat template and special tokens |
| Omyga-150m-sft | Supervised fine-tuned chat model |
| Omyga-150m-distill | SFT model + distillation from DeepSeek-V4-Flash on GSM8K |
Repository files
| Path | Contents |
|---|---|
model.safetensors, config.json, tokenizer files | Weights, configuration and tokenizer |
*.py (root) | Transformers-compatible model code (trust_remote_code) |
omyga_walkthrough.png | Diagram: data pipeline, batching, architecture, training and inference |
training/metrics.jsonl | Per-step training and validation log |
training/training_summary.json | Final training summary |
training/model_config.json, training/src/ | Original training-time model code and config |