Gowtham2962S/146M_Model

MODEL_146M

A 146M-parameter GPT-style causal language model, trained from scratch on 10B tokens of FineWeb-Edu and instruction-tuned on Magpie-Pro-300K-Filtered.

Model Details

• Architecture: Decoder-only transformer (LLaMA-style block)
  • 17 transformer blocks
  • 12 attention heads, 4 KV heads (grouped-query attention)
  • 768 hidden size
  • RoPE positional embeddings
  • RMSNorm
  • SwiGLU MLP (gated, hidden dim = 8/3 × embedding dim)
  • Weight-tied input/output embeddings
• Context length: 2048 tokens
• Tokenizer: GPT-2 BPE (tiktoken, vocab size 50,257)
• Parameters: ~146M
• Optimizer: Muon (2D body weights) + AdamW (embeddings, norms, LM head), cosine LR schedule with warmup
• Precision: bf16/fp16 mixed precision (autocast), fp32 master weights
• Framework: PyTorch (custom training loop, no external trainer)

Training Data

• Pretraining: 10B tokens sampled from FineWeb-Edu, a filtered, education-quality subset of FineWeb/CommonCrawl.
• Supervised fine-tuning (SFT): Magpie-Align/Magpie-Pro-300K-Filtered, a filtered set of synthetically generated instruction-response pairs.

Intended Use

This is a small-scale research/educational language model, useful for:

• Studying LM pretraining and instruction-tuning at a compact, single-GPU-friendly scale
• Experimenting with the Muon optimizer
• Baseline comparisons for architecture/data ablations

It is not intended for production deployment, factual question-answering, or any use case requiring reliability, safety filtering, or up-to-date knowledge.

How to Use

# Requires the model's GPT class + safetensors checkpoint loading utilities
# (see the project's Train.py / infer.py for the reference implementation)

import torch, tiktoken
from Train import GPT, GPTConfig, clean_state_dict_keys, load_checkpoint

config, state_dict, info = load_model_checkpoint_for_inference("sft_model_056250.safetensors", "cpu")
model = GPT(config)
model.load_state_dict(clean_state_dict_keys(state_dict))
model.eval()

enc = tiktoken.get_encoding("gpt2")
prompt = "General relativity is"
ids = torch.tensor([enc.encode_ordinary(prompt)])
with torch.inference_mode():
    logits, _, _ = model(ids)