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 and was updated to run at M2 MacBook Air at 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 framework, adapted for macOS:

1. Base Pretraining (~5 days):

   • Data: 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 (10k examples)
     • GSM8K (8k examples)
     • ARC-Easy & Challenge (~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 library to run, as it uses a custom architecture implementation optimized for educational clarity and hackability.

License

MIT