Please read the main post at https://confirmlabs.org/posts/catalog.html.
Further details and reproduction commands here.
The datasets:
- pile_scan_4
- pile_bigrams
- pile_bigram_prefixes
- pile_top_bigrams
- pile_trigrams
- pile_trigram_prefixes
- pile_top_trigrams
We use the deduplicated Pile. The data has been tokenized and pre-processed into the MosaicML Streaming format (MDS). This was not necessary but I had already done it for a separate project and found it convenient to be able to stream the tokenized dataset rather than downloading the entire thing and re-tokenizing it.
The group by id0, id1, id2 operations necessary to construct bigrams/trigrams
are memory hungry. The default DuckDB group by operation ran out of memory so
I used a sharding strategy with Parquet files and DuckDB. This was probably
overengineering where I should've just used an instance with 3TB of memory, but
it works and it wasn't very hard to setup. Using DuckDB over other parquet
handling tools seemed to make sense because:
- It's very fast and uses memory more sparingly compared to Pandas.
- Many bulk parquet-processing tools like Spark are much more involved to setup. DuckDB is just a pip install.
- I eventually put the tables into a DuckDB format anyway for a nice SQL interface.
To perform the sharding, we stream The Pile from S3 and split into Parquet
files by id0 % 64. Separating into chunks by id0 makes the later
group by id0, id1, id2 step simple to do out-of-core.
We run prompt scanning only on the first shard of
EleutherAI/the_pile_deduplicated.
There are 1650 shards total. Shard #1 contains ~112.5 million tokens. Running
on the first shard provides plenty of data for our current purposes but scaling
up to a larger portion of the dataset would be straightforward. I would
recommend using multiple GPUs because a single shard required ~19 hours to run
on a single A6000.
Set up the environment:
mamba env create # or conda env create, but mamba is much faster
conda activate catalog
poetry install
Construct the MDS format streaming
dataset in S3. out-root can be set to
a local directory if you prefer.
python convert_dataset_hf.py --out-root="s3://caiplay/the_pile_deduplicated_16"
Run the bigram/trigram dataset construction code. On a r6a.16xlarge instance
with 64 cores and 512GB RAM, the total runtime is ~3 hours, dominated by step 1
at 2 hours. The total disk size required is ~1TB. The final tables are ~80 GB
in DuckDB format.
python catalog/build_ngrams.py all \
--work-dir=data \
--db-filename=data/ngrams.db \
--batches 1000
Run prompt scanning on the first shard of The Pile. On a single A6000 GPU, the total runtime is ~19 hours.
python catalog/scanner.py all --work-dir=./data/scan_2.8b --batches 1000