Filtered Pile 13B

A cluster-scale pipeline that cleans web text for language-model pretraining, and a controlled experiment to measure what the cleaning is worth: two GPT-2-small runs on token-matched filtered and unfiltered data drawn from the same source files.

Built on Stanford CS336 Assignment 4 (Data). Every number below comes from a file in this repository; run python metrics/show_results.py to recompute them.

Highlights

Same prompt, both models (seed and sampling identical; full text in the CSV, prompt 16):

Prompt: The human heart pumps blood by

Filtered model (GPT-2 ppl 6.3): …pumping it into the veins. This is called heart pump activity. The heart pump activity is a means of pumping blood to the lungs from a place in the body called the heart.

Unfiltered model (GPT-2 ppl 16.4): …the end of the day. This is a real problem, and it will never pass. I think I have to mention that since I wrote this in 2009, I have been able to get a good feeling for your situation.

Key results

Pipeline

flowchart LR
    A[Pile shard / CC WET file] --> B[HTML detect + extract]
    B --> D[Language ID<br/>fastText lid.176]
    D --> E[Gopher rules]
    E --> F[Quality classifier<br/>fastText]
    F -.optional.-> G[NSFW / toxicity]
    G --> C[PII masking<br/>email, phone, IP]
    C -.optional.-> H[MinHash-LSH dedup]
    H --> I[GPT-2 tokenise<br/>+ token matching]
    I --> J[Hugging Face]

PII masking now runs last, on kept documents only, since it is the most expensive stage; the datasets and benchmarks below were produced with the earlier order (PII masking before language ID), and the new order has not been re-benchmarked yet. Inputs can be Pile parquet shards, JSONL or Common Crawl WET files (--input-format wet). Each shard is processed by its own worker process (--workers), and shards are spread across cluster nodes with SLURM job arrays. Every stage is timed per document and every removal is counted by reason; --metrics-json writes a report with throughput, per-stage CPU time, latency percentiles, removals by reason and peak memory (cs336_data/pipeline_metrics.py).

Where 2.44 million documents went

Production run on shards 180–209 (30 shards, 4 workers, 2,860.7 s, 853.7 rows/s). Evidence: metrics/production_run_180-209_summary.txt.

458,929 PII spans (emails, phone numbers, IP addresses) were masked.

Benchmarks

Controlled runs on shards 300+ (not part of the training data) on the cluster nodes described in Hardware, with production filter settings unless stated.

Hardware

All data processing ran as SLURM jobs on the u22 partition of IIIT Hyderabad's Ada cluster (CPU details checked with lscpu on a node that ran production tasks):

Training ran separately on one NVIDIA RTX 3090 (24 GB, bf16) in a workstation with 16 CPU threads and 94 GB RAM.

Worker scaling

16 shards × 20,000 rows. Evidence: metrics/bench/scale_w{1,2,4,8,16}.json.

Per-worker throughput falls as workers are added. Likely causes, not yet isolated: SLURM allocates hardware threads, so workers share physical cores; each worker loads its own 2.7 GB of fastText models from network storage; and the nodes are shared with other jobs.

What each stage costs

One change at a time against production settings: 8 workers, 8 shards × 20,000 rows. Evidence: metrics/bench/stage_*.json.

Findings:

• PII masking is the most expensive filter: 39.5% of CPU time. It ran on every document although only 45% survive filtering; it now runs after the filters (not yet re-benchmarked).
• Language ID is largely redundant on the Pile: without it, the Gopher rules reject non-English documents anyway and the kept share is unchanged.
• NSFW / toxicity filtering is poor value here: +32% runtime to remove 0.14% of documents.
• The two quality filters overlap: without Gopher, the classifier removes twice as many documents, yet the kept share still rises from 45.2% to 64.9%.

Per-document latency by stage

Evidence: metrics/bench/stage_production.json.

A CPU profile of one worker attributes about 38% of its time to Python regex calls (metrics/bench/profile_tottime.txt).

Optimising MinHash dedup

Dedup was 24.5× slower than all filtering combined. The cause was the signature step in cs336_data/lsh/hashing.py: for each of ~1,700 shingles per document and each of 256 hash functions, it computed (a·x + b) mod (2^61 − 1) with Python integers (about 434,000 operations per document at ~170 ns each), on a single core because the pipeline never passed its worker count to the dedup step.

Two changes:

1. Vectorised, exact arithmetic. The hash is computed for all shingles and hash functions at once with NumPy. Operands are split into 32-bit halves and reduced using 2^64 ≡ 8 and 2^61 ≡ 1 (mod 2^61 − 1), so every intermediate fits in 64 bits and signatures are bit-identical to the original (tests/test_minhash_vectorized.py).
2. Parallel signatures. process_html_pipeline.py now passes --workers to the dedup step.

Same benchmark before and after (8 workers, 8 shards × 20,000 rows). Evidence: metrics/bench/stage_plus_dedup.json, metrics/bench/stage_plus_dedup_v2.json.

The signature step is still about 87% of dedup time; per-document Unicode normalisation and one SHA-1 per shingle are the next targets. 99.99% of LSH candidates are false positives, but checking all 3.8M took only about 37 s.

Token-matched datasets

For every source shard i, pile_filtered_13B holds the pipeline's output for shard i, and pile_unfiltered_13B holds a seeded random sample of raw documents from the same shard, drawn until it has exactly as many GPT-2 tokens (the last document is truncated). Filtering is the only difference between the two (cs336_data/prepare_matched_data.py).

Evidence: metrics/matched_data_summary.json, produced by cs336_data/verify_matched_data.py, which checks every file, every per-shard token count and every token file size on both repositories.

251 shards were filtered with language ID, PII masking, Gopher rules and the quality classifier; 19 early shards also had NSFW/toxicity filtering. Each shard's manifest records its pipeline columns. Each repository contains tokens/train-XXXXX.bin (uint16, <|endoftext|> after every document), text/train-XXXXX.parquet and manifests/train-XXXXX.json.

Training

Both runs use the assignment's GPT-2-small configuration (12 layers, d_model 768, 162M parameters, context 512, global batch 256 sequences, cosine schedule, seed 0); only the dataset differs. Shards stream from Hugging Face during training (cs336-basics/cs336_basics/streaming_data.py), so a run needs about 1 GB of disk instead of 27 GB.

Both runs trained for 25,000 steps (3.28B tokens each, ~17.7 h each on one RTX 3090): W&B runs filtered_25k and unfiltered_25k. Every 2k-step checkpoint and the final model are on the Hub.

Validation loss during training (Paloma C4)

Evidence: metrics/train_25k_eval_curves.csv, exported from W&B.

The filtered model is ahead at all 13 evaluations (gap 0.088–0.148). Each estimate averages 1,000 batches, so its uncertainty is roughly ±0.01, an order of magnitude below the gap. The unfiltered model's final loss is what the filtered model reached at about 17.4k steps (between 3.590 at 16k and 3.545 at 18k), so on this benchmark the filtered data is worth roughly 1.4× the compute. Each arm has one seed.

Final models on two held-out sets, and generations

inference_check/compare_models.py evaluates both final models on the same 1.64M-token sample of each set (random 512-token windows) and generates from 50 prompts with the same seed per prompt (temperature 0.7, top-k 50, 100 new tokens). Each generation is scored by the original GPT-2 (124M) as a neutral fluency judge: the perplexity GPT-2 assigns to the generated continuation. Evidence: inference_check/results/compare_25k_summary_20261002_004204.json; all 50 prompts with both models' outputs are in inference_check/results/compare_25k_20261002_004204.csv.

What this shows:

• Filtering helps on clean web text and hurts on raw Pile text. The filtered model is better on Paloma C4 and writes text GPT-2 finds more natural, but the unfiltered model is much better on held-out raw Pile. The filters removed most code, LaTeX, tables and other structured text; that content is common in raw Pile and very predictable (repetitive syntax), so the model that trained on it wins there. Filtering is a choice about the target distribution, not a free improvement.
• Repetition is similar for both (about 6% repeated 4-grams on average), so the GPT-2 fluency gap is not explained by one model looping more.

GPT-2 perplexity also rewards repetition (one filtered generation that loops on a single sentence scores 1.5), which is why the repeated-4-gram rate is reported alongside it. A same-prompt sample from both models is shown in Highlights.

Throughput on the training GPU

Measured on the training GPU, an RTX 3090 in bf16 (evidence: metrics/train_smoke_rtx3090.txt, W&B run 12fd2vz3):

Baseline: the earlier model

Trained on 0.82B tokens (batch 16 × 512 × 100k steps) on an RTX 2080 Ti in 10.2 h (W&B run s1ml9n3n). Re-evaluated with inference_check/infer_prev.py: Paloma C4 loss 3.797 (perplexity 44.6), matching the 3.81 logged during training. Evidence: inference_check/results/prev_model_summary_20260927_172053.json.

24 generations (8 prompts × 3 samples, temperature 0.7, top-k 50) are in inference_check/results/prev_model_samples_20260927_172053.csv. The text is grammatical and holds a register (news, recipe, encyclopedia) for a sentence or two, then drifts or loops, for example:

The best way to cook pasta is to make it a part of the recipe, such as a salad, a salad or a salad.

The code prompt def fibonacci(n): produces no code, consistent with prose-only filtered training data.

How to run

Commands run from the repository root. The SLURM scripts use the paths of the IIIT Ada cluster; change the paths at the top of each script for another cluster.

1. Environment, models and tests

uv sync

mkdir -p model
wget -O model/lid.176.bin https://dl.fbaipublicfiles.com/fasttext/supervised-models/lid.176.bin
wget -O model/jigsaw_fasttext_bigrams_nsfw_final.bin https://huggingface.co/allenai/dolma-jigsaw-fasttext-bigrams-nsfw/resolve/main/model.bin
wget -O model/jigsaw_fasttext_bigrams_hatespeech_final.bin https://huggingface.co/allenai/dolma-jigsaw-fasttext-bigrams-hatespeech/resolve/main/model.bin

cd urls && python train_quality_classifier.py --positive positive.txt --negative negative.txt --output quality_model && cd ..

uv run pytest

fastText 0.9.3 needs NumPy < 2.

2. Filter shards on one node, with a metrics report

python cs336_data/process_html_pipeline.py --start-index 300 --max-files 8 --max-rows 20000 --workers 8 --no-content-classification --output-dir out/ --metrics-json metrics/run.json

Add --no-dedup to skip MinHash-LSH, or --push-to-hf --hf-repo-id USER/REPO to upload each shard.

3. Production run on the cluster

sbatch run_pipeline.sh
sbatch jobs/matched_new_shards.sh
sbatch --array=0-8 jobs/matched_reuse_shards.sh
sbatch --export=ALL,SHARDS=136-155 jobs/matched_retry_resumable.sh
python cs336_data/verify_matched_data.py --summary-json metrics/matched_data_summary.json

run_pipeline.sh filters 30 shards and uploads them; the matched_* jobs build token-matched pairs from new and already-filtered shards; matched_retry_resumable.sh rebuilds and uploads with resumable, time-bounded retries.

4. Benchmarks

bash jobs/submit_benchmarks.sh
python metrics/show_results.py

5. Train one arm of the comparison

tmux new -s a4_filtered 'cs336-basics/scripts/launch_matched_run.sh filtered'
tmux new -s a4_unfiltered 'cs336-basics/scripts/launch_matched_run.sh unfiltered'

The launcher derives micro-batch, gradient accumulation and bf16/fp16 from the visible GPUs so the global batch stays at 256. Use MICRO_BATCH=16 on smaller GPUs and Hydra overrides with ++, e.g. ++training.train_steps=25000.

6. Evaluate and sample a model

python inference_check/infer_prev.py /path/to/model_dir

It uses the GPU only when at least 3 GB is free and otherwise the CPU, and writes a CSV of prompts and generations plus a summary JSON to inference_check/results/.

To compare the two trained models (loss, perplexity and next-token accuracy on Paloma C4 and held-out raw Pile, plus 50 prompts scored by GPT-2):

python inference_check/compare_models.py work_dir/ /path/to/paloma_c4_100_val.bin

Repository layout

Engineering notes

• Worker crashes were swallowed and jobs exited 0; failures are now listed and fail the SLURM job.
• #SBATCH directives placed after the first command were silently ignored; job scripts now declare resources correctly.
• fastText 0.9.3 breaks on NumPy 2 (np.array(..., copy=False)); NumPy is pinned below 2.
• Exact line dedup opened its output files in read mode; fixed, and its test passes again.
• Uploads through Hugging Face's Xet backend hung indefinitely from the cluster while plain LFS ran at full link speed; uploads now use LFS with per-attempt timeouts, resumable retries, and staging kept until success.
• Evaluation during training ran without torch.no_grad() and ran out of GPU memory; fixed without changing the computed loss.

Acknowledgements

The assignment, handout and staff model implementation in cs336-basics come from Stanford CS336 (Spring 2025). NSFW and hate-speech classifiers are from AllenAI Dolma; language identification uses fastText lid.176. Source data is EleutherAI's deduplicated Pile. See also the OLMo 3 paper.