CodeParrotをスクラッチからトレーニングする

Hugging Faceは、Pythonコードのオートコンプリート機能を提供するためにスクラッチからトレーニングされた大規模なGPT-2モデル、CodeParrotを開発しました。このプロジェクトは、GitHub Copilotの背後にある技術と同様のコード生成モデルを、合理化されたトレーニングパイプラインと厳選されたソースコードのデータセットを使用して構築する方法を示しています。

データセットのキュレーションとクリーニング

CodeParrotは、GoogleのBigQueryで利用可能なGitHubのダンプから派生した、Pythonファイルに特化してフィルタリングされたデータセットでトレーニングされました。初期の生データセットは、合計180 GBの2,000万個のファイルで構成されていました。

モデルのパフォーマンスにとって、データのクリーニングは極めて重要でした。Hugging Faceは、重複が結果に深刻な影響を与えることを発見しました。分析の結果、重複コンテンツの非常に高い集中が明らかになりました：

0.1%のユニークなファイルが、全ファイルの15%を占めていました。
1%のユニークなファイルが、全ファイルの35%を占めていました。
10%のユニークなファイルが、全ファイルの66%を占めていました。

これを解決するために、重複が削除され、Codexの論文から得られたクリーニングのヒューリスティックが適用されました。その結果、50 GBのクリーニング済みデータセット（codeparrot-cleanとして利用可能）が作成されました。

モデルのアーキテクチャとトークナイゼーション

カスタムトークナイザー

コードのトークンを効率的に分割するために、Pythonデータセットに特化してトレーニングされた新しいトークナイザーが使用されました。これは、GPT-2のトークナイザーを取得し、train_new_from_iterator()メソッドを使用してソースコードの分布に適応させることで実現されました。

モデル構成

CodeParrotは、15億のパラメータを持つGPT-2 largeのハイパーパラメータを利用しています。数値的な安定性と新しいトークナイザーとの互換性を維持するために、以下の調整が行われました：

Embedding Layer: カスタムトークナイザーに適合するように調整されました。
Attention Scaling: scale_attn_by_layer_idxフラグが有効化され、レイヤーIDによってアテンションをスケーリングするように設定されました。
Precision: 数値的な問題を回避するために、reorder_and_upcast_attnフラグを使用して、アテンションをフル精度で計算するように設定されました。

トレーニングの実装

トレーニングは🤗 Accelerateライブラリを使用して実装されており、コードを変更することなく、単一のノートパソコンからマルチGPU環境へとパイプラインを拡張することが可能です。

技術的なトレーニングの詳細

Hardware: モデルは16 x A100 GPUマシンでトレーニングされました。
Training Duration: 110Mパラメータのモデルはトレーニングに1日、1.5Bパラメータのモデルは1週間を要しました。
Memory Optimization: GPUメモリのフットプリントを削減するために、勾配チェックポインティング（Gradient checkpointing）が有効化されました。
Data Handling: データセット全体をダウンロードするのではなく、50GBのデータセットをIterableDatasetを使用してストリーミングしました。トークンの使用を最大化するために、複数の例をEOSトークンで結合し、その後固定のコンテキストサイズに分割（chunk）しました。
Distribution: セットアップにはDistributedDataParallel (DDP)が使用され、各GPUワーカーがモデルのコピーを保持し、勾配をアグリゲートして重みを更新します。

評価とパフォーマンス

CodeParrotは、約200のコーディング課題をテストするOpenAIのHumanEvalベンチマークを使用して評価されました。パフォーマンスはpass@kメトリックによって測定されます。これは、k個の候補生成結果のうち、少なくとも1つが与えられた問題のユニットテストに合格する確率を表します。

GPT-neo (300 billion) や Codex (400 billion total) と比較して、大幅に少ないトークン数（約25-30 billion）でトレーニングされているにもかかわらず、CodeParrotは競争力のあるダウンストリーム・パフォーマンスを示し、トレーニング効率の面で高い「コストパフォーマンス」を提供しています。

機能とユースケース

CodeParrotは、プロンプトから機能的なPythonコードを生成できます。これには以下が含まれます：

Function Implementation: 関数名とdocstringに基づいて関数の本体を生成すること（例：os.path.getsizeを使用してファイルサイズを取得する関数を作成すること）。
Boilerplate Generation: unittestライブラリを使用してユニットテストの構造を作成すること。
API Usage: 正しいtransformersライブラリのようなライブラリの（例：BERT分類器を初期化すること）の実装パターンを生成すること。

ユーザーは、Hugging Face Spacesを通じて、コード生成とハイライト表示のためにCodeParrotにアクセスしたり、あるいはtransformersライブラリのtext-generationパイプラインを直接使用したりすることができます。