从零开始训练 CodeParrot

Hugging Face 开发了 CodeParrot，这是一个从零开始训练的大型 GPT-2 模型，旨在提供 Python 代码自动补全功能。该项目展示了如何使用精简的训练管道和经过策划的源代码数据集构建代码生成模型——类似于 GitHub Copilot 背后的技术。

数据集策划与清洗

CodeParrot 在从 Google BigQuery 上可用的 GitHub 转储中派生的数据集上进行训练，专门过滤出 Python 文件。初始原始数据集包含 2000 万个文件，总计 180 GB。

数据清洗对模型性能至关重要，因为 Hugging Face 发现重复内容严重影响了结果。分析显示重复内容高度集中：

• 0.1% 的唯一文件占所有文件的 15%。
• 1% 的唯一文件占所有文件的 35%。
• 10% 的唯一文件占所有文件的 66%。

为解决此问题，移除了重复内容，并应用了 Codex 论文中的清洗启发式方法，得到一个清洗后的数据集，大小为 50 GB（可作为 codeparrot-clean 获得）。

模型架构与分词

自定义分词器

为了高效地拆分代码标记，专门在 Python 数据集上训练了一个新的分词器。这是通过获取 GPT-2 分词器并使用 train_new_from_iterator() 方法来适应源代码分布实现的。

模型配置

CodeParrot 使用 GPT-2 large 的超参数，具有 15 亿个参数。为了保持数值稳定性并与新分词器兼容，进行了以下调整：

• 嵌入层：调整以适应自定义分词器。
• 注意力缩放：启用了 scale_attn_by_layer_idx 标志，以按层 ID 缩放注意力。
• 精度：使用了 reorder_and_upcast_attn 标志来在全精度下计算注意力，以避免数值问题。

训练实现

训练使用 🤗 Accelerate 库实现，使得管道能够在不更改代码的情况下从单台笔记本电脑扩展到多 GPU 环境。

技术训练细节

• 硬件：模型在 16 块 A100 GPU 机器上进行训练。
• 训练时长：1.1亿参数的模型训练用了一天，而15亿参数的模型训练用了一周。
• 内存优化：启用了梯度检查点以减少 GPU 内存占用。
• 数据处理：使用 IterableDataset 来流式传输 50GB 数据集，而不是完全下载它。为了最大化 token 使用率，多个示例被连接并加上 EOS token，然后被分割成固定大小的上下文块。
• 分布：该设置使用了 DistributedDataParallel (DDP)，其中每个 GPU 工作进程保持模型的一份副本并聚合梯度以更新权重。

评估与性能

CodeParrot 使用 OpenAI 的 HumanEval 基准进行评估，该基准测试近 200 个编程挑战的代码生成。性能通过 pass@k 指标衡量，该指标表示对于给定问题，至少有 k 个候选生成中的一个能够通过单元测试的概率。

尽管 CodeParrot 的训练 token 数量远少于 GPT-neo（3000 亿）或 Codex（4000 亿总计）——大约只有 250-300 亿——但它在下游任务上表现出竞争力的性能，提供了高训练效率的“性价比”。

能力与使用场景

CodeParrot 能够根据提示生成功能性的 Python 代码，包括：

• 函数实现：根据函数名和文档字符串生成函数体（例如，使用 os.path.getsize 创建文件大小检索函数）。
• 样板代码生成：使用 unittest 库创建单元测试结构。
• API 用法：为如 transformers 之类的库生成正确的实现模式（例如，初始化一个 BERT 分类器）。

用户可以通过 Hugging Face Spaces 访问 CodeParrot 进行代码生成和高亮，或者直接通过 transformers 库使用 text-generation 管道。