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

トレーニングするために、スクラッチからCodeParrot 🦜を作成します

このブログポストでは、GitHub CoPilotの背後にある技術を構築するために必要なものについて説明します。GitHub CoPilotは、プログラマがコードを書く際に提案を行うアプリケーションです。このステップバイステップガイドでは、ゼロから完全にトレーニングされた大規模なGPT-2モデルであるCodeParrot 🦜を訓練する方法を学びます。CodeParrotはPythonのコードを自動補完することができます – こちらで試してみてください。さあ、ゼロから構築してみましょう!

ソースコードの大規模なデータセットの作成

まず必要なものは、大規模なトレーニングデータセットです。Pythonのコード生成モデルを訓練することを目指して、GoogleのBigQueryで利用可能なGitHubのダンプにアクセスし、すべてのPythonファイルに絞り込みました。その結果、180GBのデータセットがあり、2000万のファイルが含まれています(こちらで入手可能)。初期のトレーニング実験の結果、データセットの重複はモデルの性能に深刻な影響を与えることがわかりました。データセットを調査すると、次のことがわかりました:

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

詳細は、このTwitterスレッドで調査結果について詳しくご覧いただけます。重複を削除し、CoPilotの背後にあるモデルであるCodexの論文で見つかった同じクリーニングヒューリスティックを適用しました。CodexはGitHubのコードでファインチューニングされたGPT-3モデルです。

クリーニングされたデータセットはまだ50GBの大きさであり、Hugging Face Hubで利用可能です:codeparrot-clean。これで新しいトークナイザーを設定し、モデルを訓練することができます。

トークナイザーとモデルの初期化

まず、トークナイザーが必要です。コードを適切にトークンに分割するために、コード専用のトークナイザーをトレーニングしましょう。既存のトークナイザー(例えばGPT-2)を取り、train_new_from_iterator()メソッドで独自のデータセットでトレーニングします。それから、Hubにプッシュします。コードの例からインポートや引数のパース、ログ出力は省略していますが、前処理やダウンストリームタスクの評価を含めた完全なコードはこちらで見つけることができます。

# トレーニング用のイテレーター
def batch_iterator(batch_size=10):
    for _ in tqdm(range(0, args.n_examples, batch_size)):
        yield [next(iter_dataset)["content"] for _ in range(batch_size)]

# ベースのトークナイザー
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
base_vocab = list(bytes_to_unicode().values())

# データセットのロード
dataset = load_dataset("lvwerra/codeparrot-clean", split="train", streaming=True)
iter_dataset = iter(dataset)

# トレーニングと保存
new_tokenizer = tokenizer.train_new_from_iterator(batch_iterator(),
                                                  vocab_size=args.vocab_size,
                                                  initial_alphabet=base_vocab)
new_tokenizer.save_pretrained(args.tokenizer_name, push_to_hub=args.push_to_hub)

Hugging Faceコースでトークナイザーについて詳しく学ぶことができます。

その目立たないstreaming=True引数に気づきましたか?この小さな変更には大きな影響があります。50GBのフルデータセットをダウンロードする代わりに、必要なときに個々のサンプルをストリーミングで提供するため、ディスクスペースを大幅に節約できます。ストリーミングについての詳細は、Hugging Faceコースをご覧ください。

さて、新しいモデルを初期化します。GPT-2 large(15億パラメータ)と同じハイパーパラメータを使用し、埋め込みレイヤーを新しいトークナイザーに合わせて調整し、いくつかの安定性の調整も追加します。 scale_attn_by_layer_idxフラグは、アテンションをレイヤーIDでスケーリングすることを保証し、reorder_and_upcast_attnは主に数値的な問題を回避するためにアテンションを完全精度で計算することを保証します。新しく初期化したモデルをトークナイザーと同じリポジトリにプッシュします。

# Pythonコードのトークン化のためにトレーニングされたcodeparrotトークナイザーをロード
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name)

# 設定
config_kwargs = {"vocab_size": len(tokenizer),
                 "scale_attn_by_layer_idx": True,
                 "reorder_and_upcast_attn": True}

# 設定とモデルのロード、ハブにプッシュ
config = AutoConfig.from_pretrained('gpt2-large', **config_kwargs)
model = AutoModelForCausalLM.from_config(config)
model.save_pretrained(args.model_name, push_to_hub=args.push_to_hub)

効率的なトークナイザと新たに初期化されたモデルを持つことで、実際のトレーニングループを開始できます。

トレーニングループの実装

🤗 Accelerateライブラリを使用してトレーニングを行います。このライブラリを使用すると、コードの一行も変更することなく、ノートパソコンからマルチGPUマシンにトレーニングをスケーリングすることができます。アクセラレータを作成し、いくつかの引数のハウスキーピングを行うだけです。

accelerator = Accelerator()
acc_state = {str(k): str(v) for k, v in accelerator.state.__dict__.items()}

parser = HfArgumentParser(TrainingArguments)
args = parser.parse_args()
args = Namespace(**vars(args), **acc_state)
samples_per_step = accelerator.state.num_processes * args.train_batch_size
set_seed(args.seed)

これでトレーニングの準備が整いました!新しいトークナイザとモデルを持つリポジトリを huggingface_hub クライアントライブラリを使用してクローンします。この実験では新しいブランチにチェックアウトします。このセットアップで、多くの実験を並行して実行し、最終的に最良の結果をメインブランチにマージすることができます。

# モデルリポジトリをクローン
if accelerator.is_main_process:
    hf_repo = Repository(args.save_dir, clone_from=args.model_ckpt)

# リポジトリで新しいブランチにチェックアウト
if accelerator.is_main_process:
    hf_repo.git_checkout(run_name, create_branch_ok=True)

ローカルリポジトリからトークナイザとモデルを直接ロードすることができます。大規模なモデルを扱っているので、トレーニング中のGPUメモリの使用量を減らすために勾配チェックポイントをオンにすることがあります。

# モデルとトークナイザをロード
model = AutoModelForCausalLM.from_pretrained(args.save_dir)
if args.gradient_checkpointing:
    model.gradient_checkpointing_enable()
tokenizer = AutoTokenizer.from_pretrained(args.save_dir)

次はデータセットです。固定のコンテキストサイズで例を生成するデータセットを使用して、トレーニングをより簡単に行います。データを無駄にすることなく(一部のサンプルは短すぎたり長すぎたりすることがあります)、多くの例をEOSトークンで結合してチャンク化することができます。

一緒に準備するシーケンスが多ければ多いほど、前の図のグレーのトークン(廃棄されるトークン)の割合が小さくなります。データセットを事前にすべて準備するのではなく、データセットをストリーミングするために IterableDataset を使用します。完全なデータセットクラスは以下のようになります:

class ConstantLengthDataset(IterableDataset):
    def __init__(
        self, tokenizer, dataset, infinite=False, seq_length=1024, num_of_sequences=1024, chars_per_token=3.6
    ):
        self.tokenizer = tokenizer
        self.concat_token_id = tokenizer.bos_token_id
        self.dataset = dataset
        self.seq_length = seq_length
        self.input_characters = seq_length * chars_per_token * num_of_sequences
        self.epoch = 0
        self.infinite = infinite

    def __iter__(self):
        iterator = iter(self.dataset)
        more_examples = True
        while more_examples:
            buffer, buffer_len = [], 0
            while True:
                if buffer_len >= self.input_characters:
                    break
                try:
                    buffer.append(next(iterator)["content"])
                    buffer_len += len(buffer[-1])
                except StopIteration:
                    if self.infinite:
                        iterator = iter(self.dataset)
                        self.epoch += 1
                        logger.info(f"Dataset epoch: {self.epoch}")
                    else:
                        more_examples = False
                        break
            tokenized_inputs = self.tokenizer(buffer, truncation=False)["input_ids"]
            all_token_ids = []
            for tokenized_input in tokenized_inputs:
                all_token_ids.extend(tokenized_input + [self.concat_token_id])
            for i in range(0, len(all_token_ids), self.seq_length):
                input_ids = all_token_ids[i : i + self.seq_length]
                if len(input_ids) == self.seq_length:
                    yield torch.tensor(input_ids)

バッファ内のテキストは並列にトークン化され、その後結合されます。チャンク化されたサンプルは、バッファが空になり、プロセスが再開されるまで出力されます。 infinite=True を設定すると、データセットのイテレータが終了地点から再開します。

def create_dataloaders(args):
    ds_kwargs = {"streaming": True}
    train_data = load_dataset(args.dataset_name_train, split="train", streaming=True)
    train_data = train_data.shuffle(buffer_size=args.shuffle_buffer, seed=args.seed)
    valid_data = load_dataset(args.dataset_name_valid, split="train", streaming=True)
    
    train_dataset = ConstantLengthDataset(tokenizer, train_data, infinite=True, seq_length=args.seq_length)
    valid_dataset = ConstantLengthDataset(tokenizer, valid_data, infinite=False, seq_length=args.seq_length)
    
    train_dataloader = DataLoader(train_dataset, batch_size=args.train_batch_size)
    eval_dataloader = DataLoader(valid_dataset, batch_size=args.valid_batch_size)
    return train_dataloader, eval_dataloader

train_dataloader, eval_dataloader = create_dataloaders(args)

トレーニングを始める前に、オプティマイザと学習率スケジュールを設定する必要があります。バイアスとLayerNormの重みには重み減衰を適用したくないので、それらを除外するためにヘルパー関数を使用しています。

def get_grouped_params(model, args, no_decay=["bias", "LayerNorm.weight"]):
    params_with_wd, params_without_wd = [], []
    for n, p in model.named_parameters():
        if any(nd in n for nd in no_decay): params_without_wd.append(p)
        else: params_with_wd.append(p)
    return [{"params": params_with_wd, "weight_decay": args.weight_decay},
            {"params": params_without_wd, "weight_decay": 0.0},]

optimizer = AdamW(get_grouped_params(model, args), lr=args.learning_rate)
lr_scheduler = get_scheduler(name=args.lr_scheduler_type, optimizer=optimizer,
                             num_warmup_steps=args.num_warmup_steps,
                             num_training_steps=args.max_train_steps,)

残っている大きな質問は、すべてのデータとモデルが複数のGPUに分散される方法です。これは複雑なタスクのように聞こえますが、実際には🤗 Accelerateの1行のコードだけで済みます。

model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader, eval_dataloader)

内部ではDistributedDataParallelが使用されるため、各GPUワーカーにバッチが送られ、モデルの各ワーカーには独自のコピーがあります。そこで勾配が計算され、モデルが各ワーカーで更新されます。

また、定期的に検証セットでモデルを評価する必要もありますので、そのための関数を作成しましょう。これは分散型で自動的に行われ、ワーカーからすべての損失を収集するだけです。また、パープレキシティを報告したいです。

def evaluate(args):
    model.eval()
    losses = []
    for step, batch in enumerate(eval_dataloader):
        with torch.no_grad():
            outputs = model(batch, labels=batch)
        loss = outputs.loss.repeat(args.valid_batch_size)
        losses.append(accelerator.gather(loss))
        if args.max_eval_steps > 0 and step >= args.max_eval_steps:
            break
    loss = torch.mean(torch.cat(losses))
    try:
        perplexity = torch.exp(loss)
    except OverflowError:
        perplexity = float("inf")
    return loss.item(), perplexity.item()

さて、メインのトレーニングループを書く準備が整いました。これは通常のPyTorchのトレーニングループとほとんど同じように見えます。ここでも、ネイティブのPyTorchではなく、アクセラレータの関数を使用していることがわかります。また、評価後にモデルをブランチにプッシュしています。

# モデルをトレーニングする
model.train()
completed_steps = 0
for step, batch in enumerate(train_dataloader, start=1):
    loss = model(batch, labels=batch, use_cache=False).loss
    loss = loss / args.gradient_accumulation_steps
    accelerator.backward(loss)
    if step % args.gradient_accumulation_steps == 0:
        accelerator.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        completed_steps += 1
    if step % args.save_checkpoint_steps == 0:
        eval_loss, perplexity = evaluate(args)
        accelerator.wait_for_everyone()
        unwrapped_model = accelerator.unwrap_model(model)
        unwrapped_model.save_pretrained(args.save_dir, save_function=accelerator.save)
        if accelerator.is_main_process:
            hf_repo.push_to_hub(commit_message=f"step {step}")
        model.train()
    if completed_steps >= args.max_train_steps:
        break

wait_for_everyone()unwrap_model()を呼び出すことで、すべてのワーカーが準備完了であり、prepare()で追加されたモデルレイヤーが削除されることを確認しています。また、簡単に実装できる勾配蓄積と勾配クリッピングも使用しています。最後に、トレーニングが完了した後に最後の評価を実行し、最終モデルを保存してハブにプッシュします。

# 最後のチェックポイントを評価して保存する
logger.info("トレーニング後のモデルの評価と保存")
eval_loss, perplexity = evaluate(args)
log_metrics(step, {"loss/eval": eval_loss, "perplexity": perplexity})
accelerator.wait_for_everyone()
unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(args.save_dir, save_function=accelerator.save)
if accelerator.is_main_process:
    hf_repo.push_to_hub(commit_message="final model")

完了しました!これが150行だけで完全なGPT-2モデルをゼロからトレーニングするためのすべてのコードです。コードを少しコンパクトにするために、スクリプトのインポートとログは表示していません。さて、実際にトレーニングを開始しましょう!

このコードを使用して、TransformersとNLPに関する私たちの新刊書のためのモデルをトレーニングしました。110Mおよび1.5Bのパラメータを持つGPT-2モデルです。これらのモデルをトレーニングするために、16 x A100 GPUマシンを使用しました。それぞれ1日と1週間かかりました。コーヒーを飲んだり、本を読んだりするのに十分な時間です!

評価

事前トレーニングにはまだ比較的短い時間ですが、類似のモデルと比較して良好な下流パフォーマンスが既に観察されます。私たちはOpenAIのHumanEvalベンチマークでこれらのモデルを評価しました。これはCodex論文で紹介されたコーディングの課題に対するコード生成モデルのパフォーマンスを測定します。

注意すべきは、CodeParrotを約25-30Bトークンでトレーニングし、GPT-neoは300Bトークン、Codexは300B(GPT-3チェックポイント)+ 100B(コード)トークンでトレーニングされたことです。それだけの価値があります。

では、このpass@kメトリックは具体的には何でしょうか?簡単に言えば、k個の候補生成から少なくとも1つのプログラムがプログラミングの問題のunittestに合格する確率を測定します。自然に、候補が増えるほど(kが増えるほど)、メトリックは増加します。メトリックはdatasetsと一緒に使用できます:

from datasets import load_metric

code_eval = datasets.load_metric("code_eval")
test_cases = ["assert add(2,3)==5"]
candidates = [["def add(a,b): return a*b", "def add(a, b): return a+b"]]
pass_at_k, results = code_eval.compute(references=test_cases, predictions=candidates, k=[1, 2])

print(pass_at_k)

# -> {'pass@1': 0.5, 'pass@2': 1.0}

datasetsを使用してOpenAIのHumanEvalデータセットも読み込むこともできます:

from datasets import load_dataset

dataset = load_dataset("openai_humaneval")

さて、長い間待ちわびた後に何を得るか見てみましょう!

結果

ゆっくり始めて、モデルにファイルのパスを与えてそのサイズを返す関数を書かせてみましょう:

プロンプト:

def get_files_size(filename):

完成:

    """
    ファイルのサイズを返します。

    :param filepath: ファイルのパス。
    :type filepath: str
    :return: ファイルのサイズ。
    :rtype: int
    """
    return os.path.getsize(filepath)

これまでは順調ですが、もう少し興味深いことを試してみましょう。本番のコードを構築する際のあまり興味のない側面の1つは、unittestを書くことです。CodeParrotを使用してこれを自動化できるかどうか見てみましょう!

プロンプト:

def is_even(value):
    """valueが偶数の場合はTrueを返します。"""
    return value % 2 == 0

# is_evenのためのユニットテストのセットアップ
import unittest

完成:

class TestIsEven(unittest.TestCase):
    """is_even関数のテスト。"""

    def test_is_even(self):
        """is_even関数のテスト。"""
        self.assertTrue(is_even(0))
        self.assertTrue(is_even(1))
        self.assertTrue 

テスト条件自体は正しくないように見えますが、モデルが私たちのためにすべてのボイラープレートコードを作成しました。これだけでもかなり便利で、数分の時間を節約できるかもしれません!

しかし、本当にやりたいのは、より良いトランスフォーマーモデルを構築することです。CodeParrotを使用してそれを実現しましょう:

プロンプト:

from transformers import AutoTokenizer, AutoModelForSequenceClassification

# BERT分類器を構築する

完成:

tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased')

数例を見るのは良いですが、おそらく自分で試したいと思っているでしょう。そのために、いくつかのオンラインデモを用意しました!

デモ

Hugging Face Spacesに追加した2つのデモでモデルを操作することができます。最初のデモではモデルを使用してコードを素早く生成することができ、2番目のデモではモデルを使用してコードをハイライトしてバグを見つけることができます!

  • コードの生成
  • コードのハイライト

また、transformersライブラリから直接モデルを使用することもできます:

from transformers import pipeline

pipe = pipeline('text-generation', model='lvwerra/codeparrot')
pipe('def hello_world():')

概要

この短いブログ記事では、コード生成のための大規模なGPT-2モデルであるCodeParrot 🦜をトレーニングするためのすべてのステップを説明しました。🤗 Accelerateを使用して、200行未満のコードでトレーニングスクリプトを作成し、複数のGPUで簡単にスケーリングすることができました。これにより、自分自身のGPT-2モデルをトレーニングすることができます!

この記事では、CodeParrot 🦜の概要を簡単に説明しましたが、このモデルを事前トレーニングする方法により詳しく掘り下げるには、TransformersとNLPに関する近日発売の書籍での専用の章を読むことをお勧めします。この章では、カスタムデータセットの構築、新しいトークナイザをトレーニングする際の設計上の考慮事項、およびアーキテクチャの選択について詳細が説明されています。

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!

Share:

Was this article helpful?

93 out of 132 found this helpful

Discover more

人工知能

ジョナサン・ダムブロット、Cranium AIのCEO兼共同創設者- インタビューシリーズ

ジョナサン・ダムブロットは、Cranium AIのCEO兼共同創業者ですCranium AIは、サイバーセキュリティおよびデータサイエンスチ...

人工知能

キャルレールの最高製品責任者、ライアン・ジョンソンへのインタビューシリーズ

ライアンは、初期のスタートアップからフォーチュン100の組織まで、多様なテクノロジーと製品開発のリーダーシップ経験を15年...

人工知能

ピーター・マッキー、Sonarの開発者担当責任者-インタビューシリーズ

ピーター・マッキーはSonarのDeveloper Relationsの責任者です Sonarは、悪いコードの1兆ドルの課題を解決するプラットフォー...

人工知能

「クリス・サレンス氏、CentralReachのCEO - インタビューシリーズ」

クリス・サレンズはCentralReachの最高経営責任者であり、同社を率いて、自閉症や関連する障害を持つ人々のために優れたクラ...

機械学習

「Prolificの機械学習エンジニア兼AIコンサルタント、ノラ・ペトロヴァ – インタビューシリーズ」

『Nora Petrovaは、Prolificの機械学習エンジニア兼AIコンサルタントですProlificは2014年に設立され、既にGoogle、スタンフ...

人工知能

「スノーケルAIのCEO兼共同創設者、アレックス・ラットナー - インタビューシリーズ」

アレックス・ラトナーは、スタンフォードAIラボを母体とする会社、Snorkel AIのCEO兼共同創設者ですSnorkel AIは、手作業のAI...