重みフォールディング、CUDAストリーム、そしてモデルが逆順で話し始める原因となったバグ — フィリップ・マクラドゥリ
AAI Engineer
Computing/SoftwareInternet Technology
Transcript
00:00:00.
00:00:13皆さんこんにちは。お越しいただきありがとうございます。
00:00:18それでは発表を始めます。今回の発表は
00:00:22私が書いた論文に関するものです。
00:00:27内容は非常にシンプルです。
00:00:33主張も非常に明確です。
00:00:36基本的には2行の数式を使って、
00:00:39TransformerのRMSノルム層を
00:00:42コストを抑え、高速化し、
00:00:45Transformerアーキテクチャ内の
00:00:48レイヤーとして改善するというものです。
00:00:51かつてLayer Normが標準であり、
00:00:54その後RMS Normに置き換えられたのと
00:00:57同様に、この考え方もその流れを汲むものです。
00:01:00これはその延長線上にある研究です。
00:01:04オープンソース界隈の人々と出会う機会があり、
00:01:09この論文を共同執筆しました。
00:01:11共著者のニルス・グラフ（Nils Graf）は
00:01:14このアイデアの生みの親のような人です。
00:01:17作業はその流れで進められました。
00:01:21この内容はarXivで公開されています。
00:01:23ぜひ目を通し、読んで、試してみてください。
00:01:26リポジトリも公開されています。
00:01:29そしてそのコンセプト、つまりアイデアや
00:01:31考え方は、FlashAttentionを例にすると
00:01:36最も説明しやすいでしょう。
00:01:38FlashAttentionが掛け算の実行まで
00:01:40処理を待機させ、メモリ間の通信を制限することで
00:01:45プロセス全体の高速化を図るのと
00:01:48同様の発想に基づいています。
00:01:52これによって特定の改善が行われ、
00:01:54RMS Normの処理が大幅に高速化され、
00:01:55結果としてTransformer全体が改善されます。
00:01:58ここで、「なぜRMS Normなのか？」という
00:02:01疑問が湧くかもしれません。このレイヤーは
00:02:08計算の大部分を占めているわけではないからです。
00:02:15それは事実であり、演算量の観点から見ても
00:02:19その割合はかなり小さいです。
00:02:24しかし、いわゆる実時間（ウォールタイム）の
00:02:26観点では無視できない大きさになります。
00:02:29例えば、推論が実行される際の
00:02:311つのデコードステップにおいて、
00:02:34RMS Normは33回も呼び出されることがあります。
00:02:36もちろんモデルによって異なりますが、
00:02:40論文内では具体的なモデルや
00:02:43テスト方法について記載しています。
00:02:47問題は、これをどう改善し、
00:02:50行列積の待機時間をどのようにして
00:02:52回避できるかということです。
00:02:54これが遅くなる理由は、GPUが計算能力が
00:02:59低いわけではなく、実際の計算以外の
00:03:03周辺処理が苦手だからです。
00:03:06つまり、実際の作業を開始するまでの
00:03:09オーバーヘッドが問題になります。
00:03:15実験の一部で見られたように、
00:03:20処理を33回も起動すると、
00:03:24それだけで多大な時間がかかってしまいます。
00:03:30これを回避するため、各正規化処理を
00:03:32行列積に統合（フューズ）したり、
00:03:35ウェイト・フォールディングを行って
00:03:40メモリ間のデータ移動を削減することができます。
00:03:44データ移動もGPUにとっては
00:03:46ボトルネックとなる処理です。
00:03:50また、RMS Norm層で行われる
00:03:54除算処理を後回しにすること（遅延）も、
00:03:56この待機ステップを避ける有効な手段です。
00:04:00したがって、この論文の核心は、
00:04:02RMS Normの計算方法に
00:04:06いくつかの代数的な工夫を施すことで、
00:04:09これら3つの側面すべてを改善するという点にあります。
00:04:13内容としてはそれだけです。
00:04:16数学的な裏付けとしては、
00:04:19主に最初の2つの命題に基づいています。
00:04:211つ目は重みなし正規化（Weightless Normalization）です。
00:04:25こちらに示されている通りです。
00:04:29そして2つ目が正規化の遅延（Deferred Normalization）です。
00:04:31より新しいアーキテクチャでは、
00:04:33RMS Normが2回現れる状況が存在します。
00:04:35例えば、Gemma 4などがそれに該当します。
00:04:37そのため、事前正規化を相殺することも有効です。
00:04:39これらはいずれも論文内で代数的に証明されています。
00:04:44最初の命題が表しているのは、
00:04:48ゲインと重みを1つの行列に
00:04:52折りたたむ（フォールドする）という点です。
00:04:58ここでアスタリスク付きで示されているのが
00:05:00その行列Wです。
00:05:04これはオフラインで計算されます。
00:05:06FlashAttentionが事前に一部の処理を
00:05:09計算しておくことで、
00:05:12頻繁なメモリ間通信を避けるのと
00:05:14同じようなアプローチです。
00:05:16このように、重みの折りたたみという
00:05:18ステップが事前に実行されます。
00:05:20もう一つのステップは、
00:05:24行列積におけるスカラー除算を遅延させ、
00:05:26並列実行できるようにすることです。
00:05:32通常の場合、一度計算を行ってから
00:05:35待機し、再び計算を行う必要があります。
00:05:38しかしこの方法では、
00:05:41処理を分割して並列化できるように工夫しています。
00:05:44そして3つ目のポイントとして、
00:05:48スケール不変性があるため、
00:05:51重複する部分の一方を削っても
00:05:56問題なく機能するという特徴があります。
00:05:58これは、同様のアーキテクチャを採用した
00:06:01新しいモデルに適用可能です。
00:06:04これを実際に、特に2つ目の命題を
00:06:10実用化するにはどうすればよいでしょうか。
00:06:13例えば、「transformer tricks」という
00:06:16リポジトリを使えば簡単に適用できます。
00:06:20どんなモデルにでも適用でき、動作します。
00:06:22ただし、本格的に実装するには
00:06:25CUDAカーネルレベルの作業が必要になるため、
00:06:29それほど単純ではありません。
00:06:31私自身も実際に実装を進める中で、
00:06:35ある実験結果に直面しました。
00:06:42プロンプトとして「The transformer architecture revolutionized NLP because」
00:06:46（TransformerアーキテクチャがNLPに革命をもたらしたのは……）
00:06:48を入力し、ある程度期待される出力を
00:06:51得ようとしたときのことです。
00:06:54しかし得られた出力には、
00:06:56単語の繰り返しや1ステップの遅延が見られました。
00:06:59ご覧の通り、
00:07:01「because」という単語が再び現れています。
00:07:07GPUのストリーム周辺で何かが起きており、
00:07:10何が原因なのか突き止めようと試みました。
00:07:15その1ステップの遅延によって、
00:07:18過去の出力が混ざり込むような現象が起きていました。
00:07:21デバッグを重ねる中で気づいたのですが、
00:07:26先ほど説明したような最適化や演算の遅延を
00:07:32CUDAで実装する際、主に2つの機能を利用できます。
00:07:351つは行列積の一部を担当するテンソルコア、
00:07:39もう1つは要素ごとの演算やリダクション、
00:07:44平方根などを実行するCUDAコアです。
00:07:47これらを並列実行し、論文で説明したメリットを活かして
00:07:52コンセプトを実際に検証することが狙いでした。
00:07:58このように配置する予定でした。
00:08:00順次処理を行う場合
00:08:03ベクトルユニットがRMSとスケーリングを計算している間に、アイドルの待機時間が発生します。
00:08:10その後に行列積が行われます。
00:08:12そこで論文の手法であるフラッシュノームを使えば
00:08:16これらを並列で実行できるはずです。
00:08:19行列ユニットが積算を計算し、ベクトルユニットがRMSを計算します。
00:08:23このようにして時間を節約できます。
00:08:26しかし、これをPythonだけで行うことはできません。
00:08:28もう少し低レイヤーを触る必要があります。
00:08:30そして、このようなCUDAコードでそれを実現しました。
00:08:35当時は、全体的にこれでうまくいっているように見えました。
00:08:42しかし、少し間違ったことをしていることに気づきました。
00:08:47それは、2つのストリームを結合するはずの最後の結合が、私の場合では暗黙的になっていたということです。
00:08:57これをテストしたところ、単体テストは機能しました。
00:09:01パープレキシティテストなどの品質も似たようなものでしたが、それは単に似たような生成が行われるためです。
00:09:08しかし、長時間の生成を行うと、この問題が見えてきました。
00:09:11何が起きているのか全く分かりませんでした。
00:09:14その理由は、暗黙的な結合を行っていたため、基本的に一方のストリームの処理が完了していなかったからです。
00:09:24その結果、未完了の行列積から過去のデータを読み込んでしまうような競合状態が発生しました。
00:09:31これを修正するためのアイデアは、結合を明示的に行う必要があるという点に行き着きました。
00:09:40そして、結合時に過去のデータを読み込んでいないことが確実になるまで、一方の処理の完了を待つようにしました。
00:09:48CUDAストリームの探索における気づきは以上のようなものです。
00:09:54そして、このように処理を実装していました。
00:09:57結合が暗黙的になっていたのです。
00:10:00そのため、スケール後の処理で古いバッファ値を読み込んでしまいました。
00:10:06これを修正する方法は、基本的に行列積の終了をマークするというものです。
00:10:14次にRMSの終了をマークします。
00:10:17そしてスケール後の処理で最初のストリームを待ちます。
00:10:21さらに2番目のストリームを待ちます。
00:10:24それによりバグが修正され、論文の内容が正しく機能し、モデルが逆向きではなく前向きに出力できるようになりました。
00:10:33それがおそらく、学術的な観点からの面白い部分でした。
00:10:38でも、実際に試してみたいですよね。
00:10:40デプロイしてテストし、よりプロダクションに近い環境でどのように機能させられるかを確認したいと思いました。
00:10:47論文を読んですべてのテストを確認することもできます。
00:10:50Llamaモデルを中心に検証されていますが、これは他のアーキテクチャでも同様に機能します。
00:10:56この特定の論文に対してできることは、例えば、私が説明したウェイトフォールディングや命題1を、リポジトリ内のコードを使って実行することです。
00:11:10flashifyのように呼び出すだけでそれが実行されます。
00:11:13ただし、2つ目に挙げた項目に関しては、それをやりたい場合、カーネルの作業を少し行う必要があります。
00:11:19私が例で説明した通りです。
00:11:22これらはLlamaモデルに基づいた結果であり、他にも確認できるさまざまな詳細があります。
00:11:29遅延正規化のみを行った場合どうなるか、完全に融合したカーネルを使用した場合どうなるかなどです。
00:11:36そのため、すべての命題をテストするために、より低レイヤーでの多くの実験が行われています。
00:11:41これらが、さまざまなレベルの検証と詳細における私たちの結果です。
00:11:48しかし、ウェイトフォールディングのようなシンプルなものでも、一定の改善が見られます。
00:11:54これは、モデルで使用する日常的なツールとも連携します。
00:11:59そのため、車輪の再発明をしたり、ゼロから何かを作ったりする必要はありません。
00:12:05新しいチェックポイントのようなものなので、torch compileとも問題なく連携します。
00:12:13FlashAttentionも、異なるレイヤーで同様の工夫を行っています。
00:12:16また、量子化モデルとも互換性があります。
00:12:18そのため、これを実際に適用して、新しい正規化レイヤーを備えたモデルを手に入れるのは素晴らしいことです。
00:12:27実際に実行するための詳細やコードは、このtransformer-tricksリポジトリから入手できます。
00:12:34私が説明したさまざまな代数的テクニックや、言及した論文が含まれています。
00:12:41また、GitHubではありませんが、いくつかのモデルでこれを実装したHugging Faceのモデルリポジトリもあります。
00:12:50そのモデルへのHugging Faceリンクから試すことができます。
00:12:54これらのHugging Faceモデルを使って、本番環境にデプロイすることも可能です。
00:13:00これを考えていたとき、科学的な検証が終わり、Hugging Faceモデルへのリンクが用意できたら
00:13:11Superlinkの推論エンジンは、あらゆるHugging Faceモデルをデプロイする優れた方法であることに気づきました。
00:13:19ハッカソンでもこれを行っており、参加者がファインチューニング済みのカスタムなHugging Faceモデルを持ち寄りました。
00:13:27モデルを改善して自身の研究アイデアをテストするために、こうした代数的テクニックのバージョンがある場合でも
00:13:37実際に試してクラスター上にデプロイし、デプロイに関する複雑なグルーコードを心配する必要がありません。
00:13:48それは非常に便利です。ポイントは、クラスター全体とモデル推論がオープンソースであれば
00:13:58カーネルの操作やフラッシュノームのような、より斬新な研究アイデアを実際にテストできるということです。
00:14:12推論を自分で所有していないレンタルされたエンドポイントでは、これを行うのがはるかに難しくなります。
00:14:18そのため、こうした作業を行える、移植性と柔軟性に優れたものが必要です。
00:14:23同時に、大規模なテストも行えるほど実用的な準備が整っているものが望ましいです。
00:14:27例えば、siteを使ってこれを他のモデルと組み合わせることもできます。
00:14:33左上に示されているように、必要に応じてフラッシュ化されたモデルを他のさまざまなモデルと組み合わせてエージェント的なタスクを実行し
00:14:43エンドツーエンドの大規模なユースケースを実現できます。
00:14:46siteの仕組みとして、この本番用クラスターがモデルのデプロイをサポートしてくれます。
00:14:52詳細については、siteのリポジトリも参考にしてください。
00:14:57また、特に小規模なモデルを扱う場合に役立つ、よりスマートなキューイング機構もあります。
00:15:03フラッシュノームの作業を行っていた際、小さめのLlamaモデルやHugging Faceの小さなエージェントを使用していました。
00:15:11GPUコストに多額の費用をかけずに済むよう、小さなモデルを同じGPU上で動作させ
00:15:24必要に応じて小規模なモデルを切り替える仕組みがあるのは非常に便利でした。
00:15:30また、APIやクラスターを通じてモデルの設定を制御することもでき
00:15:35オープンソースの研究をサポートしてくれるインフラ担当者がいなくても、非常に手軽に行えます。
00:15:40そこも素晴らしい点です。
00:15:43さらに自身のクラウドを所有できるため、オープンウェイト、オープンモデル、オープンソースを重視する場合に役立ちます。
00:15:50また、siteには私が言及したものだけでなく、さまざまなモデルのカタログが用意されています。
00:15:57ぜひ確認してみてください。
00:15:59そのようなシステムを構築している場合は、再ランキング用の埋め込みモデルもあります。
00:16:03これで、トランスフォーマーを改善する技術に関する論文を共著し、私の研究の道のりの話は一区切りとなります。
00:16:15同時に、これを本番環境に持ち込み、テストし、オープンソースモデルで自由に遊ぶ方法を見つけることができました。
00:16:24質問やコントリビューションがあれば、LinkedInでお気軽にご連絡ください。
00:16:30ここで言及した内容の多くは、VLLMやHugging FaceへのPRなどに含まれています。
00:16:36あちこちで見つかるでしょう。
00:16:38論文もぜひチェックしてみてください。
00:16:40そこにarXivのリンクが記載されています。
00:16:43また、siteのリポジトリや私のLinkedInもあります。
00:16:47ご参加いただき誠にありがとうございました。
00:16:58質問があれば声をかけてください。
00:16:59ここにいますので。
00:17:00すぐ近くにいます。
00:17:13（音声なし）
Community Posts
No posts yet. Be the first to write about this video!
Write about this video