Team Tanuki Phase2 - Tanuki-8B, 8x8B 開発の経緯

はじめに

2024/08/30(金)、GENIAC松尾研LLM開発プロジェクトにて開発した大規模言語モデルTanuki-8B、Tanuki-8x8Bが公開されました！そのマルチモーダル版であるTanuki-8B-vision、Tanuki-8x8B-vision-expも同時に公開されております。

本稿では、Tanuki-8B, 8x8B開発の経緯を概要レベルで示します。

全体の流れ

Phase2の開発を「今」振り返ると、概ねこのような流れでした。

ここでは活動を、以下の３つに分類して表現しています。リンクは、Team Tanukiメンバーによる記事へとリンクしております。

分類	概要	活動
1. モデル開発	提出・公開するモデルの開発。 スケジュール死守必須、 要リスク管理。	・モデルアーキ／開発環境の選定・構築 ・新トークナイザー構築 ・事前学習 (Llama3 38B, 8B, 8x8B MoE) ・事後学習 (SFT, DPO)
2. 手法開発	1. に投入する 学習・性能向上手法の開発。	・合成データ開発環境構築 (Nemotron等) ・FT・PT手法の試行と評価 ・マージ検討 ・MoE up cycling検討
3. 拡張ソリューション開発	モデルに他の要素を追加する タイプの開発。	・マルチモーダル(LLaVA) ・エージェント

理解しやすくするために上記のように活動を分類しましたが、実際には「事前学習データと事前学習」、「ポストトレーニング検討とファインチューニング」、「合成データ手法開発と指示データ構築」のように、一体と呼んでよい緊密に連携していた活動の組み合わせも多く存在します。

プロジェクトの推移

ここでは、プロジェクトの推移を５段階にわけて解説したいと思います。

事前準備：(5/28(火)～5/31(土))

Phase2開始前の話です。

Team Tanuki Phase1メンバーは、もしもPhase1で優勝できるとしたら、Phase2の準備を早く始めないと間に合わないのでは？という危機感を持っており、事前に準備を開始しておりました。実際には、5/28(火)よりももっと前から活動していた話もあったりします。

・2024/05/28(火) Phase2事前準備開始

Phase1の各種提出を終えたのち、Phase2が始まる前に以下を進めました。

• 事前学習データセット構築
• 新トークナイザー構築
• 数十Bクラスの学習環境準備
• マージの取り組み
• 合成データの作成テスト
• 数学のDPOデータ作成、DPOの実験、DPO同士のモデルをMoE

立ち上げ：モデルアーキと開発環境検討 (6/01(土)～6/10(金))

・2024/06/01(土) Phase2正式開始(17:00?)、ノード21個が空くと即座に活動開始

プロジェクト第一フェーズ結果発表会の開始前の午前中も、行われたその夜も、活発に活動が行われていました。そのため、既に家で次の準備を進めていた方も多く、結果発表会への現地参加が少なかった、ということも発生しました 笑。

当日、結果発表会に参加したメンバは全員、東大松尾研究室の会議室をお借りしての（恐らく公式にはPhase2最初の）作業に現地で参加、Zoomを立ち上げたので遠隔からも多くの方に参加頂けました。

・効率よく学習できるモデルと環境探し

最初に行ったのは、Phase2で開発するモデルアーキテクチャと学習環境の選定です。候補は以下の３つでした。

1. Llama35B
2. DeepSeek v2立ち上げ
3. Llama 30-50B with transformer engineの立ち上げ

最初に立ち上がったのは 1. Llama35Bです。こちらは試行でとりあえず学習を開始することができ、5日間学習しました。
他方、Scaling Lawを考慮すると、少しでも学習をたくさん回す方が性能向上に有利ですので、「更に高速な学習の実現」を目指して 2. と 3. の立ち上げと高速化に取り組みました。目安を「1000B tokenを7月中旬？ぐらいまでに事前学習できるスピードを確保する」として、450～500FLOPSを目標として活動しました。

Phase2計算量の試算
大規模言語モデルの学習に必要な計算量を試算する｜Kan Hatakeyama (note.com)

序盤：transformer engineによる38Bモデルの事前学習＋各種検討 (6/11(土)～6/27(木))

・2024/06/11(土) モデルアーキ＋開発環境を38B with transformer engineに決定

チーム内で投票を行った上で判断し、Llama 30-50B with transformer engineの立ち上げに決定しました。この組み合わせは、2node(16GPU)で540FLOPSを記録しており、３つの候補の中で最も多くのトークンを学習できることが見込める、ということで決めました。

	1. Llama35Bの再開 (5日間学習後休止中)	2. DeepSeek v2立ち上げ	3. Llama 30-50B with transformer engineの立ち上げ
2 node学習速度(tokens/gpu/sec)	1200	1400	2520
2 node学習速度(TFLOPS)	430TFLOPS	230TFLOPS	540 TFLOPS (2nodeの結果がうまくスケールした場合)
学習できるトークン数の見込み*	800B(47+5 days: 既に学習が先行しているので､+5days)	820B (47 days)	1500 B(47 days: 条件は同上)

その後、transformer engineによる38Bモデルの事前学習を進めました。

・並行して2. 手法開発／ポストトレーニング手法検討・環境開発を推進

• 2024/06/26(水) SFT実験の手引きが完成
  この手引きの完成により、各種のデータを用いたSFTの効果測定が開始されました。
  
• 2024/06/26(水) GPT-4o評価対応のNejumi Leaderboard Neo向け評価コードをリリース
  ポストトレーニングチームは、その後の評価活動の軸となる「GPT-4o評価対応のNejumi Leaderboard Neo向け評価コード」も6/26(水)にリリースしています。
  

・事前学習、6月下旬はLoss spikeに苦しむ

6月下旬はLoss spikeに苦しみます。38Bモデル事前学習にて局所解に陥ったと思われるLoss spikeが発生し、なかなか抜け出せない状況が続きました。ここから、学習見守り・立ち上げ直しの各種対応と並行して、プロジェクトとしてこのまま38Bモデルの学習を継続すべきかどうかの吟味が続きました。

Loss spikeに陥った要因としては、学習効率を高くできるパラメータ設定を取ったことでリスクが上がったことが当時の状況から示唆されます。例えばROPEのthetaを500,000にしたり高速化のためにTransformer Engineの各種パラメータを8bitにしたことで学習の成功確率が下がった可能性や、global batch sizeを大きめの2k程度としたことで局所解に陥るリスクが上がった可能性などが考えられます。

中盤：戦略変換し、8x8B MoEモデルの学習 (6/28(金)～7/11(木))

・2024/06/28(金) 8B x Nへの戦略変換に踏み切る

前述の吟味の結果、この時点で「8B x Nを構築する」という戦略変換に踏み切りました。50Bクラスのモデルを8月初旬で構築できる、できる限り確実性の高い方法を我々なりに選んだ結果です。その戦略は以下の通りです。

1. Phase1 8Bモデルに継続事前学習を行う(1-2 T token程度)
2. 1で構築したdenseモデルをベースとして、8B x N (当時はN=4or8の２択と想定）のMoEを構築。いわゆるup cycling戦略。
3. 2で構築したMoEを継続事前学習(数百B token程度?)

この戦略に基づき、6/28(金)より学習を開始しました。この後しばらく、Lossが下がり切るまで8Bの事前学習を続けると共にMoEの準備を行いました。

・2024/07/8(月) MoEモデルアーキを8x8B(N=8)に、開発環境をllmjp環境に決定。

マージMoE特有の問題点を解決するため、MixtralをベースにSkywork-MoE論文を参考にしたゲーティングロジットの正規化を追加しました。Expert数は4x8B, 8x8B, 11x8Bの３つを比較検討、開発環境はAbeja環境とllmjp環境を比較検討しました。その結果、8x8Bモデルを、transformer engineに対応しているllmjp環境を用いて学習することに決定しました。

・2024/07/10(水) 8x8B MoE学習開始

7/10(水)、8bでnanエラーが発生したのと、tanuki-8bのlossの下がりが悪くなってきたことで、MoE化するタイミングだと判断、当時の最新checkpoint (0126000)をベースとするMoEを作成して継続学習開始しました。

・2024/07/10(水) 8Bモデルに対する本格的なSFT/DPO実験開始

6/26(水)のSFT環境・手引リリース後、それらを用いてSFTの「実験」が進みノウハウが蓄積されていました。そこにLossがある程度下がった8Bモデルができてきた、ということで、8B事前学習モデルをベースモデルとした本格的なSFT/DPOの実験が開始されました。

・2024/07/14(日) 合成データ用モデルNemotron-4 340Bが本番環境で稼働開始

Nemotron-4 340Bは、Phase2開始後の6/14(土)にNVIDIAよりリリースされた、合成データ生成に特化したLLMです。

Team Tanukiではこの新LLMを合成データ生成に生かすべく当初から取り組み、GCP本番環境とNemotron動作用Dockerとの整合取りに難航したものの、この日、遂に本番環境での推論に成功しました。こちらはその後の合成データ生成に大きく貢献することになります。

終盤：データ作成・事前学習・SFT・DPOが一体となっての総力戦 (7/12(金)～8/13(火)AM08:00)

・2024/07/12(金) JMT-Benchにおいて、8B、8x8Bの両部門で「トップスコア」を目指すことを決定

この段階で改めて、チームとして性能向上していくための方策が考察されました。

その結果、数値目標を設定してチームの力をその方向に注ぐのがよいのではないか、という結論となり、7/5(金)から１週間ほど、目標を練りました。そして7/12(金)、目標が表題の通りに決まりました。

この目標を翌日の7/13(土)全体定例にてメンバー間で共有し、この後は、性能向上が見込める活動に資源を集中しました。

・事前学習・SFT・DPOのメンバーが密に連携して総力戦を展開

この時期の悩みの一つは、論理・数学・コードの能力が伸び悩んでいたことでした。この解決策として、以下が行われました。

• 事前学習にも論理・数学・コードのインストラクションデータを投入
• 事前学習データに性質が近くてより品質の高いデータでSFT・DPOを実施

これを効果的に実行するために、事前学習・SFT・DPOのメンバーが密に連携してデータ作成・事前学習投入・その結果を受けてのSFT・DPO実行に取り組みました。

・データ構築・事前学習・SFT・DPOの取り組みが実を結び、スコア進展

8Bのスコアは、以下のように進展していきました。(gpt4-o評価)

• 7/16 4.30 (初の4点台)
• 7/24 5.3 (初の5点台)
• 7/27 6.0 (初の6点台)

8x8Bのスコアは、7/25にベースモデルで5.281を記録後、以下のように進展していきました。 (gpt4-o評価)

• 7/26 5.79
• 7/30 5.813
  

・2024/08/10(土) 総力戦総仕上げの中で、Tanukiとしての8x8B最高性能を観測

総力戦の後半は、能力が上がってきたモデルに対して、その能力を壊すことなく賢くできる可能性が高いと言われているDPO中心に試行を進めました。

延長となった8/10(土)～8/12(月)も粘り強く試行を継続し、最終的に8/10(土)、Tanuki 8x8Bとしては最高性能となる7.906を、Nejumi Leaderboard Neo, GPT-4oでの測定で観測しました。

成果

最終的には以下の４つのモデルを提出し、一定の成果を残すことができたと思います。LLMモデルの技術的な考察についてはブログ「大規模言語モデルTanuki-8B, 8x8Bの位置づけや開発指針など」、「フルスクラッチで開発した大規模言語モデルTanuki-8B, 8x8Bの性能についての技術的な詳細」をご参照ください。

■ LLM

名称	総パラメータ数(アクティブパラメータ)
Tanuki-8B-dpo-v1.0	7.5B
Tanuki-8x8B-dpo-v1.0	47B(13B)

■ Vision(LLaVA)

名称	総パラメータ数(アクティブパラメータ)
Tanuki-8B-vision-v4-checkpoint-18000	7.5B
Tanuki-8x8B-vision-v0-stage2-ckpt900	47B(13B)

今後

他方、もう一つの目標である「10B、50BクラスLLMを開発するための知見を蓄積し共有する」のところは、チーム内でかなりの蓄積はできているものの、共有はこれからです。今後も継続的に頑張っていきたいと思います。

感想

このプロジェクトは、目標や計画を立てるのが(技術的に)とても難しかったと感じます。その理由は大雑把にいうと

• モデル／学習データ／学習方式の組み合わせから、学習の成否や得られる性能を予測するのが困難
  だからだと思います。

そのような予測困難性がある中、このプロジェクトでは

• どんどんアイデアが出てきて
• アイデアが出ると誰かが試行し
• 試行結果から「うまくいく学習方法」や「性能向上の方法」を見出し
• 実現していく
  というサイクルが、ほとんどタイムラグなく素早く行われていました。

これを、常に高いモチベーションを保って粘り強く開発を進めたメンバーの皆さん、尊敬してます m(_)m

謝辞

このような機会に携わらせて頂き、本当にありがとうございました。関係者の皆様に深く御礼申し上げます。皆さんとご一緒できたことは一生の財産です。

この成果は、NEDO（国立研究開発法人新エネルギー・産業技術総合開発機構）の助成事業「ポスト５Ｇ情報通信システム基盤強化研究開発事業」（JPNP20017）の結果得られたものです。