はじめに
本日(2025/03/10)、Swallow ProjectよりLlama-3.3-Swallow-70B-v0.4とLlama-3.3-70B-Instruct-v0.4をリリースさせていただきました。
今回リリースしたLlama-3.3-Swallow-70B-v0.4は日本語理解・生成タスク(アカデミックタスク)において、Qwen2.5-72Bの平均スコアをわずかに上回り、同程度の能力をベンチマーク評価上示しました。また、Llama-3.3-70B-Instruct-v0.4は、前回リリースした Llama-3.1-Swallow-70B-Instruct-v0.3よりもわずかに高いJapanese MT-Benchの性能を示しており、GPT-4oやQwen2.5-72B-Instructには及ばないものの高い性能を示しています。
公式ページより
公式ページより
Llama-3.3-Swallow-70B-v0.4のライセンスはMeta社のLlama 3.3ライセンスを継承しています。加えて、Llama 3.3ライセンスに従ったうえで、Gemma利用規約の利用制限を遵守する必要があります。このライセンスおよび利用制限に従う限りにおいては、研究および商業目的での利用が可能です。
公式プロジェクトページはこちらです。 本モデルの開発は、産総研、東京科学大学 岡崎研究室、横田研究室の合同プロジェクトにて行われました。
このブログでは、公式ページでは触れることが出来なかった開発に関する詳細について紹介します。
今回のモデル開発においてAWSの各種サービスを利用するにあたり、AWSの柳生様、佐々木様、渡辺様、宮本様、井阪様には様々なサポートをいただきました。この場をお借りして感謝申し上げます。
リリースモデル
リリースモデルの性能に関しては公式プロジェクトページを参照ください。
また、Swallowで評価しているLLMの評価結果については以下を参照ください。
ベースモデル
チャットモデル
モデルサイズ別
5B-15B
15B-100B
学習環境
Swallow LLMシリーズの学習は、産総研のABCI 2.0、東京工業大学(現 東京科学大学)のTSUBAME 4.0などのスーパーコンピュータを利用して学習を行ってきました。しかしながら、産総研のABCI 2.0が次期システムへの移行のために2024年10月末で運用を終了したこと、そして、TSUBAME 4.0の利用枠を確保できなかったことから、今回のモデルの学習は、Amazone Web Services(AWS)のSageMaker HyperPodを利用して行いました。
学習環境は以下のような構成になっています。今回の学習環境では、計算ノード(Compute Nodes)としてH100インスタンス(p5.48xlarge)を32ノード利用して学習を行いました。
(AWS SageMaker HyperPod catalogより)
使い慣れているスパコンと同様にログインノード(Head node)を設け、ログインノードにsshを行った後、ジョブスケジューラ(今回はSlurmを採用)にジョブをsubmitし、ジョブスケジューラが要求されたリソースをCompute Nodeから確保するという仕組みを採用しています。
Sagemaker HyperPod環境整備
Permissionの設定
デプロイには、AdministratorAccessが必要です。AWSアカウントの管理者に頼み、以下のようにAdministratorAccessを付与してもらってください。
CLIツールのインストール
まず、お使いの環境に合わせてAWS CLIをインストールします。
また、デプロイ後にClusterにSSHするには、SSM Session Manager Pluginが必要です。
以下のドキュメントを参考にお使いの環境に合わせてインストールをしてください。
macOS(arm64)の場合を以下に示します。
curl "https://s3.amazonaws.com/session-manager-downloads/plugin/latest/mac_arm64/sessionmanager-bundle.zip" -o "sessionmanager-bundle.zip"
unzip sessionmanager-bundle.zip
sudo ./sessionmanager-bundle/install -i /usr/local/sessionmanagerplugin -b /usr/local/bin/session-manager-plugin
FSx for Lustreのセットアップ
こちらのリンクのCloudFormationからFSx Lustreを構築しました。(注意: このリンクは、us-east-2リージョンにてデプロイするためのリンクになります。適時、それぞれの環境に合わせて変更してください。)
今回は以下のような設定でデプロイを行いました。なお、PerUnitStorageThroghputについては、想定される読み書きに応じて適切な値を設定してください。
デプロイを開始すると以下のようにEventsが生成されます。
今回の構成でのFSx for Lustreの構成完了までに約10分ほど要しました。参考程度にしてください。
データ転送準備
デプロイされたLustreにアクセスし、以下のようにData repositoryを選択します。Data repository associcationsを作成することで、S3にアップロードしていたデータをLustreに紐づけ、簡単に転送することが可能です。
データ転送作業
aws s3コマンドによるデータ転送において、デフォルト値のままでは転送速度が遅く、データの移送が間に合わない計算であったため、こちらを参考にmax_concurrent_requests、max_queue_sizeなどを変更しました。
また、Swallowで利用するデータサイズは3TB以上存在するため、これらを移送するために15個のジョブを作成し、それぞれ別々のデータセットを担当させるようにしました。これにより、TSUBAMEからS3への転送を1日未満で行うことができました。
HyperPod Clusterのデプロイ
easy-steup.shを利用し、以下のようにデプロイを行いました。
bash easy-setup.sh --region <region> \
--instance-type p5.48xlarge \
--instance-count <node-count> <cluster-name>
データの実体化
以下のドキュメントに従い、S3からData Repository機能を利用しFSx for Lustreにmeta dataだけコピーされたデータにアクセスし、データをプリロードします。これにより、学習開始時にデータのロードが発生し、学習がなかなか始まらない事態を回避できます。
find <data/path> -type f -print0 | xargs -0 -n 1 -P 8 sudo lfs hsm_restore
学習設定
ハイパーパラメータ
継続事前学習におけるハイパーパラメータ
継続事前学習において利用したハイパーパラメータと分散学習設定は以下のようになります。
LRに関しては、1.25E-5, 1.50E-5, 1.75E-5でablation実験を行い、学習安定性や下流タスクでの評価から1.25E-5に決定しました。しかし、50B Token学習時点では、どのLRを選択しても大きな差は存在しませんでした。
| LR | min_LR | LR warmup iter | weight decay | grad clip |
|---|---|---|---|---|
| 1.25E-5 | 1.25E-6 | 1000 | 0.1 | 1.0 |
| global batch size | micro batch size | DP | TP | CP | PP |
|---|---|---|---|---|---|
| 1024 | 1 | 16 | 8 | 1 | 2 |
| adamw beta1 | adamw beta2 | adamw eps |
|---|---|---|
| 0.90 | 0.95 | 1E-8 |
| attention dropout | hidden dropout | RMSNorm eps | RoPE base |
|---|---|---|---|
| 0.0 | 0.0 | 1E-5 | 500000 |
分散学習設定は、後述する memory消費量予測式を利用し、学習可能な分散学習設定をすべて考慮しました。学習可能な設定の中で、AWS EFAを利用したことによるPyTorch管理外のbuffer memory消費量増加に伴いOOMにより学習ができない設定を除き、最終的に最も高速なものを採用しました。
指示チューニングにおけるハイパーパラメータ
指示チューニングモデルを学習する際に利用したハイパーパラメータは以下のようになります。
指示チューニングのハイパーパラメータは、以前にLlama-3.1-Swallow-70B-Instruct-v0.1を開発した際にablationした結果を踏襲しており、再度ハイパーパラメータチューニングを行うことは、計算コストの観点から行っていません。
| LR | min_LR | LR warmup ratio | weight decay | grad clip |
|---|---|---|---|---|
| 1.75E-5 | 1.75E-6 | 0.1 | 0.1 | 1.0 |
| global batch size | micro batch size |
|---|---|
| 256 | 1 |
| adamw beta1 | adamw beta2 | adamw eps |
|---|---|---|
| 0.90 | 0.95 | 1E-8 |
また、今回の指示チューニングでは、先行研究 Nemotron-4-340Bで有効性が示されたStaged Supervised Fine-tuningを採用しています。これは、コードSFTデータとGeneralなSFTデータを同時に学習するのでなく、coding dataで最初にSFTを行い、その後、coding dataを少量含むGeneralなデータでSFTを行うことで、コーディングタスクを含むdownstreamタスクを上昇させることができるという知見です。
そのため、LRの推移は以下のようになります。
Swallow Projectでは、コーディングに特化した高品質なSFTデータの用意が2024年1月まで行えておらず、この先行研究の有効性を検証することが出来ていませんでした。しかし、Swallow-Code-v0.3(プロジェクトページ参照のこと)として高品質なコードコーパスが作成できたため、このデータからInstruct形式にデータを加工し、ある程度の品質を有するコードSFTデータを用意することができました。今回は、この合成コードSFTデータを利用し、2段階のSFTとmixed SFT(1 staged SFT)を比較し有効性を検証しました。
| Training Strategy | coding | extraction | humanities | math | reasoning | roleplaying | stem | writing | 1st turn | 2nd turn | JMT Avg | Ja char ratio |
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| 日本語+英語データ (General) | 0.513 | 0.727 | 0.821 | 0.533 | 0.500 | 0.746 | 0.675 | 0.773 | 0.718 | 0.604 | 0.661 | 0.695 |
| 日本語+英語データ (General) + コード (Coding) | 0.578 | 0.713 | 0.815 | 0.543 | 0.518 | 0.774 | 0.666 | 0.769 | 0.737 | 0.606 | 0.672 | 0.728 |
| 2 Staged SFT (コード -> General) | 0.549 | 0.732 | 0.826 | 0.576 | 0.549 | 0.786 | 0.695 | 0.797 | 0.756 | 0.621 | 0.689 | 0.723 |
Japanese MT-Benchにて評価したところ上記のように、codingを除く全てのタスクでmixする設定よりも2 staged SFTの方が有効である結果が示されました。また、JHumanEval、HumanEvalについては、2-staged SFTによりそれぞれ4ポイント強上昇することを確認しました。
メモリ消費量予測
Swallow Projectでは、2024年11月以降、学習時のGPUあたりのメモリ消費量を正確に予測し、学習可能な設定をすべて考慮した上で最速の学習設定を半自動的に決定するツールにより示された分散学習設定を利用して実験を行っています。
Llamaアーキテクチャを4D Parallelismで学習する際のメモリ消費量は、私の過去の研究から見積もれるため、これを利用し、学習するモデル、そしてGPU環境を指定するだけで、探索するべき設定と予想されるPyTorchが消費するメモリ消費量を表示する独自ライブラリを作成しました。
なお、このライブラリは将来的に使い方を含め、公開する予定です。
(memory消費量予測ライブラリ(自作)により自動生成)
communication overlap
Swallow Projectでは、Data Parallelismにおける通信と計算のOverlapだけでなく、Tensor Parallelismにおける通信と計算のOverlapを有効化しています。これにより、通信待ちにより計算をブロッキングすることを最小限に抑え、効率的な学習を実現しています。
Communication and Computation Overlap(通信と計算のOverlap)で、なぜ学習を高速化することができるのかを模式的に示すと以下のようになります。
完全に直列に通信と計算が並んでいる上段の図を下段のように変更し、一部のbackwardが終了し通信をしても問題なくなった時点で通信を行うことで、大幅に同様の計算を行うまでの時間を短縮することができています。
PyTorch DCPの利用
checkpointを保存する時間を削減するためにDistributed Checkpoint(DCP)を利用しました。
従来のMegatron-LMでは、Tensor Parallel size、Pipeline Parallel sizeをそれぞれ、tp_size, pp_sizeとおくと、tp_size x pp_size個のプロセスからcheckpointを保存していました。そのため、大量のGPUを利用していても、checkpointを保存するためにwriteを行うプロセス数が限られ、分散ストレージの性能を十分に利用し尽くせていない事態が生じていました。
そこで、DCPを利用することで、使用しているGPU数(=プロセス数)からcheckpoint saveを行うことで、checkpointの保存(save)と読み込み(load)をより並列化して行うことが可能になりました。
また、加えて、DCPのAsynchronous Saving(非同期保存)機能を利用することで、さらにcheckpoint save時間を削減しました。これまでは、checkpoint saveのために学習を完全にブロックしていましたが、async save機能により、最低限のブロッキングのみで、学習をすぐに再開できるようになりました。これにより、2024年7月にリリースしたLlama-3-Swallow-70B-v0.1を学習したときと比較して1/10未満の時間でcheckpoint saveを行うことが可能になりました。
学習管理
Swallow Projectにおける継続事前学習、及び指示チューニングの実験管理体制について紹介します。
まず、利用しているツールは主にSlack Canvases、Google SpreadSheet、Weights & Biasesになります。学習の計画段階から、不測の事態への対応まで順に説明を行います。
学習計画表
まず、実験計画を立てるために、前もって、次の実験サイクルにて使用する事前学習ライブラリ、指示チューニングライブラリのバージョンを固定化します。(バージョンが切られていないライブラリを使用している場合は、branchやtag等で対応します)
使用するバージョンを決め次第、学習で使用する可能性があるGPUノード数にて短時間の学習速度測定を行い、メモリ消費量予測ライブラリから探索する必要があると示された分散学習設定の中から最速の設定を探します。また、調達済みの計算資源量から、どの程度のモデルサイズの学習が何B Token可能なのか、そしてリソースを占有している場合は、予備ノードと計算用のノードの割当とスケジューリングを行います。この際に、実測したtokens/secを利用し、学習にかかる時間を算出しています。
次に、Swallow Projectの合同会議にて提案する実験の草案をSlackのCanvasに記す、もしくは、合同チャンネルに実験案を提起します。急ぎの提案でない限り、毎週の合同会議にて、内容が議論され、実験を行うか否か、またその規模について決定されます。(規模について議論するために、前述の作業は欠かせません。実行可能な実験数を正確に把握せず、行き当たりばったりな実験をしていては、GPUノードが余ってしまい、無駄が生じてしまったり、GPUノードの予約等が必要な場合は、予約を行うのが直前になってしまい、課金係数が高くなり予算を無駄に消化してしまう事態につながります。このような事態に回避するためにも、ライブラリのメンテナンス作業と最速設定の探索、そして実測を実験サイクルにあわせて行う必要があります。)
行う実験数と実験規模が決定したため、事前に作成していたスケジューリングを決定した実験をふまえて変更します。この際、以下のようにGoogle SpreadSheet上に実験スケジュールを大まかに記載します。これにより使用するノード数と日数が一目で分かるようになります。
(実験スケジュール管理表 (Google SpreadSheet))
また、上記の実験スケジュールは1日単位の大まかなスケジュール管理しかできないため、実測値から学習が終了する時間を割り出し、実験管理担当である私のカレンダーに実験名と終了時刻を予定として記録します。(また、約半日ごとに学習トーク数を確認し、想定通り実験が推移しているか、スケジュールに変更が必要でないかを判断しています。)
実験準備
前述の実験管理表をGoogle SpreadSheetに作成でき次第取り掛かるのが、実験準備です。
新しいコーパスを試す場合は、コーパスのトークナイズ作業から行う必要があります。トークナイズ作業は他のメンバーの方に頼む場合もあれば、私が行うこともありますが、トークナイズをCPUノードにて行います。
トークナイズ作業が完了次第、学習トークン数を確認し、Google SpreadSheetのトークン管理表に数値を入力します。また、コーパスの作成者とトークナイズ作業を行った人間が別の場合は、コーパス作成者にトークン数について報告を行い、想定の範囲内か確認を行います。
(実験ごとのトークン数内訳の管理(Google SpreadSheetより))
次に実験計画に従い、それぞれのコーパスごとに学習するトークン数を決定します。この際も、各実験ごとにSpreadSheet上に表を作成し、総トークン数が想定通りか、日、英、コード比率は想定通りか、などを確認します。実験発案者と準備作業担当者が別の場合は、このタイミングで認識に齟齬がないか再度調整を行います。
そして最後に、実験スクリプトに情報を落とし込みます。
spreadsheetの情報から実験スクリプトの一部を自動作成するツールを作成しているため、これによりコーパスごとの学習トークン数等を自動で埋めます。次に、実験のハイパーパラメータ等を埋めていきます。学習スクリプトが完成したら、一晩あけて再度、実験スクリプトと実験案に相違がないか確認します。その後、小規模なノードで動作確認を行い、ジョブスケジューラに投入しておけば、リソースが確保され次第自動で流れることを確認して、実験準備作業は完了です。
不足の事態への対応
ジョブが更新なく停止、もしくはexitした場合は、実験管理者のスマホに通知が飛ぶように設定しています。ジョブの更新が停止した場合に通知を行う機能についてはWeights & Biasesの標準機能としてサポートされていないため独自スクリプトを常時稼働することで対応していますが、ジョブが正常終了、もしくは異常終了した際の通知は、Weights & Biasesの標準的な機能としてメール通知または、Slackへの通知が選択できるため、Slackへの通知をONにしています。
また、SlackbotチェンネルにWeights & Biasesの通知を飛ばすように設定しておくと他の通知と混ざらず便利かと思います。
これにより、外出時等でPCにアクセスできない時間を除き、深夜帯であろうと常に実験管理者に通知が飛び、すぐに対応できる体制を構築することができます。(精神的に疲弊するため、あまりオススメしませんが...)
今後について
2023年12月にLlama-2から継続事前学習を行ったSwallowシリーズをリリースして早1年が経過しました。日本語理解・生成タスク(アカデミックタスク)では、GPT-4oに近い水準に到達するなど、評価結果上は非常に高い性能を有するLLMを継続事前学習にて開発することが可能になったという意味では非常に喜ばしいことだと思います。しかし、instructモデルと対話すると、GPT-4o, Claude 3.7Sonnet, DeepSeekなど最先端モデルの性能には、届いていないことが感じ取れます。特にReasoningタスクと呼ばれるような複雑な推論課程を要するタスクにおいては、Swallow Projectが開発したモデルでは十分に対応できていません。
Swallow Projectは当初の"日本語に強いLLMを開発する"という目標をアカデミックタスクに関しては、相当程度完了しているため、今後の研究開発は、これまでとは異なった、よりチャレンジングな研究開発を行っていくことになるかと考えています。
しかし、Swallow Projectが使用可能な計算資源量は、Google, OpenAI, Anthropicなどの企業と比べるとはるかに少ないですし、日本国内のLLM開発組織として代表的なPFE(PFN)、SB Intutionsなどと比べても1/10未満の計算資源量しか有していません。そのため、from scratchからの事前学習において何か競争力のあるモデルを開発したりすることは、Swallow Projectでは不可能であり、計算資源の制約から、継続事前学習、annealing学習、Post-Training(SFT、DPO、強化学習など)が研究対象になるかと思います。
from scratchからの学習により競争力のあるLLMを開発されているPFE(PFN)、SB Intutiions等の企業の取り組みを少しでもサポートできるような研究を行うことで、今後の日本国内におけるLLM開発に貢献できればと、私は考えています。
Discussion