OpenAIの4番目の大規模言語モデル(LLM)であるGPT-4のトレーニングには、推定50ギガワット時を要し、これは5,000世帯のアメリカの家庭の年間電力消費量に相当する。これは2023年のことだ。それ以来、フロンティアLLMのトレーニングに使われる計算リソースはさらに増加しているが、直接的な電力使用量の数字は入手しにくい。

現在、トゥウェンテ大学オランダ)の研究グループは、GPUのクロック周波数を計算中に巧妙に調整することで、LLMのトレーニングで使用されるエネルギーの最大14%を、速度を犠牲にすることなく節約できることを示したJeffrey Spaan氏(トゥウェンテ大学の博士課程候補生で、本論文の筆頭著者)は先月、シチリア島カターニアで開催されたComputing Frontiersカンファレンスでこの結果を発表した。

「私の研究は、コンピューティングの無駄を見つけることです」とSpaan氏は言う。「それはハードウェアの利用不足に似ていますが、ソフトウェアをハードウェアに最適化するのではなく、ハードウェアをソフトウェアに最適化しようとしています。」

GPUのクロックを制御する

Spaan氏と共同研究者たちは、動的電圧・周波数スケーリング(DVFS)と呼ばれる手法を用いてこれを実現した。あらゆるチップ——フロンティアモデルのトレーニングに一般的に使われるGPUを含む——は、計算を制御するために少なくとも1つのクロックを使用する。チップ内の各操作はクロックパルスによってトリガーされる。このクロックが刻む周波数は、チップの動作速度と消費電力を制御する。

現代のGPUには2つのクロックがあり、1つは計算コア用、もう1つはメモリ用である。コアが数値計算に忙しく働いている間は、計算速度を確保するためにクロック周波数を高く保つ。しかし、DVFSを使えば、その間にメモリクロックを遅くして消費電力を抑えることができる。原理的には、チップのメモリ部分を完全にオフにすることも可能だが、GPUの設計上、そのオフスイッチをソフトウェアで制御することはできず、計算の途中で再びオンにするのに時間がかかりすぎる。同様に、コアがメモリからデータを読み込むのを待っている間は、コアクロック周波数を極めて低くし、メモリクロック周波数を上げることもできる。

DVFSは少なくとも1990年代から知られている手法である。しかしSpaan氏によれば、他の研究者はLLMトレーニングに有用に適用できなかった。なぜなら、彼らの手法は計算速度を大幅に低下させたり、エネルギー使用量を改善するのに十分な細かさがなかったりしたからだ。

以前のDVFSの試みでは、トレーニングプロセスの各イテレーションごとに周波数を調整していた。LLMのトレーニングでは、各イテレーションは2つの部分で構成される。すなわち、データが重みを使ってモデルの層を順方向に流れるフォワードパスと、フォワードパスの結果に基づいて重みを層ごとに調整するバックプロパゲーションである。そのため、以前の研究では、フォワードパスでは1つの周波数値を保持し、バックプロパゲーションでは別の値に調整していた。

Spaan氏と共同研究者たちは、より短い時間スケールでクロック周波数を調整した。GPUのワークロードは、カーネルと呼ばれる小さな計算単位に分解される。たとえば、単一のベクトル-ベクトル乗算が1つのカーネルになる場合がある。カーネルはGPUに送られ、並列に何度も処理される。Spaan氏の実装では、ディープニューラルネットワークの単一層の計算が約40個のカーネルに分割される。カーネル単位でクロック周波数を調整することで、チームははるかに大きなエネルギー節約を実現できた。

GPUはまた、チップ内部のシステムが需要の高低を検知すると自動的にDVFSを実行する、とSpaan氏は指摘する。「だから『GPUに任せればいい』と思う人もいるかもしれない」と彼は言う。「しかし、GPUにはどのカーネルが実行されるかという我々の先見性がないため、その場限りの最善の推測でしか動作できず、同じレベルの節約を達成することはできない。」そこで手動調整が必要になる。

エネルギーを減らし、時間を同じに

チームは、13億パラメータのモデルであるGPT-3-XLをNvidia RTX 3080 Ti GPUでトレーニングする実験を行った。時間を節約するため、モデルの単一層のトレーニングに焦点を当てた。この設定で、トレーニング時間をわずか0.6%しか遅らせることなく14%のエネルギー節約を実現する周波数調整の組み合わせを見出した。モデルの性能は計算速度とエネルギー使用量の両方に依存する。

課題が1つある。クロック周波数を下げるのはコアのオン/オフを切り替えるよりもはるかに高速だが、それでも瞬時には行えない。彼らの実験では、研究者たちは周波数切り替え速度を考慮せず、1つのカーネルずつ評価した。そのため、14%のエネルギー節約は最良のシナリオである。実際のところどれほど問題になるかは、使用するGPUに大きく依存するとSpaan氏は述べている。Blackwell GPUのような新しいハードウェアは、古いバージョンよりもはるかに高速な切り替え速度を持ち、完全なエネルギー節約を実現できるはずだ。

現在、チームは特定のワークロードに対して最適な周波数スケーリングを自動的に実装できるツールを開発している。Spaan氏は、この手法が業界のリーダーにとって採用に値するほど魅力的になることを期待している。「我々は性能を損なうことなくエネルギーを節約することを最適化しています」とSpaan氏は言う。「現実の世界では、性能こそが聖杯なのです。」