Nebulixlabs/Nutral-GPT2-16M

Tiny But Mighty😉

🚀 Nutral-GPT2-16M

Nutral-GPT2-16M is a custom, highly-efficient 16 Million parameter conversational AI model built from scratch by Nebulixlabs. Built upon the raw and elegant nanoGPT architecture, this model represents a complete end-to-end AI lifecycle—from data tokenization and pre-training to instruct fine-tuning and Safetensors deployment.

🧠 Model Details & Architecture

Despite being a "nano" model, it features a complete transformer backbone:

• Architecture: Custom nanoGPT (Decoder-only Transformer)
• Parameter Count: ~16 Million
• Layers (n_layer): 4
• Attention Heads (n_head): 4
• Embedding Size (n_embd): 256
• Context Length (block_size): 256
• Vocabulary Size: 50,304 (Standard GPT-2 tiktoken Encoding)
• Format: Safetensors (Optimized for fast loading)

💻 Supported Devices & Hardware Requirements

Because of its ultra-lightweight size (~65 MB in RAM), Nutral-GPT2-16M can run literally anywhere. It does not require massive cloud GPUs!

• CPU: Can run flawlessly on any standard laptop or PC CPU (Windows/Mac/Linux).
• Mobile/Edge: Small enough to be deployed on Android/iOS devices, Raspberry Pi, or microcontrollers.
• Web Browser: Can be converted to WebAssembly (WASM) or WebGPU to run locally inside Google Chrome.
• GPU: Runs blisteringly fast on NVIDIA GPUs (T4, RTX series) or Apple Silicon (M1/M2/M3) using PyTorch MPS.

📊 Training Datasets & Pipeline

The model was trained on Kaggle using Distributed Data Parallel (DDP) across 2x NVIDIA T4 GPUs using mixed-precision (float16).

1. Pre-Training (Knowledge & Grammar):

   • Dataset: HuggingFaceFW/fineweb-edu
   • Tokens: 200 Million Tokens
   • Goal: Teaching the AI fundamental English grammar, reasoning, and basic worldly facts.

2. Fine-Tuning (Conversational Alignment):

   • Dataset: HuggingFaceH4/ultrachat_200k
   • Tokens: 60 Million Tokens
   • Goal: Teaching the AI how to behave like an assistant and format its replies to user prompts.

3. Training Performance:

   • Final Fine-Tuning Loss: 3.8892 🏆 (A highly competitive and healthy loss for a 16M parameter architecture).

---

🛠️ How to Use (Inference Code)

Since this model is built on a pure nanoGPT framework, you need the model.py file to initialize the neural network structure. Running this script will automatically download the necessary files from Hugging Face and start a chat session!

Prerequisites

pip install torch transformers safetensors huggingface_hub

Python Chat Script

import json
import torch
from transformers import AutoTokenizer
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

# Repository details
REPO_ID = "Nebulixlabs/Nutral-GPT2-16M"

print("Downloading necessary files from Hugging Face...")
# This automatically downloads the architecture script, config, and weights
hf_hub_download(repo_id=REPO_ID, filename="model.py", local_dir=".")
config_path = hf_hub_download(repo_id=REPO_ID, filename="model_config.json")
model_path = hf_hub_download(repo_id=REPO_ID, filename="model.safetensors")

# Import the architecture (now that model.py is present in the directory)
from model import GPTConfig, GPT

# 1. Load the GPT-2 Tokenizer
tokenizer = AutoTokenizer.from_pretrained(REPO_ID)

# 2. Load Config and Initialize Model
with open(config_path, "r") as f:
    config_args = json.load(f)

gptconf = GPTConfig(**config_args)
model = GPT(gptconf)

# 3. Load Model Weights
state_dict = load_file(model_path)
model.load_state_dict(state_dict)

# Hardware setup (CUDA, MPS for Mac, or CPU)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
if torch.backends.mps.is_available(): device = 'mps'
model.to(device)
model.eval()

print(f"✅ Model Loaded Successfully on {device.upper()}! Ready to chat.\n")

# 4. Chat Function
def chat(user_message):
    prompt = f"user: {user_message}\nassistant:"
    input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
    
    with torch.no_grad():
        generated_ids = model.generate(
            input_ids, 
            max_new_tokens=80, 
            temperature=0.7, 
            top_k=40
        )
    
    output_text = tokenizer.decode(generated_ids[0].tolist(), skip_special_tokens=True)
    
    # Extract only the AI's response
    if "assistant:" in output_text:
        return output_text.split("assistant:")[-1].strip()
    return output_text

# Test the AI!
question = "What is artificial intelligence?"
print(f"👤 You: {question}")
print(f"🤖 AI: {chat(question)}")

⚠️ Limitations

As a 16M parameter model, Nutral-GPT2-16M is primarily an educational research project. While it successfully understands chat formats and can construct coherent English sentences, it cannot compete with multi-billion parameter LLMs in deep factual recall, complex mathematics, or coding tasks. It might occasionally hallucinate or generate repetitive text.

🙏 Acknowledgements

A massive thanks to Andrej Karpathy and the open-source nanoGPT project. Without this brilliantly simple and educational PyTorch repository, training this model from scratch wouldn't have been possible. All core framework credits go to the nanoGPT codebase.