Pico Train

The pico-decoder is a LLAMA-style auto-regressive transformer designed for research in the 1M to 1B parameter range. It implements the following components:

• RMSNorm: Root Mean Square Layer Normalization.
• RoPE: Rotary Positional Embeddings.
• Multi-head self-attention: Includes KV-cache support.
• SwiGLU: SwiGLU activation functions.

The implementation can be found in src/model/pico_decoder.py.

Pico Train produces enriched checkpoints that include snapshots of activations and gradients. To perform advanced post-training interpretation, use the companion library pico-analyze.

Pico Analyze can process Pico Train checkpoints to calculate metrics such as:

• CKA (Centered Kernel Alignment)
• PWCCA (Projection Weighted CCA)
• Gini and Hoyer sparsity metrics
• Representation shifts and gradient comparisons

Pico Train uses a config-driven setup where you specify architecture, optimizer, dataset, and logging settings in a YAML file.

To start training:

1. Create or edit a configuration file (e.g., configs/demo.yaml).
2. Execute the training command via Poetry.

By default, checkpoints (including model states, optimizer states, gradients, and activations) are stored in runs/YOUR_RUN_NAME/checkpoints/.

poetry run train --config_path configs/demo.yaml

To use Pico Train, clone the repository, configure your environment variables, and run the setup script to initialize a Poetry virtual environment.

1. Clone the repository

git clone https://github.com/pico-lm/pico-train
cd pico-train

2. Configure Environment

Create a .env file at the root and provide your Hugging Face and Weights & Biases tokens:

export HF_TOKEN=your_huggingface_token
export WANDB_API_KEY=your_wandb_key

3. Install Dependencies

Run the provided setup script to install necessary tools and set up the Poetry environment:

source setup.sh

Pico Train is designed to be extensible for new evaluation metrics. To add a custom metric, follow these steps:

1. Create a Task Module: Implement your evaluation logic in a new module within the evaluation/tasks directory.
2. Expose a Run Function: Your module must expose a function following the pattern run_<metric_name>(model_path, config).
3. Update Configuration: Ensure your EvaluationConfig includes a configuration object for your new metric.
4. Register in Pipeline: Add a conditional branch in the run_evaluation loop in src/evaluation/__init__.py to handle your new metric name.

Currently, the package only supports paloma out of the box, but the architecture is optimized for third-party libraries by storing models in the HuggingFace format.

The extraction process relies on PyTorch forward hooks to capture internal model states without modifying the model's forward pass logic.

1. Hook Setup: CheckpointStateExtractor identifies layers matching the layer_suffixes defined in LearningDynamicsCheckpointingConfig and registers forward hooks on them.
2. Activation Capture: During the forward pass, the hook captures the output of the module at a specific sequence_idx. It uses fabric.all_gather to collect these activations from all distributed processes.
3. Weight/Gradient Capture: Weights are captured once per layer. If compute_gradients=True, a backward pass is performed, and gradients are extracted from the parameters (using deepspeed.utils.safe_get_full_grad if using DeepSpeed) and moved to CPU.

The Trainer supports automatic training recovery.

• Resuming: If checkpointing.training.auto_resume is set to true in your configuration, the Trainer will attempt to load the latest checkpoint using load_checkpoint. This restores the model, optimizer, learning rate scheduler, and the initial_batch_step.
• Data Fast-Forwarding: When resuming, the Trainer automatically fast-forwards the data iterator to the correct batch step to ensure continuity in the training sequence.
• Saving: Checkpoints are saved at intervals defined by checkpointing.save_every_n_steps.

The Trainer uses Lightning Fabric to provide multi-GPU and distributed training support.

• Setup: The Trainer initializes Fabric using the training configuration and wraps the model and optimizer via self.fabric.setup(self.model, self.optimizer).
• Gradient Accumulation: During the training loop, self.fabric.no_backward_sync is used to disable gradient synchronization during accumulation steps, optimizing communication overhead.
• Synchronization: The trainer uses self.fabric.barrier() to synchronize processes, particularly after data fast-forwarding or during cleanup.
• Logging: Metrics are gathered across all processes using self.fabric.all_reduce before being logged to ensure accurate global averages.

The Trainer is driven by a YAML configuration file. The initialization process expects several top-level configuration keys to set up the pipeline:

• model: Configuration for model architecture.
• training: Parameters for the training loop, including max_steps and optimization (e.g., gradient_accumulation_steps).
• data: Configuration for datasets and dataloaders (e.g., batch_size).
• checkpointing: Controls where and how often to save checkpoints, including auto_resume and save_to_hf (HuggingFace).
• monitoring: Settings for logging, including save_to_wandb and logging.log_every_n_steps.
• evaluation: Defines metrics and intervals for model evaluation.

The repository is organized as follows:

• src/model/pico_decoder.py: Core LLAMA-style decoder implementation.
• src/training/trainer.py: Main training loop, distributed/multi-node management, and metric collection.
• src/checkpointing/: Logic for saving model/optimizer states, gradients, activations, and Hugging Face uploading.
• src/config/: Dataclass-based configuration system for hyperparameters and logging.
• configs/: Directory containing YAML configuration files (e.g., demo.yaml).