TRL documentation
トレーニングの高速化
トレーニングの高速化
このガイドでは、TRLでのトレーニングを高速化するさまざまな方法について説明します。各手法には、より詳細なドキュメントへのリンクを伴う最小限の例が含まれています。
オンライン手法における高速な生成のためのvLLM
GRPO や Online DPO などのオンライン手法では、モデルが補完を生成する必要があり、これは多くの場合遅い処理であり、トレーニング時間に大きな影響を与える可能性があります。 生成を高速化するには、PagedAttention などにより高速な生成を可能にするライブラリであるvLLMを使用できます。TRL のオンライントレーナーは vLLM をサポートしており、トレーニング速度を大幅に向上させます。詳細については、vLLM 統合を参照してください。
vLLMを使用するには、まず以下を使用してインストールしてください。
pip install trl[vllm]
まず、…を実行してvLLMサーバーを起動します。
VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
--weight-transfer-config '{"backend": "nccl"}' \
--logprobs-mode processed_logprobs \
--max-logprobs -1次に、学習スクリプトを実行し、学習引数にuse_vllm=Trueを渡します。
from trl.experimental.online_dpo import OnlineDPOConfig
training_args = OnlineDPOConfig(..., use_vllm=True, vllm_mode="server")まず、…を実行してvLLMサーバーを起動します。
VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
--weight-transfer-config '{"backend": "nccl"}' \
--logprobs-mode processed_logprobs \
--max-logprobs -1次に、学習スクリプトを実行し、学習引数にuse_vllm=Trueを渡します。
from trl import GRPOConfig
training_args = GRPOConfig(..., use_vllm=True, vllm_mode="server")追加の引数を渡すことで、サーバー設定をカスタマイズできます。詳細については、vLLM統合を参照してください。
vLLM を使用する際は、リソースの競合を避けるため、学習と生成に割り当てる GPU を分けてください。例えば、学習に 4 つの GPU、vLLM の生成にもう 4 つの GPU を使用する場合、
CUDA_VISIBLE_DEVICESを使用して GPU の割り当てを指定できます。vLLMの生成用にGPU 0-3を設定します:
CUDA_VISIBLE_DEVICES=0,1,2,3 VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \ --weight-transfer-config '{"backend": "nccl"}' \ --logprobs-mode processed_logprobs \ --max-logprobs -1学習にはGPU 4-7 を使用します:
CUDA_VISIBLE_DEVICES=4,5,6,7 accelerate launch train.py
まず、…を実行してvLLMサーバーを起動します。
VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
--weight-transfer-config '{"backend": "nccl"}' \
--logprobs-mode processed_logprobs \
--max-logprobs -1次に、学習スクリプトを実行し、学習引数にuse_vllm=Trueを渡します。
from trl import RLOOConfig
training_args = RLOOConfig(..., use_vllm=True, vllm_mode="server")追加の引数を渡すことで、サーバー設定をカスタマイズできます。詳細については、vLLM統合を参照してください。
vLLM を使用する際は、リソースの競合を避けるため、学習と生成に割り当てる GPU を分けてください。例えば、学習に 4 つの GPU、vLLM の生成にもう 4 つの GPU を使用する場合、
CUDA_VISIBLE_DEVICESを使用して GPU の割り当てを指定できます。vLLMの生成用にGPU 0-3を設定します:
CUDA_VISIBLE_DEVICES=0,1,2,3 VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \ --weight-transfer-config '{"backend": "nccl"}' \ --logprobs-mode processed_logprobs \ --max-logprobs -1学習にはGPU 4-7 を使用します:
CUDA_VISIBLE_DEVICES=4,5,6,7 accelerate launch train.py
最適化されたアテンションの実装
TRLは、メモリ使用量を削減しながらトレーニングを大幅に高速化できる、さまざまな最適化されたアテンション実装をサポートしています。Kernels Hubから事前最適化されたカーネルを直接使用するか、手動で構築したアテンションバックエンドを使用できます。
Hubから事前に最適化されたアテンションカーネルを使用することで、手動でのコンパイルなしで利用できます:
from trl import SFTConfig
training_args = SFTConfig(..., model_init_kwargs={"attn_implementation": "kernels-community/flash-attn2"})その他のオプションには、kernels-community/vllm-flash-attn3やkernels-community/paged-attentionがあります。
最適化されたアテンションは、すべてのTRLトレーナーで動作します。詳細については、Kernels Hub Integrationを参照してください。
最適化されたアテンションバックエンドを手動で構築することは複雑で時間がかかるため、絶対に必要でない限り推奨されません。前節で説明したように、代わりにHubからのKernelsの利用を検討してください。
Flash Attention 2 などの最適化されたアテンションバックエンドを手動でインストールした場合は、トレーニング引数で指定できます:
from trl import SFTConfig
training_args = SFTConfig(..., model_init_kwargs={"attn_implementation": "flash_attention_2"})混合精度トレーニング
bf16またはfp16を使用した混合精度トレーニングにより、モデルの品質への影響を最小限に抑えながら、トレーニングの高速化とメモリ使用量の削減を実現できます。
from trl import SFTConfig
training_args = SFTConfig(..., bf16=True) # or fp16=True for older GPUsAmpere GPU（A100、RTX 30xx）以降の GPU では bf16=True を、それ以前の GPU では fp16=True を使用してください。混合精度トレーニングは、すべての TRL トレーナーでサポートされています。