大規模日本語VLM Asagi-VLMにおける合成データセットの構築とモデル実装

大規模日本語VLM「Asagi」を構築した経験から得られた知見を紹介しています．
- Megatron-LMを用いた大規模モデルの学習知見
- 英語VLM・日本語LLMを用いたデータセット合成方法の紹介
- 訓練コードのチュートリアル

## モデルへのリンク
- https://huggingface.co/MIL-UT/Asagi-2B
- https://huggingface.co/MIL-UT/Asagi-4B
- https://huggingface.co/MIL-UT/Asagi-8B
- https://huggingface.co/MIL-UT/Asagi-14B

## 関連リンク
プロジェクトページ：https://uehara-mech.github.io/asagi-vlm
コード：https://github.com/mil-tokyo/Megatron-VLM

Transcript

1. 大規模日本語VLM Asagi-VLMにおける 合成データセットの構築とモデル実装 上原 康平, 黒瀬 優介 (東大/理研), 安道 健一郎 (理研/東大), Jiali Chen, Fan Gao, 金澤 爽太郎, 坂本 拓彌, 竹田 悠哉, Boming Yang, Xinjie Zhao (東大), 村尾 晃平, 吉田 浩 (NII), 田村 孝之 (ROIS), 合田 憲人 (ROIS/NII), 喜連川 優 (ROIS/東大), 原田 達也 (東大/理研/NII) ※ 言語処理学会第31回年次大会 （NLP2025）でポスター発表予定 プロジェクトページ（モデル・コード・論文へのリンクあり） https://uehara-mech.github.io/asagi-vlm
2. • 日本語の汎用大規模 Vision&Language モデルを構築するプロジェクト • 要件 • オープンなモデルを目指す → データセット合成などにGPTなどは使えない（出力物の利用がライセンスで制限されているため） • 100B超までスケールする実装を目指す • モデルは Huggingface Hubからダウンロードして利用できます！ • https://huggingface.co/MIL-UT/Asagi-2B • https://huggingface.co/MIL-UT/Asagi-4B • https://huggingface.co/MIL-UT/Asagi-8B • https://huggingface.co/MIL-UT/Asagi-14B • 関連リンク • プロジェクトページ：https://uehara-mech.github.io/asagi-vlm • コード：https://github.com/mil-tokyo/Megatron-VLM 1 研究概要
3. 目次 1. Asagi: 大規模モデルの構築 ▶ 1.1 事前知識・概要 1.2 データセット合成 1.3 実装・実験 2. 訓練コード解説
4. • 基本的にはエンコーダー・デコーダーモデル • エンコーダー：画像エンコーダー（CLIPなど） • Adapter ：画像特徴量をデコーダーにつなぐ • デコーダー ：訓練済みLLMを利用 3 大規模 Vision&Language モデル（VLM）とは Image Encoder Adapter LLM A dog is sitting next to the teddy-bear. Please describe this image.
5. • LLMの訓練同様，訓練を複数ステージに分けて行うことが多い • 事前訓練（stage1） → Instruction Tuning（stage2） • 各ステージごとに，訓練に使うデータの性質が大きく異なる 4 大規模 V&Lモデルの訓練 ~Billions ~Millions 事前訓練データ Instruction-tuning データ 質より量 webクロールデータなど 量より質 指示文つきデータ
6. • 事前訓練（stage1） • 画像エンコーダーとLLMはフリーズ，アダプターのみ訓練 • アダプターによる画像特徴量とテキスト特徴量をAlignmentを学習 • Instruction Tuning（stage2） • 画像エンコーダーはフリーズ，アダプターとLLMを訓練 • Instructionに応じたテキスト生成を学習 5 大規模VLMの訓練概要 Image Encoder Adapter LLM A dog is sitting next to the teddy-bear. Please describe this image. Image Encoder Adapter LLM A dog is sitting next to the teddy-bear. Please describe this image.
7. シンプルなMLPをアダプターとして用いる • 画像エンコーダー：CLIP，LLM: vicunaなど • アダプター：MLP 構造・実装がシンプルであるため，様々な派生研究でベースとして用いられている 6 大規模 V&Lモデルの例：LLaVA Liu et al. "Visual Instruction Tuning." NeurIPS 2023.
8. • 超大規模なVLM学習は，オープンな実装では見たことがなかった（2024年4月当時） • 1GPUで普通に学習できるのは7Bモデルくらいが限度 • 既存の実装（LLaVA）は，DeepSpeedでマルチGPU・マルチノード学習を実現 • DeepSpeedはモデルパラメータの通信が多くなるため， 超大規模なモデルの学習を十分高速に実行できるか確信が持てず…… → Megatron-LMをベースとした実装を採用 • ✅ メリット • LLM-jpでの利用実績があり，172BクラスのLLM学習は可能 • ❌ デメリット • 2024年4月当時，legacy → mcore への転換の真っ最中 • 既存のexampleがあまり使えず…… • VLMの実装はおろか，エンコーダー・デコーダー型モデルの実装も整備されていなかった 7 日本語VLM開発の課題①
9. • 日本語V&Lデータは圧倒的に不足している • 数百万件程度 • 回答も簡潔であり，VLMの学習に必ずしも適さない（下図参照） → 合成データメインでの学習を採用 • LLaVAデータセットなど，合成データの有効性はよく言及されている （そもそも人手で大規模に作るのは時間的・コスト的に不可能……） • GPTは使えないので，サイバーエージェントのCALM3と，MicrosoftのPhi-3.5-visionを中心に データ合成 8 日本語VLM開発の課題② STAIR Captionsのキャプション例： ⻑い机と椅⼦が、並んで置いてある テーブルの上に、植物が置いてある 背もたれが丸い椅⼦がたくさん並べられているテーブルの真ん中に⼤きなガラスの花瓶に活けられた花が 飾られている ⻑いテーブルにお揃いのイスが置かれている 10⼈掛けのテーブルの上には分厚い本やフラワーアレンジメントが置かれている Yoshikawa et al. "STAIR Captions: Constructing a Large-Scale Japanese Image Caption Dataset." ACL 2023.
10. • 実装パイプラインをシンプルにするため，多数の試行錯誤を要する処理は避けたい • LLaVAにおけるデータ合成 • 「物体検出結果の座標情報」+「キャプション情報」などをGPTに入力 • 欠点：bbox情報をオープンなLLMが正確に処理できるか不明であり，試行錯誤が必要 → よりシンプルなデータ合成 • (A) Phi-3.5-visionで生成した英語キャプション • (B) もとのデータの日本語テキスト情報（ALTテキストやWikipediaの記事など） • （A）+（B）をCALM3に入力し，合成するように指示 • （A）（B）ともに普通のテキストなので，LLMが扱いやすいだろうと予想 9 日本語VLM開発の課題③
11. 目次 1. Asagi: 大規模モデルの構築 1.1 事前知識・概要 ▶ 1.2 データセット合成 1.3 実装・実験 2. 訓練コード解説
12. • 合成データ作成の大きな方針 • キャプションデータ・VQAデータを合成 • VQAデータは，キャプションデータをもとにLLMで合成 • なお，OCR系データセットの合成は難しい（工夫の余地が大きすぎる）ので今回は諦める 11 データセット合成（概要） 英語 VLM ⽇本語 LLM ⽇本語 LLM 英語キャプション ⽇本語キャプション ⽇本語 VQA ALTテキストなど 画像
13. • 新規に収集したwebクロールデータを利用 • HTML換算で数千万〜数億件程度 • 「画像とHTMLが紐づいていないデータ」が多く存在する • 「画像にaltやfigcaptionがついていないデータ」がほとんど • 使えないデータのフィルタリング • 画像ベース • 解像度・アスペクト比によるフィルター • 色数の少ない画像の除去（PIL.ImageOps.posterizeを利用）（バナー等を除去するため） • 文書画像等の除去：DocVQAなどの画像とCLIP特徴量を比較し，文書っぽい画像を除去 • テキストベース • NGワードの除去（HojiCharフィルター） • 文字数・文字種などに基づくフィルタリング ※ なお，画像サイドでのNSFWフィルタリングは行っていない （目視で数千件確認したが，怪しい画像がなかった・テキストではフィルタリングしているため十分と判断） 12 データセット前処理（フィルタリング）
14. • fastwarc • warcファイル（Webクロールデータ用のファイル形式，CommonCrawlなどでも採用されている）の 処理用ライブラリ • warcio より高速 • selectolax • HTMLパーサー • beautifulsoupより軽量かつ高速 • trafilatura • HTMLからテキストデータなどを簡単に抽出できる • TensorRT • 画像特徴量に基づくフィルタリングは CLIP を用いたが，CLIP 推論高速化のために使用 13 データセット前処理に使ったライブラリなど
15. CLIP CLIP CLIP • 「画像に alt や figcaption がついていないデータ」の処理 • ほとんどの画像にはalt等が存在しなかった • このままではデータのほとんどが使えない → 画像を含んでいたHTMLから，使えそうなテキストを抽出したい • 画像に対する適切なテキストの選択 • 日本語CLIP（LINEヤフー）を用いて，画像とHTMLテキスト各文の類似度を計算 → もっとも類似したテキストを採用 • 類似度が閾値を下回る場合は，HTMLテキストは不採用 14 データセット前処理（ALTテキストの選択） HTMLテキスト： 富⼠には⽉⾒草が良く似合う 私は、その男の写真を…… ⽂壇の、或る⽼⼤家が…… 富士には月見草が良く似合う ← 採用
16. • Webクロールデータには， 「HTMLが紐づいていない画像データ」も大量に存在していた • 全部捨てるのはもったいない → データが「画像のみの場合」「画像とALTなどが使える場合」で合成処理を分ける 15 データセット合成（英語キャプションの活用）
17. • 共通処理：Phi3-Visionで英語キャプションを生成 • 画像のみの場合：単純に英語キャプションを翻訳 16 データセット合成（画像のみの場合） The image captures a serene scene of a traditional Japanese street .... quiet, peaceful moment frozen in time. この画像は⽇本の伝統的な路地の落ち着いた …… ⽇本語 LLM 英語キャプション ⽇本語 （翻訳）キャプション 英語 VLM
18. • 共通処理：Phi3-Visionで英語キャプションを生成 • ALTがある場合：ALTと英語キャプションをまぜあわせるようにCALM3に指示 • 画像の視覚的コンテンツと，背景情報（ALT）の両方を含むキャプションが得られることを期待 17 データセット合成（画像 + HTMLの場合） ⾦沢市の にし茶屋街 を撮影した写真 The image captures a serene scene of a traditional Japanese street .... quiet, peaceful moment frozen in time. ⾦沢市のにし茶屋街を撮影したこの写真は ……落ち着いた雰囲気を感じさせます。 ⽇本語 LLM 英語キャプション ⽇本語キャプション ALTテキストなど 英語 VLM
19. • VQAデータセットの合成 • 先述の手順で，画像に対するキャプションが得られている • CALM3を用い，キャプションからQAデータを合成 • また，JMMMUのように多肢選択式VQAベンチマークが存在することを踏まえ， 多肢選択式のVQAデータも合成 • 合成VQAデータに対し，追加で誤答選択肢を3つ程度生成 18 データセット合成（VQA） （データセット合成） ⽇本語 LLM ⾦沢市のにし茶屋街を撮影した この写真には、……、落ち着い た雰囲気を感じさせます。 ⽇本語 ⻑⽂キャプション ターゲット：撮影場所 質問：この写真はどこで撮影されたものですか？ 回答：⾦沢市のにし茶屋街で撮影されたものです。 ⽇本語 VQA
20. • 新規収集データ以外のデータセットの利用 • CommonCanvas：画像 + BLIP-2英語キャプションが大量 • LLM-jp-pair ：画像とALT的なテキストが大量 • WIT ：Wikipediaの画像・記事のデータセット（日本語記事のみ利用） • ArtBench ：絵画と作者情報のデータセット．作者についてのWikipedia記事を 追加で利用（パブリックドメインの画像のみ利用） • これらにも前述の合成手順を適用 • CommonCanvasは画像のみ利用 • WIT，ArtBenchは，Wikipediaの本文をALTとみなした 19 データセット合成に使用したデータセット
21. • CALM3 の指示追従能力は極めて優秀 • 少数（3個程度）の例示を加えることで，出力がかなり安定する • プロンプトなどは論文に載せています • LLM や CLIP を利用したデータセット合成は，実行に時間がかかる • vLLMライブラリや，TensorRT を利用した高速化 • 特に，CLIPの推論は TensorRT で大幅に高速化できる • 並列処理 • CPUで完結する処理は，ray（pythonの分散処理ライブラリ）を使ってマルチノードで実行 • GPUを要する処理は，H100が8枚搭載されたノードを8ノード程度使い，並列で前処理ジョブを実行 20 データセット合成Tips
22. • LLMでデータを合成する都合上，学習に不適切な文言が紛れ込むことがある • 「申し訳ありませんが……」「いただいたテキストからは……」など • 合成データを目視で確認し，NGワードリストを作成 → NGワードを含む合成テキストを排除 • できなかったこと • 繰り返し文字列の除去 • LLM生成データは，稀に繰り返し文字列（「卵，生姜，メンマ，メンマ，メンマ，……」など）が 含まれる • 微量であるため，事前の目視チェックでは存在に気付かず…… • 除去スクリプト自体は実装済み，次の学習から適用予定 • ハルシネーションなど • 根本的な対策は難しいが，目視で確認した限りではそこまで悪くなかった 21 データセット合成（合成データのフィルタリング）
23. • 合成キャプション • 合成VQA • Open-ended • Multi-choice • 既存の日本語データセット • 英語データセットの翻訳（by CALM3） • なお，当初は合成キャプションは Stage1のみに用いる予定だった → Stage2にも使えそうなクオリティ だったためStage2訓練にも利用 22 学習データまとめ
24. • ALTをあまり反映してくれないこともある 23 学習データの例（キャプション） ALTテキスト コアラのウッドデッキで水遊び キャプション コアラのウッドデッキで水遊びをする子供たちの写真です。 ALTテキスト 2015.2.9 どんぶり会館 (55) キャプション この画像は、青い衣装を着たキャラクターが明るい店内でさ まざまなティーポットのデザインを紹介している様子を描い ています。店内の木製の棚には、異なる形やサイズ、色をし たボウルやプレート、ティーポットが並べられており、自然 光が店内を明るく照らしています。
25. 24 学習データの例（VQA） ALTテキスト なし 質問 この画像に写っている収納カートにはどのような特徴があり ますか？ 回答 この画像に写っている収納カートは、黒い2段の引き出し収 納カートで、下部に車輪が付いています。 ALTテキスト ランペドゥーザ・エ・リノーザは、イタリア共和国シチリア 州アグリジェント県にある、人口約6,600人の基礎自治体。 イタリア最南端のコムーネである。 質問 この画像の場所はどこですか？ 回答 この画像の場所は、イタリアのランペドゥーザ・エ・リノー ザです。ランペドゥーザ・エ・リノーザは、イタリア共和国 シチリア州アグリジェント県にある人口約6,600人の基礎自 治体で、イタリア最南端のコムーネとして知られています。
26. 25 学習データの例（選択式VQA） ALTテキスト なし 質問 与えられた選択肢の中から最も適切な回答のアルファベット を直接記入してください。歩道はどうなっていますか？ 1: 金属のプレートでできています。 2: レンガの石畳のようになっています。 3: アスファルトで舗装されています。 4: 砂利の道です。 回答 2 ALTテキスト なし 質問 与えられた選択肢の中から最も適切な回答のアルファベット を直接記入してください。このロケットは何ですか？ a: スペースXのドラゴンカプセルです b: ブルーオリジンのニューシェパード3号機です c: ファルコン9ロケットです d: 宇宙探査用の人工衛星です 回答 b
27. • ALTがないと日本的なものについて正しいデータを作るのは難しい 26 学習データの例 ALTテキスト なし キャプション この画像は、オレンジ色の背景に対比的なミニマルなデザイ ンのシンプルな白い鳥が特徴です。鳥は丸みを帯びた体形 で、黒い目が2つ、黄色いクチバシがついています。翼や尾 は描かれておらず、このシンプルなスタイルのイラストに特 徴的です。全体的にミニマルなデザインで、シンプルな線と 基本的な色使いが特徴です。 ALTテキスト なし キャプション この画像は、白い丸いパンの中にダークな光沢のあるフィリ ングが入っているもので、それぞれが宙に浮かんでいるよう に見えます。左のパンは完全に開かれており、ダークなフィ リングがこぼれ出していて、右のパンは部分的に開かれ、 ダークな光沢のあるフィリングが見えています。背景は真っ 白で、白いパンと黒いフィリングのコントラストを強調して います。テキストや他のオブジェクトは一切なく、パンが開 けられた瞬間が描かれています。
28. 目次 1. Asagi: 大規模モデルの構築 1.1 事前知識・概要 1.2 データセット合成 ▶ 1.3 実装・実験 2. 訓練コード解説
29. • 基本的にはLLaVA • 画像エンコーダー：SigLIP • デコーダー：LLM-jpのリリースしたLLM（1B, 3B, 8B, 13B） • プロジェクター：2層MLP • データ • Webからクロールしたデータを元に， LLMなどでデータ合成を行った • 既存データセットなどと合わせて， 合計2000万件程度の画像・テキストペアで学習 28 実装したモデル・データ この画像を⾒て、次の指⽰に詳細 かつ具体的に答えてください。 この画像について詳しく説明して ください。 画像エンコーダー ⽇本語 LLM プロジェクター層 指⽰⽂ この画像は、東京のレインボーブリッジを撮影し たものです。...
30. • 172Bクラスのモデルは当然GPU1枚に乗り切らない → 分散学習手法を用いる必要がある • 分散学習手法 • ZeROベースの手法 • DeepSpeedなどに実装 • パラメータや勾配，optimizerを各GPUに分散配置して，メモリ消費を削減 • 必要なタイミングでパラメータを送り合う • パラメータの通信頻度が多くなるため，100B超の学習には不向きと判断 • モデルパラレル（今回はこちらを採用） • モデルを各GPUに分散配置 • 各GPUへの配置を制御できるため，GPUクラスタの構成に応じた設計が可能 → 通信量を削減して高スループットを狙える • テンソルパラレル・パイプラインパラレルなど 29 実装の方針
31. • 行列計算は小行列に分割して計算することができる 例：Aを列方向に3つに分割する場合 • 行列を分割して計算 → 適切なタイミングでの集約を繰り返すように実装することで， 1GPUに乗らないモデルの計算を，マルチGPUで計算することができる 30 分散学習：モデルパラレル（Tensor Parallel） https://huggingface.co/docs/text-generation-inference/conceptual/tensor_parallelism
32. • モデルをレイヤーごとに分割し， 各GPUに順番に配置・計算 • ナイーブに実装すると，Forward・Backward計算の 待ち時間（Bubble）が発生してしまう → 様々な方法でBubbleをなるべく減らす工夫が行われ ている 31 分散学習：モデルパラレル（Pipeline Parallel） https://colossalai.org/docs/concepts/paradigms_of_parallelism Huang et al., “GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism.” NeurIPS 2019.
33. • データパラレル（DP）・テンソルパラレル（TP）・パイプラインパラレル（PP）は， 同時に適用することができる → 3D Parallelismと呼ばれる • なお，各手法をどのように適用するかは慎重に考える必要がある 例：TPは通信オーバーヘッドが大きいため，なるべくノード間通信が発生しないようにする 32 分散学習：3D Parallelism https://www.microsoft.com/en-us/research/blog/deepspeed-extreme-scale-model-training-for-everyone/
34. • Megatron-LMでの実装 • 3D parallelismによる分散学習 • テンソルパラレル・パイプラインパラレル・データパラレルをMegatron-LMベースで実装 • 問題点 • エンコーダー・デコーダーモデルの実装が存在せず，パイプラインパラレルの実装が大変 • デコーダーonlyの場合：各ランクに同数のLayerを乗せればよい • エンコーダーがある場合 • エンコーダーも適当に分割 ❌ 実装が大変（ただし現状のMegatron実装はこれをサポートしている？） • ランク0にエンコーダー + デコーダーの最初の分割を詰め込む ✅ 実装は楽 ❌ ランク0のパラメータ数が増えるため，最適とはいえない 33 実装の詳細
35. • Dataset関連は，megatron-energonライブラリを利用 • webdatasetベースのdataset utilityを提供するライブラリ • Megatronチームが公式に提供しているdatasetライブラリ • 当初，50億件の画像・テキストでの学習を行う計画と聞いていたため， pytorchのdataloaderでは対応できないだろうと判断 • 結果的に2000万件程度に収まったため，実際にはpytorch実装でも十分だった……？ • dataloader の checkpointing などがデフォルトで提供されている • マルチノード設定だとバグあり（？） • いろいろ修正はしたが，マルチノードかつmulti-worker設定だとcheckpointからのresumeに失敗する問題が 未解決 • 訓練時はworker数を必ず1に設定することで対応 • 特に大規模な場合，micro batch sizeは1が限界なので，さほど悪影響はないはず……？ • 最新版だと修正されているかも 34 実装の詳細（Dataloader関連）
36. • 学習はAWS HyperPod上で実施（H100 × 8枚 × 最大24ノード） • GPU不良など • 1週間に1回程度は発生 • nvidia-smiがエラーを吐くパターン • nvidia-smiは正常だが，pytorchで特定のGPUを認識できないパターン • 基本的にはノードの再起動で対応 → ダメならノードのreplaceで対応 • 学習が途中でclashした場合，wandb経由でslackにエラー通知を飛ばす （なぜか深夜に停止しがち……） 35 学習中の問題
37. • 高解像度画像が訓練時に多く含まれると学習が停止するこ とがある • timeout errorで終了してしまう • 1バッチのデータサイズが大きすぎると通信が詰まる？ • ひとまず全学習データをチェックし，1000px以上の画像は 事前にリサイズ • たまに（数日に一回程度）チェックポイント保存時にtimeout • 100B超の場合にのみ発生 • 正確には，保存終了後，学習を再開する際にtimeout errorと なる • 常に発生するわけではないため，原因が不明 • ストレージの一時的な問題？ 36 学習中の問題 ↑ 停止直前のバッチに含まれるデータサイズのプロット． 停止直前のデータサイズが大きくなっている
38. • Context ParallelやSequence Parallel • CLIPの画像特徴量の次元数（class token含めて577）がTP数（4や8）で割り切れないため • SigLIPなら割り切れるはずなので，実は採用できたかも……？ • FP8 • Attention層などの行列次元数が8または16で割り切れる必要があるが， その条件が満たせなかったため • Stage1 パイプラインパラレル • Stage1ではLLMは訓練しない • PP > 2に設定すると，必ず「更新するパラメータが一切含まれない」ランクが発生してしまう • optimizerにempty listが渡されてしまい，エラーとなる • 172Bクラスであっても，stage1はTP 8でギリギリ学習可能だったため，深入りせず • 最新のMegatron-LMでは実現可能かも（未検証） 37 実装上の課題
39. • ハイパラなど • LLaVAやNVLM（NVIDIAのLLaVAベースの モデル）を参考に設定 • ハイパラは全モデルサイズで共通 • なお，ここでのバッチサイズは global batch size 38 学習設定
40. • 14Bモデル（LLM13B） • 特に問題なし？ • 大きなloss spikeも観測されず 39 学習状況（14B stage2）
41. • GPT合成データを用いていないモデル（†）の中では最高スコア（表下線部） • モデルのパラメータサイズ増大に伴い，評価スコアも上昇 40 定量評価 Heron-Bench JA-VLM-Bench-In-the-Wild JA-VG-VQA-500 Model LM Size ROUGE-L LLM(/5.0) ROUGE-L LLM(/5.0) Japanese InstructBLIP Alpha † 7B 14.0 20.8 2.42 - - Japanese Stable VLM † 7B 24.2 23.3 2.47 - - LLaVA-CALM2-SigLIP † 7B 43.3 47.2 3.15 17.4 3.21 Llama-3-EvoVLM-JP-v2 8B 39.3 41.4 2.92 23.5 2.96 VILA-jp 13B 57.2 52.3 3.69 16.2 3.62 Ours 2B † 1.8B 44.7 48.8 3.26 53.7 3.69 Ours 4B † 3.7B 49.3 49.6 3.38 55.6 3.78 Ours 14B † 13B 55.8 50.8 3.44 56.8 3.84 GPT-4o - 87.6 37.6 3.85 12.1 3.58
42. 41 生成結果（14B） Q: 町並みの先に写っているお城の名前は何でしょうか？また、お城の向こ う側には何が写っていますか？ A: お城の名前は「江戸城」です。お城の向こう側には、富士山が写ってい ます。富士山は日本の象徴的な山であり、その雄大な姿が写真に写ってい ます。 Q:この画像を見て、次の質問に詳細かつ具体的に答えてください。この写 真について、非常に細かいところまで精密に説明してください。 A:この写真は、夜の渋谷スクランブル交差点を上空から捉えたものです。 無数のビルが立ち並び、その間を行き交う車のヘッドライトやテールラン プが光の軌跡を描いています。特に目を引くのは、中央にそびえ立つ巨大 なビルで、その壁面には鮮やかなネオンサインが輝いています。このビル は、渋谷の象徴的な存在であり、日本のポップカルチャーの中心地である 渋谷の活気を象徴しています。
43. 42 生成結果（14B） Q: この画像を見て、次の指示に詳細かつ具体的に答えてください。この写 真の内容について詳しく教えてください。 A: この写真は、東京のレインボーブリッジを背景にした都市景観を撮影し たものです。橋は白と灰色の塗装が施されており、水の上に優雅にアーチ を描いています。空は澄み切った青色で、少数の雲が浮かんでいます。遠 くには高層ビルが立ち並ぶ都市の風景が見えます。 Q: この画像を見て、次の指示に詳細かつ具体的に答えてください。これは 日本庭園の中でも、どのような形式に分類される庭園ですか？また、その 理由は何ですか？ A: これは日本庭園の中でも、池泉回遊式庭園と呼ばれる形式の庭園です。 池泉回遊式庭園は、池を中心に配置し、その周りを回遊できるように設計 された庭園です。この形式の庭園は、自然の美しさを最大限に引き出し、 訪れる人々に静寂と安らぎを提供することを目的としています。
44. 43 生成結果（14B） Q: この画像を見て、次の質問に詳細かつ具体的に答えてください。この写 真について評価してください。 A: この写真は非常に美しいです。東京タワーの赤と白のコントラストが鮮 やかで、青空と白い雲との対比が印象的です。周囲の緑豊かな木々と街灯 が都会的な雰囲気を醸し出しており、全体的にバランスの取れた構図と なっています。 Q: この画像を見て、次の質問に詳細に答えてください。この画像の場所を 舞台とした小説のあらすじを書いてください。 A: この画像の場所を舞台とした小説のタイトルは「階段の秘密」です。物 語は、都会の喧騒から離れた静かな路地裏にある古い石造りの階段を舞台 に展開します。主人公は、この階段を登りながら、過去と現在が交錯する 不思議な体験をしていきます。階段の途中には、歴史を感じさせる苔むし た石壁や、伝統的な木造建築の建物が立ち並んでいます。主人公は、階段 を登るたびに、過去の出来事や人々の記憶に触れ、自分自身の過去と向き 合います。そして、階段の頂上にたどり着くと、そこに隠された秘密が明 かされるのです。
45. 目次 1. Asagi: 大規模モデルの構築 2. 訓練コード解説 ▶ 2.1 環境構築 2.2 データセット形式 2.3 モデルの実装 2.4 訓練
46. • コンテナとしてApptainer（旧singularity）を使った ✅ 実行時に root 不要 ✅ イメージが単一のファイル（.sifファイル）に保存されるため，すぐれた可搬性をもつ ✅ Dockerイメージをベースイメージとして使えるため，Dockerの資産を使い回せる • リポジトリの environments/recipe_aws.def が singularity のレシピファイル • 以下のようなコマンドを実行すると，一通りのライブラリなどが揃ったイメージファイル （hoge.sif）が作成される 45 コンテナの作成
47. • ベースイメージ：nvcr.io/nvidia/pytorch:24.09-py3 • PyTorch 2.5.0, TransformerEngine1.10, CUDA12.6.1などがインストールされているコンテナ • https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-09.html • 追加でインストールしているライブラリ（一部抜粋） • EFA関連のドライバー • transformers==4.45.2 • megatron-energon==2.3.0 • wandb など（詳細は environments/recipe_aws.def 参照） 46 コンテナの中身
48. • singularity exec コマンドで，コンテナ内で任意のスクリプトを実行可能 • --nv : NVIDIAのGPUを使うためのオプション • Dockerでいうと --gpus all • -B, --bind：ホストマシンのディレクトリをコンテナ内にマウントするオプション • Dockerでいうと -v 47 コンテナの実行
49. 目次 1. Asagi: 大規模モデルの構築 2. 訓練コード解説 2.1 環境構築 ▶ 2.2 データセット形式 2.3 モデルの実装 2.4 訓練
50. • megatron-energon形式にしておく必要がある • megatron-energonとは • 公式Doc: https://nvidia.github.io/Megatron-Energon/index.html • GitHub: https://github.com/NVIDIA/Megatron-Energon • webdataset をベースとしたデータセット関連ライブラリ • データの各サンプルファイルをtarにまとめておく → 大規模なデータセットを扱うのに最適 • 以下，VQAデータ（画像・質問・回答がペアになっているデータ）を想定して説明 49 データセットの準備
51. • 各サンプルごとに，以下の形式のファイルを作成 • [idx].jpg ... 画像ファイル • [idx].txt ... 回答ファイル • [idx].json ... 質問を含むjsonファイル（例： {"question": "What is this?"}） • この形式のファイルセットを，バッチごとにtarにまとめる • つまり，（総サンプル数 / バッチ数）個のtarファイルを作成する 50 megatron-energon形式について
52. • energon prepare コマンドを実行 • 対話形式でenergon形式へと変換が可能（ energon prepare コマンド） • あらかじめいくつかのタスクが定義されている（ここでは VQAWebdataset を使用） 51 tar ファイルから energon 形式への変換 field name for image は画像ファイル拡張 子名（今回はjpg） field name for context は質問文ファイル拡 張子 + key名（今回はjson[question]） field name for answer は回答が入っている ファイル拡張子（今回はtxt）
53. • 以下のように，メタファイルを含むディレクトリ（.nv-meta）と， 各tarの情報を持つ hoge.tar.idx というファイルが作成される 52 energon 形式への変換完了
54. 目次 1. Asagi: 大規模モデルの構築 2. 訓練コード解説 2.1 環境構築 2.2 データセット形式 ▶ 2.3 モデルの実装 2.4 訓練
55. • 全体的な流れ 1. huggingface 形式でモデルを初期化する 2. huggingface形式からMegatron-LM形式へ変換する 3. 訓練 4. 訓練後のモデルをhuggingface形式に戻す 54 モデル構築
56. • huggingface 形式でモデルを初期化する • scripts/convert/build_llava_jp.py として実装 • huggingface model hubから，任意の画像モデルと，任意のテキストファイルをダウンロードして 組み合わせるスクリプト 55 モデル構築（① モデル初期化）
57. ① 構築したいモデルのhuggingface形式で定義したpythonファイルを実装 • 例： scripts/convert/configuration_llava.py, scripts/convert/modeling_llava.pyなど • HFのPretrainedModel を継承しておく • configでエンコーダー・デコーダーを変更できるようにしておくと便利 ② 訓練済みモデルのパラメータで各モジュールをアップデートする • HFの PretrainedModel を継承していれば load_state_dictメソッドで任意のパラメータで更新可能 56 モデル構築（① モデル初期化）
58. • huggingface 形式のモデルをMegatron形式に変換 • scripts/convert/convert_hf_to_megatron.py として実装 • ただし，上記はwrapperであり，実体は tools/checkpoint/loader_llava_hf.py と tools/checkpoint/saver_megatron_llava.py • loader_llava_hf.py ：Huggingfaceモデルの読み込みを担当 • saver_megatron_llava.py ：Megatron形式での保存を担当 • このとき，モデルパラレルの分割も行う 57 モデル構築（② HF → Megatron変換）
59. • Megatron-LMにおけるモデル実装 • Transformer関連のモジュールはあらかじめ用意されている • DotProductAttention ... TEDotProductAttention関数など • Linear ... TEColumnParallelLinear, TERowParallelLinearなど • これらを組み合わせてTransformer Blockを作る（ TransformerLayerSubmodules クラス） • 各モジュールの細かいパラメータは argument で調整 • 例 • 各layerのbiasの有無 • 活性化関数の種類 • normの種類 など…… • argumentで設定できるパラメータは350以上あるので，細かい設定漏れに注意！ • 実際の失敗例：RMSnormとLayerNormを勘違いしていた，あるモジュールのbiasが設定できていなかった 58 Megatron形式でのモデル実装
60. • 基本的には，Transformerの各部品が実装済み → それらをうまく組み合わせて所望のモデルを表現 + arguments で細かい調整 • Llamaの場合 • 各Transformer ブロックの部品 • input layernorm • self attention • query projection • key projection • value projection • output projection • post attention layernorm • MLP • gate projection • up projection • down projection 59 Megatron形式とは https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/te_llama/tutorial_accelerate_hf_llama_with_te.html
61. • TransformerLayerSubmodules：Transformerブロックの要素を指定する dataclass • Llamaの場合 • 各ブロックの部品 • input layernorm • self attention • query projection • key projection • value projection • output projection • post attention layernorm • MLP • gate projection • up projection • down projection 60 LLaMAの実装例（TransformerLayerSubmodules） TransformerLayerSubmodules( self_attention=ModuleSpec( module=SelfAttention, params={"attn_mask_type": AttnMaskType.causal}, submodules=SelfAttentionSubmodules( linear_qkv=TEColumnParallelLinear, core_attention=TEDotProductAttention, linear_proj=TERowParallelLinear, q_layernorm=IdentityOp, k_layernorm=IdentityOp, ), ), input_layernorm=TENorm, self_attn_bda=get_bias_dropout_add, pre_mlp_layernorm=TENorm, mlp=mlp, mlp_bda=get_bias_dropout_add, ),
62. • TransformerLayerSubmodules：Transformerブロックの要素を指定する dataclass • Llamaの場合 • 各ブロックの部品 • input layernorm • self attention • query projection • key projection • value projection • output projection • post attention layernorm • MLP • gate projection • up projection • down projection 61 LLaMAの実装例（TransformerLayerSubmodules） TransformerLayerSubmodules( self_attention=ModuleSpec( module=SelfAttention, params={"attn_mask_type": AttnMaskType.causal}, submodules=SelfAttentionSubmodules( linear_qkv=TEColumnParallelLinear, core_attention=TEDotProductAttention, linear_proj=TERowParallelLinear, q_layernorm=IdentityOp, k_layernorm=IdentityOp, ), ), input_layernorm=TENorm, self_attn_bda=get_bias_dropout_add, pre_mlp_layernorm=TENorm, mlp=mlp, mlp_bda=get_bias_dropout_add, ), ポイント ・query, key, value projectionは linear_qkv という単一のモ ジュールにまとめる ・Linear関数は，TensorParallelに対応した実装を選ぶ （ColumnParallel, RowParallelなど）
63. • Llamaの場合 62 その他のモデル実装例（TransformerLayerSubmodules） • CLIPの場合 TransformerLayerSubmodules( self_attention=ModuleSpec( module=SelfAttention, params={"attn_mask_type": AttnMaskType.causal}, submodules=SelfAttentionSubmodules( linear_qkv=TEColumnParallelLinear, core_attention=TEDotProductAttention, linear_proj=TERowParallelLinear, q_layernorm=IdentityOp, k_layernorm=IdentityOp, ), ), input_layernorm=TENorm, self_attn_bda=get_bias_dropout_add, pre_mlp_layernorm=TENorm, mlp=mlp, mlp_bda=get_bias_dropout_add, ), TransformerLayerSubmodules( self_attention=ModuleSpec( module=SelfAttention, params={"attn_mask_type": AttnMaskType.no_mask}, submodules=SelfAttentionSubmodules( linear_qkv=TEColumnParallelLinear, core_attention=TEDotProductAttention, linear_proj=TERowParallelLinear, q_layernorm=IdentityOp, k_layernorm=IdentityOp, ), ), input_layernorm=TENorm, self_attn_bda=get_bias_dropout_add, pre_mlp_layernorm=TENorm, mlp=mlp, mlp_bda=get_bias_dropout_add, ), ポイント ・TransformerBlock内であれば（このレベルでは）概ね一緒 ・Attention Maskがcausalかどうかの違い cf. megatron/core/models/gpt/gpt_layer_specs.py
64. • TransformerLayerSubmodules で大まかなモデルのブロックを実装 • arguments で細かい設定を実装 • arguments の設定におけるモデルごとの違い 63 その他のモデル実装例（arguments） • CLIPの場合 normalization: “RMSNorm” group_query_attention: True # GQAが有効な場合のみ swiglu: True add_bias_linear: True add_qkv_bias: False normalization: “LayerNorm” group_query_attention: False swiglu: False add_bias_linear: True add_qkv_bias: True ポイント ・argumentsでbiasの有無・normの種類を切り替えている ・設定漏れが起こりやすいので注意 • LLaMAの場合
65. • TransformerBlock以外の実装（embeddings, final layer norm, logit layerなど） • 基本的には普通にPyTorchでモデルを実装するのと同様，forward関数に必要な処理を書く • LLaMAの場合（一部抜粋） 64 Megatron での forward 実装 # Decoder embedding decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids) # Rotary positional embeddings rotary_pos_emb = None rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( inference_params, self.decoder, decoder_input, self.config ) rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len) # Run decoder. hidden_states = self.decoder( hidden_states=decoder_input, attention_mask=attention_mask, inference_params=inference_params, rotary_pos_emb=rotary_pos_emb, packed_seq_params=packed_seq_params, **(extra_block_kwargs or {}), ) # logits and loss logits, _ = self.output_layer(hidden_states, weight=output_weight) if labels is None: # [s b h] => [b s h] return logits.transpose(0, 1).contiguous() loss = self.compute_language_model_loss(labels, logits) return loss embeddingの計算 Transformerの計算 logitsの計算 lossの計算
66. image embeddingの計算 projection layerの計算 language embeddingの計算 embeddingsの結合処理など LLM Decoderの計算 • ここまでに作成したLLMとVision Modelを組み合わせる • 基本的には普通にPyTorchでモデルを実装するのと同様 • LLaVAの場合のforward関数（一部抜粋） 65 マルチモーダルモデルの forward 実装 image_embeddings = self.vision_model(images) image_embeddings = image_embeddings.permute(1, 0, 2) # map vision model output size to language model input size. image_embeddings = self.vision_projection(image_embeddings) language_embeddings = self.language_model.embedding( input_ids=input_ids, position_ids=position_ids ) # concat and process image embeddings and language embeddings combined_embeddings, labels, attention_mask, position_ids = self._merge_input_ids_with_image_features( image_embeddings, language_embeddings, input_ids, labels, attention_mask ) if labels is not None: labels[labels == -100] = 0 output = self.language_model( None, # input_ids None, # position_ids attention_mask, decoder_input=combined_embeddings, labels=labels, inference_params=inference_params ) return output
67. • Huggingfaceモデルから，実装したMegatronモデルへweightをコピーする • layerの名前や形状は異なるので，すべてチェックして間違えないように作業する必要あり • 今のところ，huggingfaceのlayer名とMegatronのlayer名をルールに従って変換する方法はない → 普通に全layerの名前を確認して， 手作業で変換テーブルを書く必要がある • これは本当に苦行ですが，近道はないです 諦めてひとつひとつlayer名とshapeを チェックしながら作業しましょう…… • ⚠ linear_qkvのように，もともと複数のlayer だったものが統合されている場合は， shapeをそろえる際にミスしやすいので注意 66 HF → Megatron変換 https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/te_llama/tutorial_accelerate_hf_llama_with_te.html
68. • 変換が正しくできているかの確認 ① 各layerのパラメータの平均・分散が一致しているか確認 ② 実際にダミーデータを流してみて，各layerごとに出力が一致するか確認 ③ 推論を実際に実行して確認 • ただし独自モデルのMegatron推論は極めて大変 ④ HF → Megatronに変換したモデルを，さらにMegatron → HFに変換して確認 • ただしMegatron → HFの変換にバグがあったら無意味 • 個人的には ① or ② が手軽なのでおすすめ • ただしパラメータの平均・分散が一致していることは正確な変換の保証にはならないので注意 67 HF → Megatron変換のチェック
69. • HF → Megatron への変換時に分割を行う （Tensor Parallel数 × Pipeline Parallel数）個に分割されたMegatronチェックポイントを保存する • 実装的には割とシンプル • Pipeline Parallel ... 単純に decoder layer num // PP size ごとに区切って保存 • Tensor Parallel ... torch.chunk で各layerのweightをTP sizeごとに区切って保存 • どちらかというと，PP・TP分割をすべきlayerかどうかを間違えないようにする • 例：LayerNormはTP分割しない・Word EmbeddingやViTのConv層はTP/PP分割をしない 68 モデルパラレルの設定
70. • モデルパラレルの分割は，任意に設定していいわけではない • TPの制約 • モデルのすべての隠れ層サイズの公約数である必要がある • PPの制約 • モデルの Layer数の約数である必要がある • TP×PP が global batch size の約数である必要がある • ノード間通信は極力避ける • TPの方が通信が頻繁に発生する → TPは極力ノード内にのみ配置されるようにする • つまり，TPはノード内GPU数の約数であるべき 69 モデルパラレルのチューニング
71. たとえば：70Bモデルをどう分割すべきか？ • 前提：70BのGPUメモリ消費量（訓練時）は 1800GBくらい（※ 入力データやモデルアーキテクチャによって異なります） → 1GPUあたり メモリ80GBの H100を使う場合，1800 / 80 = 22~23枚くらいあれば足りそう • 1ノードあたり8枚のGPUがある → TPは最大8に設定できる • ノード間通信はなるべく減らしたいため，TPは最大値の8に設定 • よって，PPは 23 / 8 ≒ 3 以上 • 隠れ層数（=80）は3で割り切れないため，PPは4以上 • 結論：TP8・PP4がよさそう • ただし，実際にはいろいろ試してスループットを実測しながらチューニングするのがよい 70 モデルパラレルのチューニング例
72. 目次 1. Asagi: 大規模モデルの構築 2. 訓練コード解説 2.1 環境構築 2.2 データセット形式 2.3 モデルの実装 ▶ 2.4 訓練
73. • 基本的には，megatron.training.training.pretrain関数を使う • cf. scripts/train/pretrain_vlm.py • Dataloader，model，forward stepの実装と，引数を設定するだけ • modelの実装 ：作成したモデルclassを返す関数model_provider_wrapperを書くだけなので略 • forward_step：ほぼmodelをcallするだけのwrapperなので略 72 訓練概要 pretrain( train_valid_test_dataloaders_provider, model_provider_wrapper, ModelType.encoder_or_decoder, forward_step, args_defaults={'tokenizer_type': 'HuggingFaceTokenizer'}, strict_checkpoint_loading=False, extra_args_provider=add_multimodal_extra_args, get_embedding_ranks=llava_embedding_ranks, get_position_embedding_ranks=llava_position_embedding_ranks, ) Dataloaderの実装 modelの実装 forward stepの実装
74. • Dataloaderの実装（ train_valid_test_dataloaders_provider ） • ほとんどmegatron-energonの便利関数を使うだけ • task_encoder.loader.get_train_dataset などでtrain, valのdatasetを取得 • get_savable_loaderなどで energon形式の dataloaderを取得 • ポイント：get_train_datasetなどに渡すTaskEncoderの実装 73 データセット関連の実装 from megatron.energon import ( get_loader, get_savable_loader, get_train_dataset, get_val_dataset, ) train_dataset = get_train_dataset(task_encoder=TaskEncoder(), ...) val_dataset = get_val_dataset(task_encoder=TaskEncoder(), ...) train_dataloader = EnergonDataloader(get_savable_loader(train_dataset, worker_config=worker_config)) valid_dataloader = EnergonDataloader(get_loader(val_dataset, worker_config=worker_config))
75. • TaskEncoderの実装 • Dataset作成時に，dataset classを選択していた（ここでは VQAWebdataset） • このとき，左図のような VQASampleクラスが内部では使われるようになる ① megatron.energon.task_encoder.base.DefaultTaskEncoderクラスを継承した TaskEncoderクラスを作成 ② 引数でSampleクラスを受け取るencode_sample関数を実装し， 画像やテキストに対する前処理などを実装 74 TaskEncoderの実装（データセット） @dataclass class VQASample(Sample): image: torch.Tensor context: str answers: Optional[List[str]] = None answer_weights: Optional[torch.Tensor] = None def encode_sample(self, sample: VQASample): img = self.get_visual_transform(np.array(sample.image)) question_token = self.tokenizer.tokenize(sample.context) answer_token = self.tokenizer.tokenize(sample.answers) text_sample = np.concatenate([question_token, answer_token]) return ImageTaskSample( __key__=sample.__key__, __subflavors__=sample.__subflavors__, img=img, text=text_sample, prompt_len=prompt_len )
76. • 先ほど実装した pretrain_vlm.py を実行する • コマンドの引数で，訓練に関係するargumentsを設定する • optimizer関連 • optimizerの種類とパラメータ，バッチサイズ，weight decay，... • 実はモデルのlayer数や隠れ層のサイズもargumentsで指定する • huggingfaceのmodel configから読み込んで適当に設定すればよい • コマンドの基本形 75 訓練コマンド概要 singularity の設定 分散学習の設定（by torchrun） 訓練コマンド
77. • 実際のコマンド • arguments数が異常に多いため，何らかのスクリプトを書いて手間を減らした方がよい • たとえば scripts/train/run.py 76 訓練コマンド
78. • 基本的にはargumentsでいろいろ設定できる • 例： • --use-distributed-optimizer : optimizerも分散させる．高速化というかメモリ効率UP • --bf16 : BF16で計算する．余程の理由がない限り設定すべき • --fp8 : FP8で計算する．高速化できるかと思いきやバグっていて使えない • --use-flash-attn : flash-attentionを使って計算してくれる……はず（環境によっては使えない） • --distributed-backend : ncclにすべき．これでエラーが出る場合は環境がおかしい • --recompute-activations : これもメモリ効率UP 77 高速化（arguments編）
79. • 例： • NCCL_SOCKET_IFNAME : 通信に使えるネットワークインターフェース名を調べて書く （書かないと変な通信経路が使われてしまうことがある） • FI_PROVIDER : AWSの場合は “efa” • FI_EFA_USE_DEVICE_RDMA : RDMAが有効な環境では”1” • LD_LIBRARY_PATH : EFA関連のパスを通す（ /opt/amazon/efa/lib64, /opt/aws-ofi-nccl/lib ） • その他細かい設定は scripts/train/run.py 参照 78 高速化（環境変数編）
80. • 訓練が終わったら，Megatron形式でチェックポイントが保存される • 推論用に，Huggingface形式に再度変換を行う • 基本的には HF → Megatron変換の逆を実装するだけ • HFとMegatronのlayer名の対応を調べる • weightのshapeを確認する • weightをMegatron → HFにコピー • HFの推論は簡単なので，実際に推論を走らせてチェックするとよい 79 訓練後の Megatron → HF変換
81. • HF形式に変換できていれば，推論は一般的なHFモデルの推論と同様 （以下は一部抜粋） 80 推論
82. • 大規模日本語VLM「Asagi」を開発した経験から得られた知見を紹介 • Megatron-LMを用いた大規模モデルの学習知見 • 英語VLM・日本語LLMを用いたデータセット合成方法の紹介 • 訓練コードのチュートリアル • モデルは Huggingface Hubからダウンロードして利用できます！ • https://huggingface.co/MIL-UT/Asagi-2B • https://huggingface.co/MIL-UT/Asagi-4B • https://huggingface.co/MIL-UT/Asagi-8B • https://huggingface.co/MIL-UT/Asagi-14B • 関連リンク • プロジェクトページ：https://uehara-mech.github.io/asagi-vlm • コード：https://github.com/mil-tokyo/Megatron-VLM 81 まとめ