Llama 3.3 Swallow: 高性能な日本語LLM

はじめに

本日(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利用規約の利用制限を遵守する必要があります。このライセンスおよび利用制限に従う限りにおいては、研究および商業目的での利用が可能です。

公式プロジェクトページはこちらです。
https://swallow-llm.github.io/llama3.3-swallow.ja.html
本モデルの開発は、産総研、東京科学大学 岡崎研究室、横田研究室の合同プロジェクトにて行われました。

このブログでは、公式ページでは触れることが出来なかった開発に関する詳細について紹介します。

今回のモデル開発においてAWSの各種サービスを利用するにあたり、AWSの柳生様、佐々木様、渡辺様、宮本様、井阪様には様々なサポートをいただきました。この場をお借りして感謝申し上げます。

リリースモデル

リリースモデルの性能に関しては公式プロジェクトページを参照ください。
https://swallow-llm.github.io/llama3.3-swallow.ja.html

また、Swallowで評価しているLLMの評価結果については以下を参照ください。

ベースモデル

https://swallow-llm.github.io/leaderboard/index-base.ja.html

チャットモデル

https://swallow-llm.github.io/leaderboard/index-chat.ja.html

モデルサイズ別

5B-15B
https://swallow-llm.github.io/leaderboard/category-5b-15b-base.ja.html

15B-100B
https://swallow-llm.github.io/leaderboard/category-15b-100b-base.ja.html

学習環境

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から確保するという仕組みを採用しています。

大規模学習においては、GPUノードの不良がつきものだと思いますが、32ノードでは予備ノードがないのではありませんか？

はい、おっしゃるように不良ノードの発生は避けられません。しかし、今回利用したHyperPod環境では以下の2つの理由から予備ノードを確保せずに学習を行うことが出来ました。

1. 今回利用した環境では、scontrol updateコマンドにより10-30分程で不良ノードを交換することが可能でした。そのため、学習を長時間止めることなく、32ノードで学習するジョブを流すことが可能でした。
2. --auto-resume=1フラグを利用することで、自動的にfailedしたジョブを再開することが可能となっているためです。(今回は、使用した独自build PyTorchとの関係性から利用しませんでした)

予備ノードを確保せずに学習を行うことが出来たため、同じ費用でより長い時間インスタンスを借りることが出来ました。

学習環境の整備、運用についてはどのような体制で行っていますか？

AWSの方々のサポートの下、私が1人で行っています。
sagemaker hyperpodのデプロイはこちらのsamplesを参考に、主にAWSの渡辺さんのサポートの下行いました。

ABCIやTSUBAMEでの学習においても、ライブラリのメンテナンス、必要なEnvironment Modulesの整備、Pythonの必要ライブラリの環境整備、学習の実行、不良ノード対応、学習の進捗管理は、1人で担当していますので、今回の学習では、これに学習環境のデプロイや調達作業が追加された形です。

Sagemaker HyperPod環境整備

Permissionの設定

デプロイには、AdministratorAccessが必要です。AWSアカウントの管理者に頼み、以下のようにAdministratorAccessを付与してもらってください。

CLIツールのインストール

まず、お使いの環境に合わせてAWS CLIをインストールします。
https://docs.aws.amazon.com/ja_jp/cli/latest/userguide/getting-started-install.html

sudoが利用できない環境でaws cliをインストールしたいのですが、どうすれば良いでしょうか？

以下にLinux x86-64の例を示します。
次のようにコマンドを打ち、~/binにパスを通すようように~/.bashrcまたは~/.zshrcを変更すれば利用可能です。

curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"
unzip awscliv2.zip

mkdir -p ~/local/aws-cli
mkdir -p ~/bin

./install -i ~/local/aws-cli -b ~/bin

また、デプロイ後にClusterにSSHするには、SSM Session Manager Pluginが必要です。
以下のドキュメントを参考にお使いの環境に合わせてインストールをしてください。
https://catalog.workshops.aws/sagemaker-hyperpod/en-US/01-cluster/05-ssh

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に紐づけ、簡単に転送することが可能です。

直接 FSx for Lustreにデータを移送するのではダメなのでしょうか？

いいえ、問題ありません。しかし、比較的安価なS3にデータを事前に移送しておき、その後、LustreへデータをS3から移送することで、データの転送が間に合わない事態や、比較的高価なLustreをデータ転送のためだけに1週間前からデプロイするといったこともする必要がなくなるため、上記のようにData repository associationsを利用するのが良いかと思います。

データ転送作業

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>

デプロイが完了すると、どのような表示になりますか？
以下に、ap-southeast-2 リージョンに2ノードのclusterを作成した際のログを示します。

+ python3 validate-config.py --cluster-config cluster-config.json --provisioning-parameters provisioning_parameters.json --region ap-southeast-2
✔️  Validated instance group name worker-group-1 is correct ...
✔️  Validated subnet subnet-0ff95d3693a51c0c8 ...
✔️  Validated security group sg-08422ea4331b448b1 ingress rules ...
✔️  Validated security group sg-08422ea4331b448b1 egress rules ...
✔️  Validated FSx Lustre DNS name fs-0e71a66e985b69560.fsx.ap-southeast-2.amazonaws.com ...
✔️  Validated FSx Lustre mount name 57ytrbev ...
✔️  Validated provisioning_parameter.json is valid json ...
✅ Cluster Validation succeeded
+ echo 'aws --region' 'ap-southeast-2 sagemaker create-cluster --cli-input-json file://cluster-config.json'
aws --region ap-southeast-2 sagemaker create-cluster --cli-input-json file://cluster-config.json
+ [[ DRY_RUN -eq 1 ]]
+ aws --region ap-southeast-2 sagemaker create-cluster --cli-input-json file://cluster-config.json
{
    "ClusterArn": "arn:aws:sagemaker:ap-southeast-2:128132837601:cluster/xgrgfzve735w"
}

デプロイに失敗しました。どうしてでしょうか？

過去にデプロイを行ったことがある、または、同じ名前でデプロイに失敗したことがある場合は、<cluster-name>の衝突の可能性があります。変更してみてください。

また、python package周りでエラーが発生している場合は、pip install jsonschemaを行うことで解消するかもしれません。

データの実体化

以下のドキュメントに従い、S3からData Repository機能を利用しFSx for Lustreにmeta dataだけコピーされたデータにアクセスし、データをプリロードします。これにより、学習開始時にデータのロードが発生し、学習がなかなか始まらない事態を回避できます。
https://docs.aws.amazon.com/ja_jp/fsx/latest/LustreGuide/preload-file-contents-hsm-dra.html

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を行うことが可能になりました。

Asynchronous Savingについて詳細が分かる資料はありませんか？

PyTorchの以下の記事が参考になるかと思います。
https://pytorch.org/tutorials/recipes/distributed_async_checkpoint_recipe.html

また、PyTorchの実装とは異なりますが、同様のアイデアの論文がPyTorch teamがasync checkpointを開発しているのとほぼ同時期に提出されています。

https://arxiv.org/abs/2406.10707v1

学習管理

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開発に貢献できればと、私は考えています。