從零開始訓練 CodeParrot

Hugging Face 開發了 CodeParrot，這是一個從零開始訓練的大型 GPT-2 模型，旨在提供 Python 程式碼自動補全功能。該專案展示了如何使用精簡的訓練流程和精選的原始碼數據集來構建程式碼生成模型——其技術與 GitHub Copilot 背後的技術類似。

數據集策劃與清洗

CodeParrot 是在源自 Google BigQuery 上可用的 GitHub dump 數據集上進行訓練的，並特別針對 Python 檔案進行了過濾。初始的原始數據集包含 2000 萬個檔案，總計 180 GB。

數據清洗對於模型性能至關重要，因為 Hugging Face 發現重複內容會嚴重影響結果。分析顯示重複內容高度集中：

• 0.1% 的唯一檔案佔了所有檔案的 15%。
• 1% 的唯一檔案佔了所有檔案的 35%。
• 10% 的唯一檔案佔了所有檔案的 66%。

為了落實這一點，系統移除了重複項，並應用了 Codex 論文中的清洗啟發式方法，最終得到一個 50 GB 的清洗後數據集（可透過 codeparrot-clean 使用）。

模型架構與分詞

自定義分詞器

為了確保程式碼標記（tokens）能被高效地分割，專門針對 Python 數據集訓練了一個新的分詞器。這是透過採用 GPT-2 分詞器並使用 train_new_from_iterator() 方法來使其適應原始碼的分佈情況來實現的。

模型配置

CodeParrot 使用 GPT-2 large 的超參數，擁有 15 億個參數。為了保持數值穩定性並與新分詞器相容，進行了以下調整：

• Embedding Layer: 調整以適應自定義分詞器。
• Attention Scaling: 啟用了 scale_attn_by_layer_idx 標記，以根據圖層 ID 縮放注意力機制。
• Precision: 使用了 reorder_and_upcast_attn 標記，以全精度計算注意力機制，從而避免數值問題。

訓練實作

訓練是使用 🤗 Accelerate 函式庫實作的，這使得訓練流程可以從單台筆記型電腦擴展到多 GPU 環境，而無需更改程式碼。

技術訓練細節

• Hardware: 模型是在一台配備 16 x A100 GPU 的機器上進行訓練的。
• Training Duration: 1.1 億參數的模型需要一天時間進行訓練，而 15 億參數的模型則需要一週。
• Memory Optimization: 啟用了梯度檢查點（Gradient checkpointing）以減少 GPU 記憶體佔用。
• Data Handling: 使用了 IterableDataset 來串流 50GB 的數據集，而不是將其全部下載。為了最大化標記（token）的使用率，多個範例會與一個 EOS token 串接在一起，然後切分成固定大小的上下文長度。
• Distribution: 配置使用了 DistributedDataParallel (DDP)，其中每個 GPU 工作節點都維護一個模型的副本，並聚合梯度來更新權重。

評估與性能

CodeParrot 使用 OpenAI 的 HumanEval 基準測試進行評估，該測試涵蓋了近 200 個程式碼挑戰。性能是透過 pass@k 指標來衡量的，該指標代表在 k 個候選生成結果中，至少有一個通過了給定問題的單元測試的機率。

儘管與 GPT-neo (300 億) 或 Codex (總計 400 億) 相比，訓練使用的標記數量顯著較少（大約 25-30 億），CodeParrot 展現了具備競爭力的下游性能，在訓練效率方面提供了極高的「性價比」。

能力與使用場景

CodeParrot 可以根據提示詞生成功能性的 Python 程式碼，包括：

• Function Implementation: 根據函數名稱和 docstring 生成函數主體（例如，使用 os.path.getsize 建立一個獲取檔案大小的函數）。
• Boilerplate Generation: 使用 unittest 函式庫建立單元測試結構。
• API Usage: 為 transformers 等函式庫生成正確的實作模式（例如，初始化一個 BERT 分類器）。

使用者可以透過 Hugging Face Spaces 進行程式碼生成與高亮顯示，或者直接透過 transformers 函式庫使用 text-generation 流程來存取 CodeParrot。