ProTrain: Efficient LLM Training via Memory-Aware Techniques

要約

大規模言語モデル (LLM) をトレーニングするには、非常にメモリを消費します。
この問題を解決するために、既存の研究では、ZeRO-Offload など、トレーニング プロセスに CPU と GPU の組み合わせを利用しています。
このような技術により、数十億規模のモデル トレーニングが大幅に民主化され、少数の消費者向けグラフィック カードでトレーニングできるようになります。
ただし、私たちの観察によると、既存のフレームワークは多くの場合、粗粒度のメモリ管理を提供しており、構成調整には経験豊富な専門家が必要であり、ハードウェアの使用率とパフォーマンスが最適化されていません。
この論文では、メモリ、計算、IO を調整することでメモリ使用量とパフォーマンスのバランスをインテリジェントに調整する新しいトレーニング システムである ProTrain を提案します。
ProTrain は、チャンクベースのモデル状態管理とブロック単位のアクティベーション管理を通じて、ユーザーの介入なしでメモリ認識ランタイム プロファイラーによってガイドされる適応型メモリ管理を実現します。
ProTrain はトレーニング アルゴリズムを変更しないため、精度が損なわれることはありません。
実験によると、ProTrain は SOTA トレーニング システムと比較してトレーニング スループットを 1.43$\times$ から 2.71$\times$ 向上させます。

要約(オリジナル)

It is extremely memory-hungry to train Large Language Models (LLM). To solve this problem, existing work exploits the combination of CPU and GPU for the training process, such as ZeRO-Offload. Such a technique largely democratizes billion-scale model training, making it possible to train with few consumer graphics cards. However, based on our observation, existing frameworks often provide coarse-grained memory management and require experienced experts in configuration tuning, leading to suboptimal hardware utilization and performance. This paper proposes ProTrain, a novel training system that intelligently balances memory usage and performance by coordinating memory, computation, and IO. ProTrain achieves adaptive memory management through Chunk-Based Model State Management and Block-Wise Activation Management, guided by a Memory-Aware Runtime Profiler without user intervention. ProTrain does not change the training algorithm and thus does not compromise accuracy. Experiments show that ProTrain improves training throughput by 1.43$\times$ to 2.71$\times$ compared to the SOTA training systems.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google