Training CodeParrot from Scratch

Hugging Face has developed CodeParrot, a large GPT-2 model trained from scratch to provide Python code auto-completion capabilities. The project demonstrates how to build a code-generation model—similar to the technology behind GitHub Copilot—using a streamlined training pipeline and a curated dataset of source code.

Dataset Curation and Cleaning

CodeParrot was trained on a dataset derived from the GitHub dump available on Google's BigQuery, filtered specifically for Python files. The initial raw dataset consisted of 20 million files totaling 180 GB.

Data cleaning was critical to model performance, as Hugging Face found that duplicates severely impacted results. Analysis revealed a high concentration of duplicate content:

  • 0.1% of unique files accounted for 15% of all files.
  • 1% of unique files accounted for 35% of all files.
  • 10% of unique files accounted for 66% of all files.

To resolve this, duplicates were removed and cleaning heuristics from the Codex paper were applied, resulting in a cleaned dataset of 50 GB (available as codeparrot-clean).

Model Architecture and Tokenization

Custom Tokenizer

To ensure code tokens are split efficiently, a new tokenizer was trained specifically on the Python dataset. This was achieved by taking a GPT-2 tokenizer and using the train_new_from_iterator() method to adapt it to the source code distribution.

Model Configuration

CodeParrot utilizes the hyperparameters of GPT-2 large, featuring 1.5 billion parameters. To maintain numerical stability and compatibility with the new tokenizer, the following adjustments were made:

  • Embedding Layer: Adjusted to fit the custom tokenizer.
  • Attention Scaling: The scale_attn_by_layer_idx flag was enabled to scale attention by the layer ID.
  • Precision: The reorder_and_upcast_attn flag was used to compute attention in full precision to avoid numerical issues.

Training Implementation

Training was implemented using the 🤗 Accelerate library, allowing the pipeline to scale from a single laptop to multi-GPU environments without code changes.

Technical Training Details

  • Hardware: The models were trained on a 16 x A100 GPU machine.
  • Training Duration: The 110M parameter model took one day to train, while the 1.5B parameter model took one week.
  • Memory Optimization: Gradient checkpointing was enabled to reduce the GPU memory footprint.
  • Data Handling: An IterableDataset was used to stream the 50GB dataset rather than downloading it entirely. To maximize token usage, multiple examples were concatenated with an EOS token and then chunked into fixed context sizes.
  • Distribution: The setup used DistributedDataParallel (DDP), where each GPU worker maintains a copy of the model and aggregates gradients to update weights.

Evaluation and Performance

CodeParrot was evaluated using OpenAI's HumanEval benchmark, which tests code generation across nearly 200 coding challenges. Performance is measured by the pass@k metric, which represents the probability that at least one of k candidate generations passes the unit tests for a given problem.

Despite being trained on significantly fewer tokens (roughly 25-30 billion) compared to GPT-neo (300 billion) or Codex (400 billion total), CodeParrot demonstrated competitive downstream performance, providing a high "bang for your buck" in terms of training efficiency.

Capabilities and Use Cases

CodeParrot can generate functional Python code from prompts, including:

  • Function Implementation: Generating the body of a function based on its name and docstring (e.g., creating a file size retrieval function using os.path.getsize).
  • Boilerplate Generation: Creating unit test structures using the unittest library.
  • API Usage: Generating correct implementation patterns for libraries like transformers (e.g., initializing a BERT classifier).

Users can access CodeParrot via Hugging Face Spaces for code generation and highlighting, or directly through the transformers library using the text-generation pipeline.

Sources