--- license: mit language: - en tags: - custom-architecture - pytorch - scratch-model - nanocoder --- # 🚀 Nebulixlabs/Nanocoder-Base **Nanocoder-Base** is a custom-built, ultra-lightweight, autoregressive language model trained from scratch. With approximately **19.5 Million parameters**, it is designed to be highly efficient, experimental, and capable of running on severely resource-constrained hardware (including edge devices and single standard GPUs). It was built specifically to understand basic English language structure and the foundational syntax of programming languages like Python and JavaScript. ## 📊 Model Details * **Developer:** Nebulixlabs * **Model Type:** Custom Autoregressive Decoder-Only Transformer * **Parameter Count:** 19,231,488 (~19.5M) * **Language(s):** English, Python, JavaScript * **License:** MIT ### Architecture Specifications | Component | Specification | | :--- | :--- | | **Layers (Transformer Blocks)** | 8 | | **Hidden Dimension (d_model)** | 256 | | **Attention Heads** | 8 (32 dimensions per head) | | **Context Window (MAX_SEQ_LEN)** | 256 tokens | | **Vocabulary Size** | 50,257 (Standard GPT-2 Tokenizer) | --- ## ⚙️ How It Works (Under the Hood) Nanocoder is not a standard Hugging Face `transformers` class; it is a raw, custom PyTorch implementation optimized for speed and memory efficiency. 1. **Flash Attention Integration:** Instead of standard multi-head attention math, Nanocoder uses PyTorch 2.0's native `F.scaled_dot_product_attention`. This drastically reduces VRAM usage and speeds up both training and inference. 2. **Weight Tying:** The embedding layer (`token_emb`) and the final output layer (`lm_head`) share the same weights. This is a crucial technique that saves millions of parameters while allowing the model to learn token representations more effectively. 3. **Pre-Layer Normalization:** To maintain gradient stability during training, LayerNorm is applied *before* the self-attention and feed-forward networks, rather than after. 4. **Compute-Optimal Scaling:** The model was trained using a 15x token-to-parameter ratio (~292.5 Million tokens), ensuring it extracts the maximum possible knowledge without overfitting its small parameter budget. --- ## 🎯 Capabilities & Limitations **What Nanocoder is good at:** * **Syntax Recognition:** It understands the basic visual structure of code (e.g., Python indentation, function definitions `def ... :`, and basic loops). * **Pattern Completion:** Generating short sequences of text or continuing a simple coding prompt. * **Educational Prototyping:** It is an excellent foundational model for students and researchers who want to learn how LLMs work, how to write custom PyTorch architectures, and how to execute fine-tuning pipelines locally without massive GPU clusters. **What Nanocoder is NOT good at:** * Because it only has 19.5M parameters (compared to billions in Llama or GPT), it has a strict "Capacity Wall." * It cannot execute complex mathematical logic, remember long conversational contexts, or write production-ready software. * It will hallucinate if asked complex reasoning questions. --- ## 📚 Recommended Fine-Tuning Data To make Nanocoder highly effective for your specific use case, you must fine-tune it on **high-quality, narrowly focused datasets**. Do not feed it broad knowledge; feed it specific formats. * **For a Chatbot:** Use datasets like `OpenAssistant/oasst_top1_2023-08-25`. This will teach the model the `<|im_start|>user` and `<|im_start|>assistant` conversational tags. * **For a Coding Assistant:** Use `sahil2801/CodeAlpaca-20k`. This teaches the model to read an `Instruction:` and generate the corresponding `Output:` code. * **Format is Everything:** Ensure your fine-tuning data strictly follows a uniform template. Small models learn formats much faster than they learn raw facts. --- ## 💻 Demo: How to Load and Fine-Tune Nanocoder Because Nanocoder uses a custom architecture, you cannot load it using `AutoModelForCausalLM.from_pretrained()`. You must define the architecture in your script and load the state dictionary. Here is a complete, self-contained PyTorch script to load the model and start a fine-tuning loop: ```python import torch import torch.nn as nn import torch.nn.functional as F # ========================================== # 1. DEFINE THE EXACT ARCHITECTURE # ========================================== VOCAB_SIZE = 50257 MAX_SEQ_LEN = 256 EMBED_DIM = 256 NUM_LAYERS = 8 NUM_HEADS = 8 class SelfAttention(nn.Module): def __init__(self): super().__init__() self.c_attn = nn.Linear(EMBED_DIM, 3 * EMBED_DIM, bias=False) self.c_proj = nn.Linear(EMBED_DIM, EMBED_DIM, bias=False) self.n_head = NUM_HEADS self.head_dim = EMBED_DIM // NUM_HEADS self.dropout = nn.Dropout(0.1) def forward(self, x): B, T, C = x.size() qkv = self.c_attn(x) q, k, v = qkv.split(EMBED_DIM, 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) y = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=0.1 if self.training else 0) return self.dropout(self.c_proj(y.transpose(1, 2).contiguous().view(B, T, C))) class TransformerBlock(nn.Module): def __init__(self): super().__init__() self.ln_1 = nn.LayerNorm(EMBED_DIM) self.attn = SelfAttention() self.ln_2 = nn.LayerNorm(EMBED_DIM) self.mlp = nn.Sequential( nn.Linear(EMBED_DIM, 4 * EMBED_DIM, bias=False), nn.GELU(), nn.Linear(4 * EMBED_DIM, EMBED_DIM, bias=False), nn.Dropout(0.1), ) def forward(self, x): x = x + self.attn(self.ln_1(x)) x = x + self.mlp(self.ln_2(x)) return x class NanoCoder(nn.Module): def __init__(self): super().__init__() self.token_emb = nn.Embedding(VOCAB_SIZE, EMBED_DIM) self.pos_emb = nn.Embedding(MAX_SEQ_LEN, EMBED_DIM) self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)]) self.ln_f = nn.LayerNorm(EMBED_DIM) self.lm_head = nn.Linear(EMBED_DIM, VOCAB_SIZE, bias=False) self.token_emb.weight = self.lm_head.weight # Weight Tying def forward(self, idx, targets=None): B, T = idx.size() pos = torch.arange(0, T, dtype=torch.long, device=idx.device) x = self.token_emb(idx) + self.pos_emb(pos) for block in self.blocks: x = block(x) x = self.ln_f(x) logits = self.lm_head(x) loss = None if targets is not None: loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) return logits, loss # ========================================== # 2. LOAD WEIGHTS SAFELY # ========================================== device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = NanoCoder().to(device) # Replace "nanocoder_base.pth" with your downloaded model path state_dict = torch.load("nanocoder_base.pth", map_location=device, weights_only=True) # Clean DataParallel 'module.' prefixes if they exist clean_state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()} model.load_state_dict(clean_state_dict) print("✅ Nebulixlabs/Nanocoder loaded successfully!") # ========================================== # 3. QUICK FINE-TUNING LOOP EXAMPLE # ========================================== # Setup Optimizer (Use a lower learning rate for fine-tuning) optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) # Dummy Input (Replace with your tokenized DataLoader) # Shape: [Batch Size, Sequence Length] dummy_input = torch.randint(0, VOCAB_SIZE, (4, MAX_SEQ_LEN)).to(device) dummy_target = torch.randint(0, VOCAB_SIZE, (4, MAX_SEQ_LEN)).to(device) model.train() optimizer.zero_grad() # Forward pass logits, loss = model(dummy_input, targets=dummy_target) # Backward pass loss.backward() optimizer.step() print(f"📉 Sample Training Step Complete. Loss: {loss.item():.4f}")