みなさんこんにちは!
背景、動機、たくさんの雑談なので、TL;DRまで飛ばしても大丈夫です。
少し前に、ここで[D] スケールしたRLのポストトレーニングに、どのフレームワークを使っていますか?と質問しました。それ以降、verlに取り組んでいて、仕事でも自分の時間でも使っています。
最初は、何か新しいものを作ろうとしていたわけではありませんでした。主にveRLをきちんと理解し、より良い体験をすることが目的でした。最初に、パッケージングをより現代的にし、`pyproject.toml`を使って、簡単にインストールできるようにし、使っていない依存関係を削除しました。それから、特にvllmとsglangが時々競合するので、適切な互換性マトリクスを見つけました。さらに、異なるrequirementsファイルに入っていた推移的(transitive)依存関係を削除するなどもしました。次に、コードベースから、興味のないコードを全部取り除きたくなりました。HF/Nvidiaに関連するもの(ロールアウト用のtransformers、trlのコード、ロールアウト用のtrtllm、megatronなど)に関するものは、どちらにしても非効率だったり、理解できなかったり、関心がなかったりしたので削除しました。ですが、自分がやっていることが正しいかを確認する手段が必要でした。テストが適切に行われていなかったので、pytestファイルではなくbashファイルが大量にありました。そこで、CPUで動かせるテストと、GPUが必要で自分のノートPCで直接実行したいテストを分ける必要がありました。そして「自分の」GPU(まあプロバイダ側のGPUですが)を最大限活用するためのスケジューラを書き、bashテストをきちんとしたテストファイルに変えました。さらに、フィクスチャを作り、Rayの後片付けを正しく扱って、テスト間でコンテキストが漏れないようにしました。
しかし、取り組んでいくうちに、さらに問題が見つかってもっと良くしたくなりました。すると、verlの中核は、そのオーケストレーション層と、シングルコントローラ(単一コントローラ)という設計パターンだと気づきました。そして、私の見解では、それはひどく書かれています。多くのメタプログラミング(それ自体には反対ではないのですが、うまく扱われていないと思います)、追跡しにくくする迂遠さや魔法のような仕掛けがあって、「実際に何が起きているのか」を辿りにくい。さらに、分散フレームワークでは特に、イミュータブルであることや明快さがたくさん欲しくなるはずです。
そこで、彼らのオーケストレーション層をリファクタしようと思いました。でも、それには明確な頭のモデルが必要でした。たとえば、何が気になっている点を直して、反復的に良くしていくのか、という下書きのようなものです。それで、LLMのポストトレーニングワークロード向けの、オーケストレーションの自己完結型モジュールを持つに至りました。ところが仕上げてみると、私のforkしたverlは200〜300コミットどころか、それ以上、遅れていました。
それに加えて、世の中の人々が気にしていないことにも気づきました。そもそも、どのフレームワークを使うかすら気にしていない。仮に一部が良いか悪いかなんて、ましてやオーケストレーション層なんて、論外です。結局のところ、こうしたフレームワークはML研究者向けで、アルゴリズムの正しさを重視します。GPUの利用効率や、良いMFUが出ているかなどを気にする人もいるかもしれませんが、そういうのは少数派です。そして私は、人々が単に最新のモデルと高い労力でClaudeコードやCodexをフレームワークに投げて、実験を動かしてもらうようなことをしているのを見ました。私はそれを責めるつもりはありません。ただ、そうした気づきがあって、私は考えてしまいました。自分はここで何をしているんだろう? ふふふ。
そして、u/dhruvnigam93が、この道のりをドキュメント化するよう私に提案してくれたことを思い出しました。ブログ記事にするなら価値があるかもしれないと思いましたが、主にコードの仕事についてどうやってブログ記事を書くのか、どうやって課題を説明すればいいのか。とはいえ抽象的なままで終わってしまいます。何がうまくいって何がうまくいかないか、どんなエッジケースが難しいのかを示すにはコードを動かす必要があります。では、自分のコードベースを作る過程で頭の中を巡ってきたすべてを、なぜそうしたのかという理由も含めて、どうやってブログ記事にすればいいのか。そもそも私はブログを書き慣れていません。少しは書きますが、自分のためにやることが多くて、文章はひどいです。
そこで、動画にしてみたら面白いかもしれないと思いました。さらに、コードベースをもう一度見直して、考え直せるのも良い点です。しかも、次の動画を作ろうとしているときにふと疑問が湧きました。「どうすれば、最も効率よくDPの異なるシャードに、データのバッチを分配/分割できるのか?」単にバッチ次元で単純に分割するだけではダメです。あるDPシャードには長いシーケンスが来て、他は短いかもしれません。なのでシーケンス長を考慮する必要があります。なぜ最初にそれを考えなかったのか分かりませんが、とにかく今それを実装しようとしています。幸いにも最初の段階で、コードベース内の異なるシステムに対してどこに境界を置くかをうまく設計して、修正が(ほぼ)しやすいようにしていました。とにかく、以上です。
最初の2本の動画は公開済みです。1本目は「RLポストトレーニングにおけるオーケストレーション問題」で、概念編です。PPOのパイプラインを辿り、モデルの役割をハードウェアに対応づけ、シングルコントローラのパターンを説明します。2本目は「Rayの基礎、ワーカー、GPU配置」と名付けました。こちらは実践編です。まず基本のRayタスク/アクターから始め、ワーカーレイヤを組み上げます。具体的には、ワーカーのアイデンティティ、メッシュレジストリ、そして保証された同居(co-location)のための配置グループです。
次に取り組んでいるのはディスパッチ層です。ディスパッチの原子的な単位は何にすべきか、トークンを意識(token-aware)した形にするにはどうすればいいか、DPシャード間でどう仕事を分割するのか、さらに、異なるローカル実行戦略を使っていてもワーカーが返すべき標準的な結果形式は何か。そしてドライバがそれをどうマージして、きれいな表現に戻すのか。ほとんどは終わっているのですが、トークンを意識した部分は2本目を作っているときに初めて思いついて、一部(主に、ワーカグループからデータを収集する際の、いくつか組み込まれた前提)を再考することを強いられました。
シリーズを始めた理由の背景や動機は以上です。補足として、私が言及した「コードベース」avridについてですが、シリーズの終わりにPyPIへ公開できたらと思っています。というのも、これはよりモジュールで、現時点ではほとんど何も入っていません。せいぜい3つのデータクラス程度です。動画に対してgitの履歴が忠実であるようにしたいからです。ただ、誰かがそれを調べたいのであれば、プライベートリポジトリに招待できます。
注:シングルコントローラのパターンは、数あるパターンのうちの1つにすぎません。私は、世の中にあるあらゆるポストトレーニングのコードベースについて深い知識があるわけでもありませんし、面白い/エレガントである必要すらないと思っています。OpenRLHFや、Ai2のopen-intsructが、うまく動かすために手作りで何かを用意して、それを同梱して出荷しているだけ、という可能性もあるわけです。オーケストレーションを本気で重視している別のコードベースとしてはMonarch/torchforgeがありますが、それについてはコメントするほどの経験がありません。
まず明確にしておくと、これは「verlがダメだった、直した」みたいな投稿ではありません。verlは難しい問題を解決し、効率的で、ちゃんと動きます。そして多くの人がそれをうまく使っていて、私たちもその一人です。彼らはNPUに対応していて、バックエンドがたくさんあり、ロールアウトエンジンやアルゴリズムも揃っていて、nvfp4のqatまであります。これほど速く提供できるのは正直クレイジーです。彼らはとにかく素晴らしい仕事をしていて、私は深く敬意を払っています。そして、私がたくさん学べたのも彼らのおかげです。私は単に、それをより良い形で実装してもっと学びたいと思っているだけで、ただのランダムなエンジニアです。あと、私はすべてを知っているとも、私の実装が最良だとも主張していません。私は、このシリーズ/コードベースを、ポストトレーニングLLM向けの本物のプロダクション対応コードベースに育てていけるよう頑張りたいと思っています。いつかは他のすべての人たちと競り合うことができるかもしれません。こういう質問(たとえば、いつ・なぜインフラがアイドル状態なのか、そこに対して何ができるのか、バブルをどう減らすのか等)が大好きなので、これからも掘り下げていきます。とはいえ、私はただのランダムなエンジニアなので、何か批評やより良いアイデア、成長や学びにつながって、もっと良くなるために役立つことがあれば、どんなことでも聞かせてください!
最終追記:もちろん、スパムにならないように、私がアップロードするすべての動画については投稿しません。代わりにRedditアカウントでやります。
最終最終追記(誓って):動画に広告を入れてはいけないはずなんですが、入れていない認識でいます。そうでないなら教えてください。私はGoogleアカウントを連携して動画をアップロードしたので、大丈夫だと思います。そしてもし見るなら、ぜひx2で見てください(笑)
TL;DR:
私はverlをたくさん使ってきて、その理解を深めようとしているうちに、特にシングルコントローラーパターンにフォーカスすることになりました。このパターンはすごく気に入っていますが、実装が考えづらくて、学習プロジェクトとして、よりクリーンで明示的な形でその部分を作り直し始めました。そうして動画シリーズになりました。



