---
language:
- en
license: mit
tags:
- code
- python
- causal-lm
- from-scratch
- gqa
- rope
- swiglu
- fim
- qk-norm
datasets:
- nampdn-ai/tiny-codes
- ise-uiuc/Magicoder-OSS-Instruct-75K
- iamtarun/python_code_instructions_18k_alpaca
- bigcode/the-stack-smol
- flytech/python-codes-25k
- iamtarun/code_instructions_120k_alpaca
---

# PyCraft-1: A 55M Python Code LLM Trained From Scratch on Consumer Hardware

## Model Description

PyCraft-1 is a 55.3M parameter decoder-only transformer trained
**entirely from scratch** on Python code using a single NVIDIA RTX 3050
laptop GPU (4GB VRAM). It demonstrates that a domain-specific code LLM
can be trained without cloud compute, following a quality-first data
curriculum inspired by the Phi-1 "Textbooks Are All You Need" approach.

## Architecture

| Component | Choice |
|---|---|
| Architecture | Decoder-only transformer |
| Parameters | 55.3M |
| Attention | Grouped Query Attention (8Q / 2KV heads) |
| Positional encoding | RoPE |
| QK-Norm | RMSNorm on Q and K (OLMo 2 / Qwen 3, 2025) |
| FFN | SwiGLU |
| Normalisation | RMSNorm pre-norm |
| Training objective | Causal LM + Fill-in-the-Middle (FIM, 50%) |
| Context window | 1024 tokens |
| Vocabulary | 32,000 (custom BPE, Python-tuned) |

## Training

**Pretraining:**
- 309,221 curated Python examples from 6 open sources
- Quality curriculum: scored on 5 heuristics, ordered best-first
- 4,000 steps | 1.05B tokens | Loss 1.16 | PPL 3.2

**Supervised Fine-Tuning:**
- Magicoder-OSS-Instruct-75K (Python subset, 40k examples)
- 400 steps | Loss 1.15 | PPL 3.15

**Hardware:** NVIDIA RTX 3050 Laptop GPU 4GB, Ryzen 7 6000, 16GB RAM
**Total training time:** ~22 hours

## Novel Contributions

1. **Quality-first data curriculum** — 309k Python examples scored and
   ordered by educational value (docstrings, type hints, comments,
   naming conventions, length). Validates Phi-1 hypothesis at the
   resource-constrained regime.

2. **QK-Norm** — RMSNorm applied to Q and K before RoPE, adopted from
   OLMo 2 and Qwen 3 (2025), improving training stability in small models.

3. **FIM pretraining on 4GB VRAM** — Fill-in-the-Middle objective using
   PSM format trained on a consumer GPU via gradient checkpointing and
   BF16 mixed precision.

4. **Full reproducibility** — complete open-source pipeline runnable on
   consumer hardware in under one week.

## Evaluation

| Metric | Value |
|---|---|
| Pretraining loss | 1.16 |
| Pretraining PPL | 3.20 |
| SFT loss | 1.15 |
| SFT PPL | 3.15 |
| Held-out PPL (binary search) | 1.4 |
| Held-out PPL (Stack class) | 1.7 |
| Held-out PPL (average, 5 samples) | 2.16 |

## Example Usage

```python
# Clone the repo and install dependencies first
# pip install torch safetensors tokenizers datasets

import torch
from safetensors.torch import load_file
from model.config import get_config_120m
from model.pycraft_model import PyCraftModel
from tokenizer.tokenizer_utils import PyCraftTokenizer

device    = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = PyCraftTokenizer("tokenizer/vocab/tokenizer.json")

cfg            = get_config_120m()
cfg.vocab_size = 32000
cfg.dropout    = 0.0

model = PyCraftModel(cfg).to(device)
model.load_state_dict(load_file("model.safetensors", device=device))
model.eval()

prompt = "def is_palindrome(s: str) -> bool:\n    "
ids    = tokenizer.encode(prompt)
inp    = torch.tensor(ids, dtype=torch.long).unsqueeze(0).to(device)

with torch.no_grad():
    out = model.generate(inp, max_new_tokens=80, temperature=0.7, top_k=40)

print(tokenizer.decode(out[0].tolist()))
```

## Limitations

- 55M parameters: generates plausible Python but may produce logical errors
- Context window: 1024 tokens
- Best on standard algorithms, data structures, and common library patterns
- Not a conversational assistant

## Citation

```bibtex
@misc{inamdar2026pycraft,
  title={PyCraft-1: Training a Python Code LLM From Scratch on Consumer Hardware},
  author={Inamdar, Rohan},
  year={2026},
  institution={University of Manchester, MSc Artificial Intelligence},
}
```

## Links

- GitHub: https://github.com/irohan0/pycraft-llm
- Author: https://www.linkedin.com/in/rohan-inamdar-47aa4b251