Training CodeParrot 🦜 from Scratch

Summary

Training code generation models from scratch requires careful data filtering, tokenization, and compute management. To build CodeParrot, a 1.5-billion-parameter Python generation model, engineers extracted 20 million Python files from BigQuery and cleaned the resulting 180 GB corpus down to 50 GB after discovering extreme duplication. A custom GPT-2 tokenizer was trained over streamed samples, and the GPT-2 large architecture was initialized with layer-scaled, full-precision attention mechanisms. The training pipeline leveraged Hugging Face Accelerate alongside gradient checkpointing and a streaming iterable dataset designed to yield fixed-length concatenated token sequences. CodeParrot was trained on roughly 25 to 30 billion tokens and evaluated on coding tasks using the pass@k metric.

Takeaways

1. 1

   Severe duplication in the raw 180 GB dataset impacted model performance, with 0.1% of unique files representing 15% of all files and 10% making up 66%.

2. 2

   Deduplicating and applying Codex cleaning heuristics reduced the raw 20-million-file dataset from 180 GB down to 50 GB.

3. 3

   Training stability and memory efficiency were managed using layer-indexed attention scaling, full-precision attention computation, and gradient checkpointing.