CoMERA: Computing- and Memory-Efficient Training via Rank-Adaptive Tensor Optimization

要約

LLM や DLRM などの大規模な AI モデルのトレーニングには、膨大な GPU とコンピューティング時間がかかります。
高額なトレーニング費用を負担できるのは大手テクノロジー企業のみとなり、その一方で環境への影響についての懸念も高まっています。
この論文では、ランク適応テンソル最適化によるコンピューティングとメモリの効率化を図るトレーニング手法である CoMERA について説明します。
CoMERA は、多目的最適化定式化によってランク適応テンソル圧縮 (事前) トレーニングを実現し、トレーニング プロセスで高い圧縮率と優れた精度の両方を提供するようにトレーニングを改善します。
最適化された数値計算 (最適化されたテンソル化埋め込みやテンソルネットワーク縮小など) と GPU 実装により、GPU でのテンソル化トレーニングにおける実行時のオーバーヘッドの一部が排除されます。
これにより、標準のトレーニングと比較して、トレーニング エポックあたり初めて 2 ～ 3 倍の速度向上が実現します。
CoMERA は、メモリ効率とコンピューティング効率の両方の点でも、最近の GaLore を上回っています。
具体的には、単一バッチ トレーニングでテストされた 6 つのエンコーダー トランスフォーマーにおいて、CoMERA は GaLore よりもトレーニング エポックあたり $2\倍$ 高速で、メモリ効率が $9\倍$ 優れています。
また、私たちの方法は、事前トレーニングで $4.23\times$ の圧縮率を達成しながら、BERT のようなコード生成 LLM での標準の事前トレーニングよりも $\sim 2\times$ の高速化を示しています。
HPC のさらなる最適化により、CoMERA は他の多くの LLM の事前トレーニング コストを削減できる可能性があります。
CoMERA の実装は https://github.com/ziyangjoy/CoMERA で入手できます。

要約(オリジナル)

Training large AI models such as LLMs and DLRMs costs massive GPUs and computing time. The high training cost has become only affordable to big tech companies, meanwhile also causing increasing concerns about the environmental impact. This paper presents CoMERA, a Computing- and Memory-Efficient training method via Rank-Adaptive tensor optimization. CoMERA achieves rank-adaptive tensor-compressed (pre)-training via a multi-objective optimization formulation and improves the training to provide both a high compression ratio and excellent accuracy in the training process. Our optimized numerical computation (e.g., optimized tensorized embedding and tensor-network contractions) and GPU implementation eliminate part of the run-time overhead in the tensorized training on GPU. This leads to, for the first time, $2-3\times$ speedup per training epoch compared with standard training. CoMERA also outperforms the recent GaLore in terms of both memory and computing efficiency. Specifically, CoMERA is $2\times$ faster per training epoch and $9\times$ more memory-efficient than GaLore on a tested six-encoder transformer with single-batch training. Our method also shows $\sim 2\times$ speedup than standard pre-training on a BERT-like code-generation LLM while achieving $4.23\times$ compression ratio in pre-training. With further HPC optimization, CoMERA may reduce the pre-training cost of many other LLMs. An implementation of CoMERA is available at https://github.com/ziyangjoy/CoMERA.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google