Generative AI has the ability to create entirely new content that traditional machine learning (ML) methods struggle to produce. In the field of natural language processing (NLP), the advent of
large language models (LLMs)
specifically has led to many innovative and creative AI use cases. These include customer support chatbots, voice assistants, text summarization and translation, and more—tasks previously handled by humans.
LLMs continue to evolve through various approaches, including increasing the number of parameters and the adoption of new algorithms like Mixture of Experts (MoE). The application and adaptation of LLMs are anticipated across many industries, including retail, manufacturing, and finance.
However, many models that currently top the LLM leaderboard show insufficient understanding and performance in non-English languages, including Japanese. One of the reasons for this is that the training corpus contains a high proportion of English data. For example,
only 0.11% of the GPT-3 corpus is Japanese data
. Creating LLM models that perform well in Japanese, which has less training data than English, has been immensely challenging.
This post presents insights gained from training an AI model with 172 billion parameters as part of the
Generative AI Accelerator Challenge (GENIAC)
project, using
NVIDIA Megatron-LM
to help address the shortage of high-performance models for Japanese language understanding.
LLM-jp initiatives at GENIAC
The
Ministry of Economy, Trade and Industry (METI)
launched GENIAC to raise the level of platform model development capability in Japan and to encourage companies and others to be creative. GENIAC has provided computational resources, supported matching with companies and data holders, fostered collaboration with global technology companies, held community events, and evaluated the performance of the developed platform models.
The
LLM-jp
project to develop a completely
open model with 172 billion parameters
(available on Hugging Face) with strong Japanese language capabilities was selected for the GENIAC initiative. LLM-jp 172B was the largest model development in Japan at that time (February to August 2024), and it was meaningful to share the knowledge of its development widely.
LLM-jp is an initiative launched by researchers in the field of natural language processing and computer systems, mainly at NII, to accumulate know-how on the mathematical elucidation of training principles, such as how large-scale models acquire generalization performance and the efficiency of learning, through the continuous development of models that are completely open and commercially available. The objective is to accumulate know-how on the efficiency of training.
Training the model using NVIDIA Megatron-LM
Megatron-LM
serves as a lightweight research-oriented framework leveraging
Megatron-Core
for training LLMs at unparalleled speed. Megatron-Core, the main component, is an open-source library that contains GPU-optimized techniques and cutting-edge system-level optimizations essential for large-scale training.
Megatron-Core supports various advanced model parallelism techniques, including tensor, sequence, pipeline, context, and MoE expert parallelism. This library offers
customizable building blocks
, training resiliency features such as
fast distributed checkpointing
, and many other innovations such as
Mamba-based hybrid model training
. It’s compatible with all NVIDIA Tensor Core GPUs, and includes support for
Transformer Engine (TE)
with FP8 precision introduced with
NVIDIA Hopper architecture
.
Model architecture and training settings
Table 1 provides an overview of the model architecture for this project, which follows
Llama 2 architecture
.
Parameter
Value
Hidden size
12288
FFN intermediate size
38464
Number of layers
96
Number of attention heads
96
Number of query groups
16
Activation function
SwiGLU
Position embedding
RoPE
Normalization
RMSNorm
Table 1. Overview of LLM-jp 172B model architecture
The LLM-jp 172B model is being trained from scratch using 2.1 trillion tokens of a multilingual corpus developed for the project consisting mainly of Japanese and English. The training is performed using NVIDIA H100 Tensor Core GPUs on Google Cloud A3 Instance with FP8 hybrid training using the Transformer Engine. Megatron-Core v0.6 and Transformer Engine v1.4 are used in the experiment.
Table 2 shows hyperparameter settings for training.
Parameter
Value
LR
1E-4
min LR
1E-5
LR WARMUP iters
2000
Weight decay
0.1
Grad clip
1.0
Global batch size
1728
Context length
4096
Table 2. Hyperparameters used for the model training
In addition, z-loss and batch-skipping techniques, which are used in
PaLM
, are incorporated to stabilize the training process, and flash attention is used to further speed up the training process.
To view other training configurations, please see
llm-jp/Megatron-LM
.
Training throughput and results
Pretraining for the latest LLM-jp 172B model is currently underway, with periodic evaluations every few thousand iterations to monitor training progress and ensure successful accuracy results on Japanese and English downstream tasks (Figure 1). So far, over 80% is complete, of the targeted 2.1 trillion tokens.
Figure 1. Training loss curves for pretraining with 1.7 trillion tokens using Megatron FP8 hybrid training
Notably, there is a sharp increase in TFLOP/s after approximately 7,000 iterations, corresponding to the transition from BF16 to FP8-hybrid precision. In this experiment, BF16 plus TE was used for training before 7,000 iterations, and FP8 hybrid plus TE was used after 7,000 iterations. In Megatron-LM, it is possible to enable hybrid FP8 training with the simple option
--fp8-format
‘
hybrid
‘. Note that this feature is experimental, with further optimizations coming soon.
Figure 2. Training throughput (TFLOP/s) when TE is used with BF16 and FP8 hybrid
The reason we started the training with BF16 plus TE and then switched to FP8 hybrid was not only to see the tokens/sec performance difference between BF16 and FP8, but also to make the initial training more stable. In the early stages of training, the learning rate (LR) increases due to the warm-up, leading to unstable training.
We chose to perform the initial training with BF16, and after confirming that there were no problems with the values of training loss, optimizer states, gradient norm, and so on, we switched to FP8 to speed up the training process. FP8 hybrid has improved the training speed. We observed a training speed of 545-553 TFLOP/s with Megatron-LM.
Figure 3. Weak scaling performance based on the results of the main and preliminary experiments of the LLM-jp 172B model training
Conclusion
As mentioned above, the training of LLM-jp 172B  is still ongoing using Megatron-LM. Based on the evaluation results of downstream tasks using the current checkpoint data, we suppose that the model has already acquired excellent Japanese language capabilities, but the complete model is expected to be ready early next year.Training time is often a significant challenge in pretraining LLMs, where vast datasets are required.Therefore, efficient training frameworks like Megatron-LM are crucial for accelerating Generative AI research and development.　For the 172B model trained with
Megatron-LM
, we explored FP8-hybrid training as a potential method for improving training speed, achieving a 1.4x training speed acceleration from 400 TFLOP/s to 550 TFLOP/s. We observed a performance acceleration from 400 TFLOP/s to 550 TFLOP/s, suggesting that FP8-hybrid could be a valuable approach for enhancing the efficiency of large-scale model pretraining.
Reading Time:
2
minutes
生成 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
アーキテクチャに準拠したこのプロジェクトのモデル アーキテクチャの概要です。
パラメーター
値
Hidden size
12288
FFN Intermediate size
38464
Number of layers
96
Number of attention heads
96
Number of query groups
16
Activation function
SwiGLU
Position embedding
RoPE
Normalization
RMSNorm
表 1. LLM-jp 172B モデルアーキテクチャ概要
この 172B モデルは、プロジェクト用に開発された多言語コーパス (主に日本語と英語) の 2.1T トークン (2.1 兆トークン) を使用してゼロから学習されています。学習は、Transformer Engine を使用した FP8 ハイブリッド学習で、Google Cloud の A3 インスタンス上の H100 Tensor コア GPU を使用して実行されています。実験では、
Megatron-Core
v0.6 と
Transformer Engine
v1.4 が使用されています。
学習のハイパーパラメーター設定は次のとおりです。
パラメーター
値
LR
1E-4
min LR
1E-5
LR WARMUP iters
2000
Weight Decay
0.1
Grad Clip
1.0
global batch size
1728
context length
4096
表 2. この実験で使用したハイパーパラメーターの概要
詳細な設定に興味のある方は、他の学習設定も
llm-jp/Megatron-LM
でご覧いただけます。
また、
PaLM
で採用されている　z-loss や batch-skipping テクニックを取り入れることで学習プロセスを安定化させ、flash attention を利用することで学習プロセスをさらに高速化させています。
学習結果とスループット
LLM-jp 172B モデルの事前学習は、数千イテレーションごとに、日本語と英語の下流タスクの評価結果をモニターし、学習がうまく進んでいるかどうかを確認しながら現在も進行中です。これまでのところ、目標とする 2 兆 1,000 億トークンの 80% 強まで完了しています。
Megatron の FP8 ハイブリッド学習を用いた 1.7 兆トークンの事前学習における学習損失曲線を以下に示します。この曲線は、240,000 ステップまで損失が着実に減少していることを示しています。
図 1. 240k ステップまでの学習ロス
以下のグラフは、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 ハイブリッド学習を有効にすることができます。
図 2. Transformer Engine を BF16 と FP8 のハイブリッドで使用した場合の学習スループット (TFLOP/s)
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 の学習結果は、線形スケーリングを示していることが分かります。
図 3. 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 が大規模モデルの事前学習の効率を向上させる有用なアプローチであることを強調しています。