Causal shift

Causal shift experiments codebase

Setup

Using uv (preferred):

uv sync
# --no-cache-dir
uv pip install --no-build-isolation flash-attn

Using pip:

pip install .
pip install flash-attn

Structure

src/
├── data/                  # Data pre-processing (generate memmap)
├── scripts/               # Example scripts for training
├── config.py              # Configuration file
├── eval.py                # Evaluate NTP accuracy
├── model.py               # Model definition
└── train.py               # Training script

# Note: train_gated.py and model_gated.py are for gated models.

Quick start

Data preparation

Generate memmap files (train.bin and val.bin) for training on 10B tokens from fineweb-edu dataset:

src/data/fineweb-edu/download.sh /path/to/local/dir # download data (parquet)
python src/data/fineweb-edu/prepare.py --path /path/to/local/dir --cache_dir /path/to/cache/dir # generate np memmap

Training

The training scripts will train a GPT2-like model on the fineweb-edu dataset.

See src/scripts/main.sh for an example training script. Make sure you have the data prepared before running the training script. The script runs a model training on all available GPUs using torchrun.

Also see src/scripts/gated.sh for training a gated model and src/scripts/skip.sh for training a model by pseudo-skipping residual connections.

Note: training should take less than 16 A100 GPU hours.

Acknowledgements

This repository is built upon the nanoGPT repository.