Optimizing RLHF Training for Large Language Models with Stage Fusion

要約

人間のフィードバック（RLHF）からの補強学習のための段階的融合を備えた効率的なトレーニングシステムであるRLHFUSEを提示します。
RLHFトレーニングの本質的な性質、つまり、生成段階のデータ歪度とトレーニング段階のパイプラインバブルにより、既存のRLHFシステムはGPUの使用率が低いことに苦しんでいます。
RLHFUSEは、RLHFワークフローの従来のビューを個々のタスクの構成として分割し、各タスクをより細かいサブタスクに分割し、GPU使用率を改善するためのステージ融合を実行します。
rlhfuseには2つの重要なアイデアが含まれています。
第一に、生成および推論のタスクについて、RLHFUSEはそれらをサンプルレベルのサブタスクに分割し、効率的なステージ間融合が生成段階と推論段階の実行と重複することを可能にし、したがって、長期尾のあるサンプルが支配する元の生成ボトルネックを軽減します。
第二に、トレーニングタスクの場合、RLHFUSEはそれらをマイクロバッチのサブタスクに分割し、ステージ内融合を実行して、融合パイプラインスケジュールでトレーニング段階でこれらのサブタスクを同時に実行し、パイプラインのバブルを効果的に軽減します。
この実験では、RLHFUSEが既存のシステムと比較して、トレーニングスループットが最大3.7 \ Times $を増加させることを示しています。

要約(オリジナル)

We present RLHFuse, an efficient training system with stage fusion for Reinforcement Learning from Human Feedback (RLHF). Due to the intrinsic nature of RLHF training, i.e., the data skewness in the generation stage and the pipeline bubbles in the training stage, existing RLHF systems suffer from low GPU utilization. RLHFuse breaks the traditional view of RLHF workflow as a composition of individual tasks, splitting each task into finer-grained subtasks, and performing stage fusion to improve GPU utilization. RLHFuse contains two key ideas. First, for generation and inference tasks, RLHFuse splits them into sample-level subtasks, enabling efficient inter-stage fusion to overlap the execution of generation and inference stages, thus mitigating the original generation bottleneck dominated by long-tailed samples. Second, for training tasks, RLHFuse breaks them into subtasks of micro-batches and performs intra-stage fusion to concurrently execute these subtasks in the training stage with a fused pipeline schedule, effectively mitigating the pipeline bubbles. The experiments show that RLHFuse increases the training throughput by up to $3.7\times$, compared to existing systems.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google