Megatron-LM を用いた日本語に強い 172B 大規模言語モデルの開発

生成 AI は、その卓越した能力のおかげで、従來の機械學習手法ではできなかったタスクを実行し、注目を集めています。例えば、自然言語処理の分野では、大規模言語モデル (LLM) が登場したことで、チャットボットによるカスタマー サポートや會議內容の要約など、これまで人間が擔っていた役割を AI が代わりに行うなど多くの革新的で創造的なユース ケースが生まれています。

LLM は、パラメーター數の増加や MoE (Mixture of Experts) のような新しいアルゴリズムの採用など、様々なアプローチを通じて進化し続けており、小売業、製造業、金融業など、さまざまな業界への応用と適用が期待されています。

しかし、現在 LLM リーダーボードの上位モデルの多くは、英語に比べて日本語の理解度やパフォーマンスが低い傾向にあります。その理由の一つは、學習コーパスの英語データの割合が大きいことです。例えば、GPT-3 の場合、日本語データはコーパスの 0.11% しかありません。日本の生成 AI の発展のためには、非常に困難ですが英語よりも學習データの少ない日本語で優れた性能を発揮する LLM モデルを作成することが、乗り越えるべき重要な課題です。

本稿では、GENIAC (Generative AI Accelerator Challenge) プロジェクトの一環として取り組んだ、Megatron-LM を用いた 172B 大規模言語モデルの學習から得られた知見を紹介し、データ不足の問題を乗り越えて日本語理解能力の高いモデル作成に取り組んだ際の洞察について紹介します。

GENIAC における LLM-jp の取り組み

上記で述べたような課題を解決するために、経済産業省は、日本國內のプラットフォーム モデル開発力の向上と企業等の創意工夫を奨勵するため、「Generative AI Accelerator Challenge (GENIAC)」を立ち上げました。GENIAC では、計算資源の提供、企業とデータ保有者とのマッチング支援、グローバルテック企業との連攜促進、コミュニティ イベントの開催、開発されたプラットフォーム モデルの性能評価などを現在も継続して行っています。

この取り組みに、LLM-jp の日本語対応力に優れた完全オープンな 172B モデルの開発というテーマが選ばれました。172B は當時 (2024 年 2 月から 8 月) 日本國內で最大規模のモデル開発であり、その開発ノウハウを広く共有することは非常に有意義なことでした。

LLM-jp は、NII (國立情報學研究所) を中心とした自然言語処理や計算機システム分野の研究者が中心となって立ち上げた取り組みで、大規模モデルが汎化性能を獲得する仕組みや學習の効率性といった學習原理の數學的解明に関するノウハウを、完全にオープンで商用利用可能なモデルの継続的な開発を通じて蓄積する事、および學習の効率性に関するノウハウを蓄積することを目的としています。

NVIDIA Megatron-LM

Megatron-LM は、Megatron-Core を活用して大規模言語モデル (LLM) を比類のない速度で學習する軽量な研究指向フレームワークとして機能します。主要コンポーネントである Megatron-Core は、大規模な學習に不可欠な GPU 最適化技術と最先端のシステム レベルの最適化を含むライブラリです。

Megatron-Core は、テンソル、シーケンス、パイプライン、コンテキスト、MoE エキスパート並列処理など、さまざまな高度なモデル並列処理手法をサポートしています。このライブラリは、カスタマイズ可能なビルディング ブロック、高速分散チェックポイントなどの學習回復力機能、Mamba ベースのハイブリッド モデル學習などの他の多くのイノベーションを提供します。すべての NVIDIA Tensor コア GPU と互換性があり、NVIDIA Hopper アーキテクチャで導入された FP8 精度の Transformer Engine テクノロジのサポートが含まれています。

上記のような最先端の機能を提供することで、Megatron-LM は、研究者やモデル開発者が 1,000 億のパラメーターを超えるモデルでも、高速な學習と數千の GPU スケールへのスケーラビリティを実現できるようにします。

モデル アーキテクチャと學習設定

以下は、Meta の Llama2 アーキテクチャに準拠したこのプロジェクトのモデル アーキテクチャの概要です。

この 172B モデルは、プロジェクト用に開発された多言語コーパス (主に日本語と英語) の 2.1T トークン (2.1 兆トークン) を使用してゼロから學習されています。學習は、Transformer Engine を使用した FP8 ハイブリッド學習で、Google Cloud の A3 インスタンス上の H100 Tensor コア GPU を使用して実行されています。実験では、Megatron-Core v0.6 と Transformer Engine v1.4 が使用されています。

學習のハイパーパラメーター設定は次のとおりです。

詳細な設定に興味のある方は、他の學習設定も llm-jp/Megatron-LM でご覧いただけます。

また、PaLM で採用されている　z-loss や batch-skipping テクニックを取り入れることで學習プロセスを安定化させ、flash attention を利用することで學習プロセスをさらに高速化させています。

學習結果とスループット?

LLM-jp 172B モデルの事前學習は、數千イテレーションごとに、日本語と英語の下流タスクの評価結果をモニターし、學習がうまく進んでいるかどうかを確認しながら現在も進行中です。これまでのところ、目標とする 2 兆 1,000 億トークンの 80% 強まで完了しています。

Megatron の FP8 ハイブリッド學習を用いた 1.7 兆トークンの事前學習における學習損失曲線を以下に示します。この曲線は、240,000 ステップまで損失が著実に減少していることを示しています。

以下のグラフは、Y 軸は TFLOP/s、X 軸はイテレーション回數を示しています。注目すべきは、學習を BF16 から FP8 ハイブリッドへ切り替えた約 7,000 回のイテレーションのタイミングで、TFLOP/s が急激に増加していることです。この実験では、BF16 + Transformer Engine が 7,000 回以前の學習に使用され、FP8 ハイブリッド + Transformer Engine が 7000 回以降の學習に使用されました。Megatron-LM では、単純なオプション --fp8-format ‘hybrid‘ で FP8 ハイブリッド學習を有効にすることができます。

BF16 + Transformer Engine で學習を開始し、その後 FP8 ハイブリッドに切り替えた理由は、BF16 と FP8 でどの程度の tokens/sec の性能差があるかを見るためだけでなく、初期學習をより安定させるためでもあります。學習の初期段階では、ウォームアップにより學習率 (LR) が上昇し、學習が不安定になります。そこで、初期學習は BF16 で行い、學習損失の値、オプティマイザーの狀態、勾配ノルムなどに問題がないことを確認した後、FP8 に切り替えて學習を高速化することにしました。FP8 ハイブリッドにより學習速度が向上していることが分かります。

私たちは、最終的に Megatron-LM を用いて 545-553TFLOP/s の性能を達成できることを確認しました。以下は、LLM-jp 172B モデル學習の本実験と予備実験の結果に基づく、弱スケーリング性能のグラフです。このグラフでは、Y 軸が Aggregate Throughput を表し、X 軸が學習に使用した GPU の數を表しています。Llama2 7B、Llama2 13B、LLM-jp 172B の學習結果は、線形スケーリングを示していることが分かります。

まとめ

前述の通り、LLM-jp 172B の學習は現在も Megatron-LM を用いて進行中です。現在のチェックポイント データを用いた下流タスクの評価結果から、現狀でも既に優れた日本語能力を獲得していると推察されますが、完全なモデルは來年の初めに完成予定です。

膨大なデータセットを必要とする大規模言語モデルの事前學習において、學習時間はしばしば大きな課題となります。そのため、Megatron-LM のような効率的に學習可能なフレームワークは、生成 AI の研究開発を加速させる為に非常に重要です。Megatron-LM で學習した 172B モデルにおいて、FP8-hybrid 學習が學習速度を向上させる効果的な手法であることを実証し、1.4 倍の高速化 (400 TFLOP/s → 550 TFLOP/s) を達成しました。この結果は、FP8-hybrid が大規模モデルの事前學習の効率を向上させる有用なアプローチであることを強調しています。