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 が大規模モデルの事前学習の効率を向上させる有用なアプローチであることを強調しています。