見出し画像

大規模分散学習でGPUを食べまくる人

マソラ

こんにちは！東工大修士2年/産総研RAの高島空良と申します。本記事は、研究コミュニティ cvpaper.challenge 〜CV分野の今を映し,トレンドを創り出す〜 Advent Calendar 2022の17日目の記事となります。

本記事では、私がcvpaper.challengeに研究メンバーとして参画した経緯から1年間活動してきた内容の振り返りを、ざっくばらんに記載します。メインコンテンツは「大規模分散学習」になる見込みです。
私の体験を通して「cvpaper.challengeでの研究・活動の流れ」や「cvpaper.challengeの研究コミニュティの雰囲気」なども伝えられれば幸いです。

注：勢いに任せて執筆した結果、かなり長くなってしまったので、ご興味に応じて適当に読み飛ばしていただければと思います🙇‍♂️

自己紹介

私は現在、東京工業大学の情報工学系修士2年で、高性能計算（HPC）と深層学習を幅広く研究テーマとして扱う横田理央研究室に所属しており、B4からデータ分散学習に関する研究を行ってきました。
cvpaper.challenge内では、片岡さん・山田さん等が率いる「数式ドリブン教師あり学習（FDSL）」を主な研究テーマとして扱っており、特にFDSL手法のデータセット大規模化や画像表現改良による事前学習効果向上に取り組んでいます。

研究グループ内では、大量に実装・実験を回して結果を出す「実験屋」の立ち回りをすることが多いです。山田さんの記事で言うところの“大量の実験を管理し打ち合わせに必ず間に合わせてくる人”を目指すロールと認識しています。

興味本位で私が修士過程になってから現在までの1年8ヶ月で、回した実験数とGPU時間の合計を、実験ログから算出したところ、、、
1871個のjobを投げており、合計のGPU時間は、
1,392,798 GPU時間（1つのGPUを約159年間回したものと等価）
にまで達していました。。。
（GPUはV-100を使っていました。）
これを算出してしまってから震えが止まりませんが、知らない方が幸せなこともありますね。

というわけで記事タイトル通り、GPUを食べまくっている人間です。
もし「全国GPU時間消費量選手権」が開催された場合、かなり良い順位を狙えるのではないでしょうか。

あ、これをご覧になって「その規模の計算資源を使いたい！！！」と思った方は、横田先生のアドベントカレンダー投稿記事を是非是非ご参照ください！大規模計算を行いたい人向けの公募制度についてまとめて紹介なされている超有用記事となっています。
あるいは是非、横田研究室にいらっしゃってください！！！大規模計算でBig Techと殴り合いましょう！！！

東工大×産総研の連携（大規模実験屋視点）

ここで、私がcvpaper.challengeの研究グループと関わりを持ったきっかけである「東工大と産総研の研究連携」について、大規模実験屋の視点から記載します。
本記事では特に「連携の背景にあった研究課題」と「スパコンを用いた大規模分散学習」について、少し詳しめに記述したいと思います。
この研究連携については、速水さんの投稿記事や片岡さんの投稿記事でも言及されているので、是非そちらも併せてご参照ください。

研究背景：事前学習用データセットの大規模化傾向と課題

まず、産総研との研究連携のモチベーションとなった研究背景について、簡単に説明します。

2020年初め、Transforerベースの大規模言語モデル「GPTシリーズ」を開発しているOpenAIが、「Scaling law（スケーリング則）」と呼称される重要な仮説を提唱しました *1。深層学習におけるスケーリング則とは、「データセットサイズ、モデルサイズ、学習ステップ数の3要素を適切にスケールアップさせる限り、事前学習モデルの性能は際限なく向上していく」という仮説です。この仮説については現在も様々な議論が飛び交っていますが、実験的には、言語のみならず、画像、音声、数値データなどの他モダリティを扱うモデルについてもある程度成り立つケースが報告されています。

昨今のCV分野でも、ViTなどのTransformerアーキテクチャベースの深層モデルを大量の実画像データセット（e.g. JFT-300M/3B）で事前学習することによって、高い事前学習効果を獲得することがトレンドになっています。

しかし、事前学習用の実画像データセットの大規模化は、「データセット収集の困難化」や「倫理的バイアス・権利的な問題」、「ラベリングの高コスト化、誤り率増加」、そして「一部研究機関によるデータセットの寡占化」などの問題を深刻化させ、これらは分野の健全な研究発展を阻害し得る最重要課題とみなされ始めています。

一研究テーマとして大規模深層学習を掲げている弊横田研でも、計算資源については十分に確保できていたものの、競争力がある規模のデータセットの不足に直面し、解決の糸口を探っていました。

*1 : 2020, Jared, K. et al.,『Scaling Laws for Neural Language Models』

産総研のFDSL研究を大規模化→CVPR2022共著採択

事前学習に有用そうな大規模画像データセットを暗中模索する中で光明を指したのが、産総研の片岡さん主導で研究が進められていた事前学習手法「FDSL」でした。
FDSLは、数式ベースで自動生成した人工画像をデータセットとして用いる事前学習手法であり、当時データセットサイズは1kカテゴリ1kインスタンス規模とまだ小規模だったものの、同規模の実画像データセット（ImageNet-1k）と引けを取らない事前学習効果を発揮するケースも報告されていました *2。
ラベリングされた人工画像データセットを任意の規模で自動生成可能なFDSLは、大規模事前学習とのシナジーが非常に高い可能性を有していると考えられました。

そこで、従来より行われていた片岡さんら産総研チームと東工大の井上先生の共同研究に、弊横田研究室が加わる形で、2021年6月頃から本格的な連携を行う運びになりました。（連携きっかけのより詳細は、片岡さんの投稿記事をご参照ください。）私はその新生連携チームの一員として参加しました。

研究連携の長期目的として「実画像を一切用いない大規模なFDSL事前学習が従来の実画像事前学習の性能を上回ることを示し、実画像データセットが抱える問題を根本から払拭して事前学習のデファクトスタンダードを置き換える」ことを掲げ、短期目標は「センセーショナルな結果を提げてCVPR2022（論文〆切：11月中旬）で採択される」ことでした。

研究メンバーそれぞれが"適材適所"に配置されたこともあって、本連携は非常に高速かつ高クオリティに回り、結果的にCVPR2022にて共著論文『Replacing Labeled Real-image Datasets with Auto-generated Contours』が採択されるに至りました。
論文の詳細についてはプロジェクトページにまとまっているので、ご興味ある方は是非そちらをご参照ください。

*2 : 2022, Nakashima, K. et al., 『Can Vision Transformers Learn without Natural Images?』

スパコンを用いた大規模実験屋のお仕事

私は本連携において、128基のGPUを数日単位で用いる事前学習20回以上、16-32基のGPUを用いる微調整学習（FIne-Tuning）150回以上を、2週間程度の期間で行っていました。
ここで、どのように従来の実験を大規模化し、短期間で大量に回したのかを、技術的要素も含めて少し詳しめに記載します。

まず以下に、実際にCVPR2022で行った実験大規模化のフロー概略を図解します。

学習実験を大規模化するフローチャート概略図

以降で、図中のステップ要素それぞれについて、順を追って言及していきます。

1. 実験表作成
まず最初に、「データセット・モデルをどこまでスケールアップさせるか」や「学習のハイパーパラメータはどうするか」などを、先行研究を参考に仮決めして実験表を作成し、その後研究チーム全体で叩き上げていきます。
この際、計算リソースや時間が足りずに全てを実行することが難しい可能性を考慮して、各実験に優先度を設けておくことが重要です。

また、実験表は研究メンバーがいつでも参照できるような形で作成しておくと良いです。具体的には、HackMDやGoogle Slidesなどの表形式が使用可能なオンラインドキュメントを用いるのが効率的です。私はMarkdownが好きなので当時よりHackMDを使用しています。

2. 大規模実験のスクリプト実装
次に、大規模計算機（スパコン）を効率的に利用して大規模データセットも捌けるように、従来の実験スクリプトを改造します。記事冒頭でも紹介しましたが、スパコンを利用するための制度や申請などについては、横田先生の投稿にまとめて紹介されているため、ご興味ある方は是非そちらもご参照ください！

学習データセットを大規模化する際には、学習を現実的な時間で終わらせるために、複数の計算ユニット（GPU, TPUなど）を用いたデータ分散学習が現状必須となっています。特に、21kカテゴリ1kインスタンス規模（画像数にして2100万枚）以上のデータセットでViT-Base規模（86M params）のモデルを事前学習する際には、一般的な単一計算ノードに搭載されているGPU数（実際に使用した産総研ABCIはV-100計算ノードで4つ、A-100計算ノードで8つ）では分散数が足りずに、学習時間が週・月単位でかかるようになってしまいます。

そこで、複数ノードを利用した分散学習が必要となりますが、世に出ているスクリプトは基本的に複数ノード分散学習には対応していないため、実装に工夫が必要なケースが多いです。
複数ノード並列実行に対応するためには、例えばOpenMPIモジュールを用いて複数ノードを一元的に管理する実装などが簡潔ですが、より詳細については省かせていただきます。

データ分散以外にも、大規模化したモデルがメモリに載るように工夫したり、まだライブラリに搭載されてない機能を実装したりもしました。

3. 小規模実験の再現
2. で実装した複数ノード分散学習対応の大規模実験スクリプトが、従来の小規模での実験結果をちゃんと再現するかどうかを確認し、実装にミスがないかを確認します。
このステップでしっかりスクリプトの信頼性を担保しておくことは、計算資源の消費が激しくやり直しが難しい大規模実験において最重要事項と言っても過言ではないので、図中でも実装とは別要素として扱っています。
ここで再現ができなかった場合は、2. に戻ってデバッグします。

4. 大規模データセット生成
2, 3. での実装と並列して、FDSLで用いる大規模データセットの生成を行います。CVPR2022では、FDSL用データセット（ExFractalDB, RCDB）を1kカテゴリから10k, 21k, 50kカテゴリまでスケールアップさせたものを、それぞれ生成しています。
FDSL用データセットをスケールアップさせる方法は論文やプロジェクトページをご参照ください。

5. データセットshard化（webdatasetの利用）
4. で生成した1000万枚オーダーの画像ファイルのディレクトリ群で構成されたデータセットのままでは、ファイル数が多すぎてデータロード時のIO処理に大きく負荷がかかり遅延してしまいます。
特に、複数の大規模学習を同時に走らせた際のデータロード時のIO負荷は尋常ではなく「半日待ってもデータロードが終わらない」という状況になり得ます。（実際になりました。。。）

そこで、何枚かの単位で画像ファイルをtarファイル（shard）にまとめてアクセスできるようにして、データロード時のIO負荷を大きく下げることが可能なwebdatasetを用います。
このステップでは、webdatasetを使うためにデータセットのファイル数を圧縮（shard化）することをしています。この時、学習スクリプト中のdata loaderはwebdatasetのものに差し替えます。

webdatasetの使い方の詳細については、名古屋工業大学の玉木先生が詳しく解説した記事を出していらっしゃるので、そちらもご参考になると思います。

6. プロファイリング
スクリプトを実装し終えて、学習に用いるデータセットのshard化も終えた後は、それらを用いて数イテレーションだけ学習を回してプロファイリングを行います。

プロファイリングでは、「学習スループット」や「学習過程の各操作に費やした時間」、「GPUの利用率」などを可視化します。これによって、「分散数の増加に対してスループットが妥当に高くなっているか」や「実行時間のボトルネックとなる操作が妥当か」、「GPUを効率的に利用できているか」などを確認し、施した学習高速化が実際に機能していることを定量的に評価します。プロファイリングを行う際には、TensorBoardなどの可視化ツールを用いると良いです。

ここで、明らかにおかしな操作に時間をかけていたり、スループットがあまり高くなっていないことを確認した場合には、設定や実装を見直します。
そういう場合、「プロセスあたりのバッチサイズが小さすぎて通信のオーバーヘッドが支配的になっていた」とか「実は数値精度を半精度じゃなくて倍精度で実行していた」、「データロード時のプレフェッチ設定が良くなかった」など、小規模実験再現ステップだけでは気づきにくい罠が原因としてあります（ありました）。

本番実験に先立って、実際に使用する規模のデータセット・分散数で適切にプロファイリングをすることによって、無駄な計算資源消費を未然に防ぐことは、大規模実験を行う上で欠かせない作業です。

7. 実行
6. でのプロファイリングで妥当な性能が出ていることを確認できたら、いよいよ本番実験の実行となります。

実行時のロギングでは、ローカルファイルへの書き込みに加えて、WandBの併用をおすすめします。WandBは、自動で学習曲線の描画やハイパラの保存、GPU状態の監視などをしてくれる他、その気になって設定すれば、網羅的なハイパラ探索なども自動実行してくれる超便利実験管理webツールです。プロジェクトグループを作成すれば、グループ内のメンバーが同一のプロジェクトで作業できる点も、共同研究する上ではプラスポイントです。
（私はWandBを使い始めてからというもの、matplotlibをあまり使わなくなりました。。。）

実験管理についてもWandBだけで出来ます。が、ちょっと可視化し辛いのに加えて、webツールであるWandBが利用不可になるリスクに備えたい意向もあり、実験表上の各実験項目に実験ステータス（待機中、実行中、完了など）や最終精度などを手動で記載して、二重管理していました。

並列して行える実行jobはどんどん同時に投げますが、投げすぎるとIO処理遅延によるスループット低下を引き起こすので、スループット低下が起こらないギリギリを見極めながら投げていきます。

また、スパコン上で128基や512基規模のGPUを2-3日用いる大規模分散学習では、学習途中のGPU故障などの予期せぬエラーで実行が止まったり、他の実行との兼ね合いで極端にスループットが低下したりする恐れがあります。私は「計算資源は消費しているにも関わらずプロセスは止まっている」というタチの悪い状況に遭遇したこともあります。。。
実行時には定期的にログに異常がないか確認し、異常があった場合は再起動などの対処を速やかにできるようにしておくと、効率的な計算資源の使用につながると思います。

8. 結果報告、実験表更新
実験結果がある程度出揃ったら、結果をまとめて研究メンバー全体に対して報告し、次の実験方針や優先度の調整を行います。
ここで、実験表や実験ログが人に見せる前提で適切に記載されていると、結果のまとめや報告は容易です。

次の実験方針が固まったら、1. に戻り、データセット生成と実行のサイクルを回していきます。新しく回す実験がハイパラ調整やデータセット調整などの変更で、特に学習過程について追加実装がない限りは、2, 3, 6. のプロセスは行わず、1 → 4 → 5 → 7 → 8 (→ 1)のサイクルを高速に回していきます。

…長くなってしまい恐縮ですが、以上が各ステップの説明となります。
より技術的な詳細については、来年自分が修士課程を卒業するまでに「大規模分散学習のTIPS共有資料」を研究室同期と作成して何らかの形で共有予定ですので、乞うご期待（？）ください！

本連携では、「4. 大規模データセット生成」と「5. データセットshard化」については、従来からFDSLデータセットの生成をしていた産総研サイドの山田さんと速水さん、そして横田研のポスドクEdgarさんに主に担当していただきました。
残りのステップについては、進め方について横田先生に数多くアドバイスいただいたり、横田研同期の張君にかなりサポートしてもらったりしつつ、私主導で進めていきました。

気づいたらあてがわれていた役割分担でしたが、それぞれのメンバーが遊んでしまうことなく、非常に効率的にフローを回せたことを実感しています。これは、片岡さんのチームビルド力が隅々まで行き届いたことが大きいと感じています。

大規模実験を行う上での心構え

本節の最後に、大規模実験を行う上での心構えについて、私の思うところを少しお話させてください。

私は、本連携において、ABCIグランドチャレンジ（グラチャレ）という「4000基を超えるV-100 GPUを24時間占有できる」良い意味でイカれた公募制度も利用しつつ、短期間で多くの大規模学習を回しました。

私は本連携の後にも2022年12月現在まででグラチャレを2回、計3回やっていますが、時間制限のある中での大規模学習は何回やっても結構精神力を試されるものでした。

グラチャレを利用するしないに関わらず、大規模な計算資源消費を伴う実験は、それ単体で「やり直しが難しい」「失敗してはいけない」プレッシャーが多少は伴います。
そのプレッシャーに打ち勝つためには、「3. 小規模実験の再現」や「6. プロファイリング」といった実験本番前の確認・準備を、徹底的に行うことが有効です。
慎重に入念に準備して、本番では焦らず気負わず「これで失敗したらもう仕方ない」ぐらいの気持ちで臨むのが良いと考えています。当たり前のことかもしれませんが。。。

メンタル的な側面についてもう一つ、大量の実験を回す際には「期待せず淡々と実行する」ことが大切であると思いました。
「これだけ計算資源を使ったのだから良い結果出てくれ」と期待をしてしまう気持ちもあるのですが、それで一つの結果に一喜一憂していてはメンタルがゴリゴリ削られて作業効率が落ちますし、客観的で正常な結果評価や実験方針策定もままならなくなってしまうかもしれません。

実験表を作成した後は、どんな結果が出ても淡々と表を埋めていくことが、成果につながる秘訣と考えています。残念な結果になったとしても、残念なりに傾向を掴むことができれば、そこからより良い改善方法を見出せるかもしれません。そう考えると、実験計画をしっかり練った上での実験であれば、そもそも"残念な"実験結果という概念は存在しませんね！

cvpaper.challengeでの活動がスタート

そんなわけで産総研との連携にそれなりに身を入れたこともあってか、2022年1月に片岡さんからslack DMで「cvpaper.challengeのFDSLを研究するチームで、研究メンバーとして活動してみないか？」という旨のお誘いをいただきました。
連携を通して「研究メンバーの方々ともう少し関わりたかった」、「FDSL手法がとても興味深くその行く末が気になる」と思っていた私にとって、そのお誘いは願ってもない機会でした。
二つ返事で参加意思を伝えさせていただき、晴れてcvpaper.challengeのFDDBチームの研究メンバーとして活動する運びになりました。

その後のcvpaper.challengeでの活動

前節の大規模学習の話が少し長くなってしまいました。。。以降は少し控えめにcvpaper.challenge参加後の活動について記載していきます🙇‍♂️

cvpaper.challengeでの活動が始まってからしばらくは、CVPR2022における大規模実験の後始末（より詳細な網羅的調査など）を行っていました。
大体週一開催の研究チーム内の進捗報告会（hacks）にもすぐに参加させていただき、徐々にメンバーの方々との交流ができるようになってきました。
メンバーの皆さんは、研究力はもちろん、人格的にも非常に優れている方々ばかりで、新参者の自分にもどんどん積極的にコメントしていただいて、大変ありがたかったです。

CCCSへの論文投稿、初めての査読

参加してからしばらく経った春先に、cvpaper.challengeで夏に行われる「Cvpaper.Challenge Conference Summer; CCCS 2022」内部向けへのshort paperの投稿を勧められ、一稿執筆する運びになりました。
（cvpaper.challengeでは年2回、所属学生の一部が運営主催の「Cvpaper.Challenge Conference」という学会のシミュレーションを、内部向けと外部向けに開催しています。今年もつい最近の12/13に、CCC Winter 2022の外部向けカンファレンスが開催されました。）

私は、冬から行っていたFDSL大規模事前学習に関する追加調査で得た知見や今後の展望をまとめて、一つの論文にしました。

CCCSに投稿された論文の査読も何本かやりました。論文査読は完全に初めてのことで、まだまだ学会経験の浅い自分が他者様の論文にあれこれ言うのはちょっと気が引けましたが、論文を評価する立場からフラットに読み込むことは、批判的思考力が養われてとても良い体験になりました。

結果的に、CCCS内部向けに投稿した論文は「Best Paper Award in FDSL」に選ばれて表彰状までいただいて、嬉しい気持ちになったことを覚えています。（学生主導の運営でやっている学会なのに、そんなものまで用意してるなんて凄い。。。）

基盤モデルメタサーベイ

CCCの外部向けカンファレンスでは、豪華ゲストを招待してのトークセッションの他に、cvpaper.challengeの各研究チームがそれぞれ適当な領域の網羅的メタサーベイを行って発表するメタサーベイセッションもあります。

自分が所属するFDDBチームでは、CCCS向けに「FDSLメタサーベイ」と「基盤モデルメタサーベイ」の2つのメタサーベイを出すことになりました。
しかしFDDBチームは、山田さんや中村さん、速水さんなど運営側に回っているメンバーが多く、なかなかメタサーベイにエフォートを割く人員の確保が難しい状況でした。

そのような中「基盤モデルメタサーベイ」の発表者として白羽の矢が立ったのが、基盤モデルについての知見はほぼ持ち合わせていない新参者の私でした。（基盤モデルは基本的に大規模学習を伴うので、全く知らない訳ではなかったですが。。。）確か依頼を受けたのはメタサーベイ発表の2ヶ月前ぐらいです。
当時は「落合陽一先生など超大御所がトークセッションをされるような場で、自分の専門からそこそこ逸れた基盤モデルについて30分間話し続ける」という状況のプレッシャーを想像して、正直「ヤバすぎ。。。無理。。。」と思っていました。

しかし、そこから一気に基盤モデルに関する論文や技術ブログを読み漁ってトレンドを追い、片岡さんや山田さん、中村さんに手厚いサポートもしていただき、なんとか納得のいく形で発表資料完成まで漕ぎ着け、無事CCCS外部向けカンファレンスで発表するに至りました。
メタサーベイを振り返ると、基盤モデルに関する論文を計30本ほど集中して読んだり、それを資料としてまとめることは、シンプルに苦しかったのですが、「論文を読む心理的ハードルの低下」や、「英文書を読む速度の向上」、何より「今激アツな基盤モデルについての体系的理解」につながり、かなり自分のためになったなと実感しています。

発表で用いたスライド資料は、Transformerの登場からtext2imageモデル（DALL·E 2など）のモダンなVision系基盤モデルに至るまでの経緯と現状についてまとめたものとなっており、後々SlideShareにて公開されました。
資料公開当時、stability.ai社がリリースしたtext2imageの高性能イラストジェネレートモデル「Stable Diffusion」が流行りまくっていたこともあり、本資料は波に乗ってそこそこ大きな反響をいただくことができました。
このメタサーベイには結構な工数をかけていたので、それが報われた気がしてニヤつきながらエゴサをしておりました。

ここでは、基盤モデルメタサーベイの内容詳細については省きますが、ご興味のある方は是非下記資料をご参照ください！

ついこの前開催されたCCCWでも、藤井さんが基盤モデルについてご講演なされており、大変わかりやすく簡潔に基盤モデルをまとめられた資料がもう公開されています！ご興味ある方はそちらも併せて是非ご参照ください。

また、CCCS2022の講演資料は下記ツイートにまとまっているので、基盤モデル以外の題目についてもご興味ある方は是非ご参照ください。

その他、未発表の研究

基盤モデルメタサーベイ以降の活動内容は未公開の研究内容が含まれるため詳しいことは記載できませんが、秋一杯まではCVPR2023投稿に向けて奔走していたり、今は修論執筆やICCV2023投稿に向けて奔走したりと、色々奔走しています🏃‍♂️🏃‍♂️
しかし孤独な奔走は修論ぐらいで、例えばCVPR2023以降は、速水さんの投稿記事でも触れられていたように、速水さんと"二人三脚"の共同研究体制で奔走したりしています。

これからも体調管理とスケジュール管理に気をつけつつ、どんどん奔走していければと思います！

おわりに

以上、大規模実験屋のcvpaper.challengeでの活動振り返りでした！
勢いに任せて思いついたことをつらつらと書いてしまったので、途中読みにくい箇所や冗長な箇所などあったかと存じますが、何卒ご容赦ください。。。🙇‍♂️
この記事が、「大規模分散学習をやろうとしてる方」や「cvpaper.challengeでの活動に興味のある方」の一参考となっていたら感無量です。

残りの修士生活3ヶ月強、悔いのないように頑張って参ります！！💪💪

この記事が気に入ったらサポートをしてみませんか？