Pico: A Lightweight Framework for Studying Learning Dynamics

Pico is a lightweight research framework that aims to demystify how language models learn. Built with simplicity in mind, it provides an efficient way to train and study models of different sizes. Visit our website for more information.

Pico consists of two key components:

1. Pico Training Framework (available on GitHub): A transparent, lightweight codebase for training language models. We use this framework to train a series of language models across scale that we release on this HuggingFace space.
2. Pico Analysis Framework (available on GitHub): A resarch framework to investigate and probe the learning dynamics of models trained using Pico.

This HuggingFace organization hosts our pre-trained models and datasets, while the GitHub repository provides the code to train and analyze your own model suites from scratch.

🤗 HuggingFace Resources (You Are Here)

Pre-trained Model Suite

Our complete suite of models from 1M to 500M parameters trained with Pico:

• pico-tiny (1M parameters)
• pico-small (10M parameters)
• pico-medium (100M parameters)
• pico-large (500M parameters)

🚧 Coming Soon! pico-xl (1B parameters) Watch this space or star our GitHub repository for updates!

All models are trained for 50,000 steps on the pretokenized-dolma dataset. They all see the same training data at each training step, use the same optimizatation process, and share the same model architecture; the only difference between models is the size of their hidden dimension.

Each model includes:

• Advanced training checkpoints (stored every 1,000 steps) that contain:
  • Weights and optimizer states (HuggingFace and Lightning Fabric-compatible versions)
  • Model activations and gradients
  • The batch of training data observed at the given training step
• Wandb logs tracking the learning process
• Pre-computed perplexity scores on the paloma evaluation set

Available Datasets

1. pretokenized-dolma

   • 420B tokens of pre-processed, tokenized and shuffled text extraced from the DOLMA corpus
   • We use this dataset to train our model suite

2. pretokenized-dolma-tiny

   • A smaller version of the pretokenized-dolma corpus for quick experiments

3. pretokenized-paloma

   • A tokenized and shuffled version of the Paloma evaluation corpus
   • The Paloma corpus was carefully curated to be disjoint from the Dolma corpus and provides
   • We use this corpus to evaluate the perplexity of our models

4. pretokenized-paloma-tinsy

   • A sub-sampled version of the pretokenized-dolma corpus

All datasets are tokenized using the OLMo Tokenizer

🔧 GitHub Training Framework

Want to train your own suite of models? Visit our GitHub repository to:

• Train models with custom architectures
• Experiment with different training regimes
• Modify checkpoint saving behavior
• Implement custom evaluation metrics

The training framework makes it easy to:

1. Train multiple models of different sizes
2. Ensure consistent training across all models
3. Save rich checkpoint data for learning dynamics analysis
4. Compare learning dynamics across scales

🛠️ Using the Resources

Using Pre-trained Models (HuggingFace)

from transformers import AutoModelForCausalLM

# Load our pre-trained model
model = AutoModelForCausalLM.from_pretrained("pico-lm/pico-small")

# Access specific checkpoint
model = AutoModelForCausalLM.from_pretrained(
    "pico-lm/pico-small",
    revision="step-xyz"
)

Training Your Own Suite (GitHub)

# Clone the repository
git clone https://github.com/rdiehlmartinez/pico.git && cd pico
source setup.sh

# Configure your model suite
# Edit configs/train.yaml to specify model sizes and training parameters

# Train your suite
python train.py --config configs/train.yaml

📊 Model Details

Architecture

All models use:

• LLAMA-style transformer
• RMSNorm for normalization
• RoPE positional embeddings
• Multi-head attention with KV-cache
• SwiGLU activation function

Training Configuration

Standard configuration (customizable in GitHub training):

• Sequence length: 2048
• Batch size: 1024
• Learning rate: 1e-3
• Weight decay: 0.1
• Gradient clipping: 1.0
• Mixed precision training
• Vocab size: 50280

🔬 Research Applications

Perfect for researchers studying:

• Learning dynamics across model scales
• Mechanistic interpretability
• Architecture and training effects
• Emergent model behaviors

Whether using our pre-trained models or training your own suite, Pico provides the tools needed for in-depth learning dynamics research.

🤝 Contributing

Contributions welcome on both platforms:

• HuggingFace: Model weights, datasets, and evaluation results
• GitHub: Training framework improvements, analysis tools, and documentation

📫 Contact

• GitHub: rdiehlmartinez/pico
• Author: Richard Diehl Martinez

🔍 Citation

@software{pico2024,
    author = {Diehl Martinez, Richard},
    title = {Pico: Framework for Training Tiny Language Models},
    year = {2024},
}