|
|
--- |
|
|
language: |
|
|
- en |
|
|
license: mit |
|
|
library_name: picochat |
|
|
tags: |
|
|
- pytorch |
|
|
- mps |
|
|
- macbook |
|
|
- education |
|
|
datasets: |
|
|
- HuggingFaceFW/fineweb-edu |
|
|
pipeline_tag: text-generation |
|
|
inference: false |
|
|
spaces: |
|
|
- MGow/PicoChat |
|
|
--- |
|
|
|
|
|
# PicoChat |
|
|
|
|
|
**PicoChat** is a 335M parameter language model trained entirely from scratch on a MacBook Air M2 (16GB RAM) in approximately 6 days. The code is based on Andrej Karpathy's |
|
|
[NanoChat](https://github.com/karpathy/nanochat) and was updated to run at M2 MacBook Air at [PicoChat](https://github.com/MichalGow/PicoChat). |
|
|
It serves as a "lab notebook" proof-of-concept for training capable small language models (SLMs) on consumer hardware using pure PyTorch and MPS (Metal Performance Shaders). |
|
|
|
|
|
> **Links:** |
|
|
> - **Space:** https://huggingface.co/spaces/MGow/PicoChat |
|
|
|
|
|
## Model Details |
|
|
|
|
|
- **Architecture:** GPT-style Transformer (Decoder-only) |
|
|
- **Parameters:** ~335 Million |
|
|
- **Layers:** 16 |
|
|
- **Embedding Dimension:** 1024 |
|
|
- **Heads:** 8 Query heads, 8 KV heads (GQA) |
|
|
- **Context Length:** 1024 tokens |
|
|
- **Vocabulary:** 65,536 (Custom BPE) |
|
|
- **Training Data:** ~377 Million tokens |
|
|
- **Precision:** Trained in mixed precision (bfloat16/float32) on MPS. |
|
|
|
|
|
### Key Features |
|
|
- **Rotary Embeddings (RoPE)** (No absolute positional embeddings) |
|
|
- **SwiGLU** (ReLU²) activations |
|
|
- **RMSNorm** (with no learnable parameters) |
|
|
- **Untied embeddings** (Input and Output embeddings are separate matrices) |
|
|
- **Grouped Query Attention (GQA)** supported (configured here as 1:1 for simplicity) |
|
|
|
|
|
## Training Recipe |
|
|
|
|
|
The model was trained in three phases using the [nanochat](https://github.com/karpathy/nanochat) framework, adapted for macOS: |
|
|
|
|
|
1. **Base Pretraining (~5 days):** |
|
|
- **Data:** [FineWeb-Edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) (100B subset, shuffled) |
|
|
- **Steps:** ~60,000 |
|
|
- **Tokens:** ~442M |
|
|
- **Objective:** Next token prediction |
|
|
|
|
|
2. **Midtraining (~16 hours):** |
|
|
- **Data:** Mixed pretraining data + synthetic conversation/instruction formats. |
|
|
- **Tokens:** ~33M |
|
|
- **Objective:** Adaptation to chat format and Q&A style. |
|
|
|
|
|
3. **Supervised Finetuning (SFT) (~4 hours):** |
|
|
- **Data Mixture:** |
|
|
- [SmolTalk](https://huggingface.co/datasets/HuggingFaceTB/smoltalk) (10k examples) |
|
|
- [GSM8K](https://huggingface.co/datasets/gsm8k) (8k examples) |
|
|
- [ARC-Easy & Challenge](https://huggingface.co/datasets/ai2_arc) (~3.4k examples) |
|
|
- Identity & Synthetic Spelling tasks (~1.6k examples) |
|
|
- **Steps:** 1,000 (Batch size 8) |
|
|
- **Tokens:** ~1M |
|
|
- **Objective:** Instruction following and personality alignment. |
|
|
|
|
|
## Character & Limitations |
|
|
|
|
|
- **Personality:** The model is designed to be "a bit silly, often wrong, sometimes delightful." It is not a rigid assistant but a fun research artifact. |
|
|
- **Hallucinations:** As a small 300M model trained on limited data, it will confidently hallucinate facts. |
|
|
- **Context Window:** Limited to 1024 tokens. |
|
|
- **Safety:** The model has not gone through extensive safety alignment or RLHF. It generally behaves like a base model with some instruction following capabilities. |
|
|
|
|
|
## Usage |
|
|
|
|
|
This model requires the [picochat](https://github.com/MichalGow/PicoChat) library to run, as it uses a custom architecture implementation optimized for educational clarity and hackability. |
|
|
|
|
|
## License |
|
|
|
|
|
MIT |
|
|
|
