Huggingface TransformersとRayを使用した検索増強生成

Using Huggingface Transformers and Ray for Enhanced Search Generation.

アノスケールのチームからのゲストブログ投稿：Amog Kamsetty氏

Huggingface Transformersは最近、Retrieval Augmented Generation（RAG）モデルを追加しました。これは、外部のドキュメント（ウィキペディアなど）を活用して知識を拡充し、知識集約的なタスクで最先端の結果を実現する新しいNLPアーキテクチャです。このブログ投稿では、スケーラブルなアプリケーションを構築するためのライブラリであるRayをRAGの文脈におけるドキュメント検索メカニズムに統合する方法を紹介します。これにより、検索呼び出しの速度が2倍に向上し、RAGの分散ファインチューニングのスケーラビリティが向上します。

Retrieval Augmented Generation（RAG）とは何ですか？

RAGの概要です。モデルは外部のデータセットから文脈ドキュメントを取得し、実行の一環として使用します。これらの文脈ドキュメントは元の入力と組み合わせて出力を生成するために使用されます。このGIFはFacebookのオリジナルのブログ投稿から取得されました。

最近、HuggingfaceはFacebook AIと提携して、RAGモデルをTransformersライブラリの一部として導入しました。

• シンプルな人々が派手なニューラルネットワークを構築するための簡単な考慮事項
• ハギングフェイスの読書会、2021年2月 – Long-range Transformers
• Hugging Faceを使用してWav2Vec2を英語音声認識のために微調整する

RAGは他のseq2seqモデルと同様に機能しますが、外部の知識ベース（ウィキペディアのテキストコーパスなど）から文脈ドキュメントを取得する中間コンポーネントを持っています。これらのドキュメントは入力シーケンスと組み合わせて基礎となるseq2seqジェネレータに渡されます。

この情報検索ステップにより、RAGはモデルパラメータに埋め込まれた知識と文脈のパッセージに含まれる情報という複数の知識源を活用することができます。これにより、質問応答などのタスクで他の最先端モデルを上回るパフォーマンスを発揮します。Huggingfaceが提供するデモを使用して、自分自身で試すこともできます。

ファインチューニングのスケーリング

これらの文脈ドキュメントの取得は、RAGの最先端の結果にとって重要ですが、追加の複雑さをもたらします。データ並列トレーニングルーチンを介してトレーニングプロセスをスケーリングアップする際、ドキュメントの検索の単純な実装はトレーニングのボトルネックになることがあります。さらに、検索コンポーネントで使用されるドキュメントインデックスは非常に大きいため、各トレーニングワーカーが自分自身の複製されたインデックスを読み込むことは不可能です。

以前のRAGファインチューニングの実装では、torch.distributed通信パッケージを使用してドキュメント検索部分を活用していました。しかし、この実装は柔軟性に欠け、スケーラビリティに制約がありました。

その代わりに、フレームワークに依存しないアドホックな並行プログラミングのためのより柔軟な実装が必要です。それには、Rayが完璧に適しています。Rayは一般的な分散および並列プログラミングのためのシンプルで強力なPythonライブラリです。Rayを使用して分散ドキュメント検索を行うことで、torch.distributedに比べて検索呼び出しごとに2倍の高速化と、ファインチューニングのスケーラビリティの向上を実現しました。

ドキュメント検索のためのRay

torch.distributedの実装によるドキュメント検索

torch.distributedを使用したドキュメント検索の主な欠点は、トレーニングに使用されるプロセスグループに依存していて、ランク0のトレーニングワーカーのみがインデックスをメモリに読み込んでいたことです。

その結果、この実装にはいくつかの制限がありました：

1. 同期のボトルネック：ランク0のワーカーはすべてのワーカーから入力を受け取り、インデックスクエリを実行し、その結果を他のワーカーに送信する必要がありました。これにより、複数のトレーニングワーカーでのパフォーマンスが制限されました。
2. PyTorch固有の：ドキュメント検索プロセスグループはトレーニングに使用される既存のプロセスグループに依存する必要があり、トレーニングにはPyTorchを使用する必要がありました。

Rayの実装によるドキュメント検索

これらの制限を克服するために、Rayに基づく分散検索の革新的な実装を導入しました。Rayの状態を持つアクター抽象化を使用して、トレーニングプロセスとは別の複数のプロセスを使用してインデックスを読み込み、検索クエリを処理します。複数のRayアクターを使用することで、検索はもはやボトルネックではなくなり、RAGにはPyTorchの要件もありません。

以下の結果からもわかるように、Rayベースの実装を使用すると、マルチGPUファインチューニングの検索パフォーマンスが向上します。以下の結果は、リトリーバル呼び出しごとの秒数を示しており、GPUの数を増やすとRayを使用した方がtorch.distributedよりもパフォーマンスが向上します。また、リトリーバルを実行するRayプロセスの数を増やすと、単一のリトリーバルプロセスがボトルネックではなくなるため、より多くのトレーニングワーカーでもパフォーマンスが向上します。

	2つのGPU	3つのGPU	4つのGPU
torch.distributed	2.12秒/検索	2.62秒/検索	3.438秒/検索
Ray 2つの検索プロセス	1.49秒/検索	1.539秒/検索	2.029秒/検索
Ray 4つの検索プロセス	1.145秒/検索	1.484秒/検索	1.66秒/検索

異なる検索実装のパフォーマンス比較です。各ドキュメント検索実装について、パフォーマンスを向上させるために、8つのGPUバッチサイズで500回のトレーニングステップを実行し、ランク0のトレーニングワーカーで各バッチのコンテキストドキュメントの取得にかかる時間を計測します。結果からわかるように、複数の検索プロセスを使用することでパフォーマンスが向上し、特に複数のGPUでのトレーニングをスケールさせる場合に効果的です。

どのように使いますか？

Huggingfaceは、PyTorch Lightningベースのファインチューニングスクリプトを提供しており、Rayの検索実装をオプションとして追加しました。

試すために、まず必要な要件をインストールしてください。

pip install ray
pip install transformers
pip install -r transformers/examples/research_projects/rag/requirements.txt

次に、データパスやその他の設定を指定してfinetune-rag-ray.shを実行できます。

# 分散検索のためにRayを使用してRAGをファインチューニングするサンプルスクリプト。

# lightning_base.pyにアクセスするために親ディレクトリをPythonパスに追加します。
export PYTHONPATH="../":"${PYTHONPATH}"

# シングルノードのRayクラスタを開始します。
ray start --head

# サンプルのファインチューニング実行です。data_dir、output_dir、model_name_or_pathを指定する必要があります。
# 全ての可能なオプションを表示するには、./examples/rag/finetune_rag_ray.sh --help を実行してください。

python examples/rag/finetune_rag.py \
    --data_dir $DATA_DIR \
    --output_dir $OUTPUT_DIR \
    --model_name_or_path $MODEL_NAME_OR_PATH \
    --model_type rag_sequence \
    --fp16 \
    --gpus 8 \
    --profile \
    --do_train \
    --do_predict \
    --n_val -1 \
    --train_batch_size 8 \
    --eval_batch_size 1 \
    --max_source_length 128 \
    --max_target_length 25 \
    --val_max_target_length 25 \
    --test_max_target_length 25 \
    --label_smoothing 0.1 \
    --dropout 0.1 \
    --attention_dropout 0.1 \
    --weight_decay 0.001 \
    --adam_epsilon 1e-08 \
    --max_grad_norm 0.1 \
    --lr_scheduler polynomial \
    --learning_rate 3e-05 \
    --num_train_epochs 100 \
    --warmup_steps 500 \
    --gradient_accumulation_steps 1 \
    --distributed_retriever ray \
    --num_retrieval_workers 4

# Rayクラスタを停止します。
ray stop

次は何ですか？

Huggingface transformersとRayの検索実装を使用して、より高速な分散ファインチューニングのためのRAGを活用することで、自身の知識集約型タスクで検索ベースの生成を行うことができます。

また、ハイパーパラメータのチューニングもトランスフォーマーファインチューニングの別の側面であり、精度に大きな影響を与えることがあります。スケーラブルで簡単なハイパーパラメータのチューニングには、Ray Tuneライブラリをご覧ください。Ray TuneのPyTorch Lightningとの統合、またはHuggingface transformersとの組み込み統合を使用することで、RAGモデルのパーフェクトなハイパーパラメータを見つけるための実験を実行できます。

最後に、HuggingfaceでのRAGのTensorflow実装もお楽しみに！

RAG+Rayの統合を試す予定がある場合は、Ray Discourseでの経験をお気軽に共有していただくか、RayコミュニティのSlackに参加してさらなる議論に参加していただければ幸いです！

また、https://medium.com/distributed-computing-with-ray/retrieval-augmented-generation-with-huggingface-transformers-and-ray-b09b56161b1e でも公開されています。

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