Hugging Face Trainerによる効率的なFine-tuning: 検証実験を素早く回そう

はじめに

Hugging Faceで公開されているモデルをfine-tuningする際、皆さんはどのようにコードを書いていますか？僕は基本的にフルスクラッチでコードを書いていました。
しかし、LLM（Large Language Models）をチューニングする際に、学習と評価の処理をラップするSFTTrainerを利用し、Hugging FaceのTrainerクラスの便利さを実感しました。
以来、検証段階ではTrainerを積極的に使うようになりました。
今回は、Hugging FaceのTrainerを利用することでどの程度コード量を削減できるのかを具体的に解説します。

Huggingface Trainerでどの程度コードが削減できるか

フルスクラッチでコードを書く際、以下のような学習の雛形を記載が面倒です。また、GPU転送やMixed Precisionの設定など、細かい部分の処理も面倒です。

for epoch in range(num_epochs):
    # 学習
    model.train()
    for train_batch in train_dataset:
        x, y = train_batch
        x = x.to(device)
        y = y.to(device)
        pred = model(x)
        loss = criterion(pred, y)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    # 評価
　　　　　　　　model.eval()
    with torch.no_grad():
        val_loss = []
        for val_batch in val_dataset:
            x, y = train_batch
            x = x.to(device)
            y = y.to(device)
            pred = model(x)
            loss = criterion(pred, y)
            val_loss.append(loss.item())
        val_loss = torch.mean(torch.tensor(val_loss))

しかし、Hugging FaceのTrainerを利用すると、以下のような簡潔な記述だけで学習することができます。

trainer = Trainer(
    model=model,
    args=training_args, # 学習率などのパラメータ設定　TrainingArgumentsを利用する
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)
trainer.train()

雛形の記載やGPU転送などの設定を省略でき、ミスなく簡単に実行することができますね。

さらに、Trainerには複雑な並列処理も実装されています。
例えば、複数のGPUが利用可能な環境では、自動でDataParallelで処理します(下記参照)。より高速なDistributedDataParallelも、TrainingArguments で設定することで、自分でコードを書くよりも簡単に利用できます。

if self.args.n_gpu > 1 and not getattr(model, "is_loaded_in_8bit", False):
    model = nn.DataParallel(model)

https://github.com/huggingface/transformers/blob/266c67b06a3ce19e584af4a809d1aaa598113b2b/src/transformers/trainer.py#L1338

おわりに

Hugging FaceのTrainerを利用することで、コードの記述時間を大幅に削減し、検証実験を素早く行うことができます。特に、マルチGPU環境での自動的な並列処理には、個人的にも感動しました。利用は非常に簡単なので、チュートリアルを参照し、ぜひ一度試してみてください！

Fine-tuning a model with the Trainer API - Hugging Face NLP Course

Trainer