ms-swiftによるMegatron-LMベースのQwen3のファインチューニング
はじめに
本記事は執筆時点（2025/05/10）の情報に基づいています。閲覧時点では情報が古い可能性があるので、公式ドキュメント等を併せてご確認ください。
この記事では、ms-swiftというフレームワークを使ってMegatron-LMをベースにLLMをファインチューニングする方法について解説します。
公式のリポジトリはこちらです。
経緯
4月末にQwen3シリーズがリリースされました。これを手元でファインチューニングしようとして海外の先行事例等を探していた際、以下のベストプラクティスを発見しました。
どうやらAlibaba Cloudが開発を主導しているms-swiftというフレームワークによるQwen3のファインチューニングガイドのようです。QwenもまたAlibabaの開発なので、相性が良さそうだなと思ってこの学習フレームワークを採用しました。
学習の結果出来上がったモデルはこちらです。
今回の記事は基本的に上記のベストプラクティス記事に基づいています。
概要
ms-swiftは、Alibaba CloudがHugging Face競合として運営するサービスであるModelScopeのコミュニティによって開発されている、LLMのファインチューニング・デプロイ用フレームワークです。
継続事前学習やSFT、DPO、ORPO、PPO、GRPOなど様々な学習手法に対応しています。
似たような学習フレームワークはaxolotlやLLaMA-Factoryなど様々なものがありますが、これらとの大きな違いとして、ms-swiftはMegatron-LMをバックエンドにしたファインチューニングに対応しています（Megatron-SWIFTと呼ばれているようです）。これにより、一部のモデルにおいて通常のHuggingface transformersをバックエンドとした学習よりも高速・大規模に学習を行うことができます。
特にMoEアーキテクチャのモデルにおいて高速化の恩恵が大きいようです。以下は公式ドキュメントに記載のあるいくつかのパフォーマンス比較表です。
Qwen3-30B-A3B MoE — NVIDIA A100 80 GB × 8 GPU × 2 ノード
| Megatron-LM | DeepSpeed-ZERO2 | DeepSpeed-ZERO3 | |
|---|---|---|---|
| Training Speed | 9.6 s/it | – | 91.2 s/it |
| GPU Memory Usage | 16 × 60 GB | OOM | 16 × 80 GB |
Qwen2.5-14B — NVIDIA A800 × 8 GPU (1ノード)
| Megatron-LM | DeepSpeed-ZeRO2 | DeepSpeed-ZeRO3 | |
|---|---|---|---|
| Training Speed | 9.04 s/it | 10.32 s/it | 10.56 s/it |
| GPU Memory Usage | 8 × 64 GB | 8 × 80 GB | 8 × 58 GB |
Qwen1.5-MoE-A2.7B — NVIDIA A800 × 8 GPU (1ノード)
| Megatron-LM | DeepSpeed-ZeRO2 | DeepSpeed-ZeRO3 | |
|---|---|---|---|
| Training Speed | 2.93 s/it | 6.02 s/it | 24.30 s/it |
| GPU Memory Usage | 8 × 66 GB | 8 × 72 GB | 8 × 50 GB |
Megatronベースの学習に対応しているモデルの一覧については、下記ドキュメントの「
Support Megatron」の項目をご確認ください。
以下、このMegatron-SWIFTを使ったQwen3のファインチューニングについて解説します。
環境構築
pipを使った環境構築
上述したベストプラクティス記事のドキュメント通りに進めていきます。
まずはms-swiftをインストールします。
git clone https://github.com/modelscope/ms-swift.git
cd ms-swift
pip install -e .
pip install liger-kernel transformers -U
この時点でswiftコマンドによるtransformersをバックエンドとした基本的な学習を実行可能なようです。下記は記事中にある学習コマンドの一例です。
# Training GPU memory: 22GB
# You can specify `--dataset AI-ModelScope/alpaca-gpt4-data-zh` to run the experiment
CUDA_VISIBLE_DEVICES=0 \
swift sft \
--model Qwen/Qwen3-8B \
--train_type lora \
--dataset '<dataset-path>' \
--torch_dtype bfloat16 \
--num_train_epochs 1 \
--per_device_train_batch_size 1 \
--per_device_eval_batch_size 1 \
--learning_rate 1e-4 \
--lora_rank 8 \
--lora_alpha 32 \
--target_modules all-linear \
--gradient_accumulation_steps 4 \
--eval_steps 50 \
--save_steps 50 \
--save_total_limit 2 \
--logging_steps 5 \
--max_length 2048 \
--output_dir output \
--warmup_ratio 0.05 \
--dataloader_num_workers 4 \
--packing true \
--use_liger_kernel true \
--attn_impl flash_attn
今回はMegatron-LMベースの学習にフォーカスを当てた記事なので、そのために下記ドキュメントに従い追加の環境準備をしていきます。
上記ドキュメントに従い、TransformerEngineやAPEX、Megatron-Coreなどの必要なライブラリをインストールします。
# Recommended PyTorch version: 2.5 / 2.6
pip install pybind11
# transformer_engine
# If an installation error occurs, you can refer to this issue for resolution: https://github.com/modelscope/ms-swift/issues/3793
pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable
# apex
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./
# megatron-core
pip install git+https://github.com/NVIDIA/Megatron-LM.git@core_r0.12.0
この中で、特にTransformerEngineのビルドで躓くことが多いと思います。実際私もいくつかビルドエラーに遭遇し、様々な対処が必要になりました。
私の環境で発生した問題に対して行っていた対応を記載しておきます。
-
このissueにあるようにインストール時のパス設定を行う
SITE_PACKAGES=$(python -c "import site; print(site.getsitepackages()[0])") && echo $SITE_PACKAGES && \ CUDNN_PATH=$SITE_PACKAGES/nvidia/cudnn CPLUS_INCLUDE_PATH=$SITE_PACKAGES/nvidia/cudnn/include \ pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable - ビルドの際に
MAX_JOBS=Nを大きくしすぎると環境によってはビルド中にクラッシュするのであまり大きな値に設定しない - ビルド時と実行時のPyTorchのバージョンは統一する
これらのライブラリのインストールが問題なく完遂すれば環境構築は無事完了です。
なお、Megatron-LM本体は後述するモデル変換や学習実行の際に自動でキャッシュディレクトリにgit cloneされますが、MEGATRON_LM_PATHという環境変数で他のディレクトリに存在するものを利用することもできます。
Dockerを使った環境構築
公式から配布されているDockerイメージを元に環境構築することもできます。TransformerEngine等がすでにビルド済みで入っているので、かなり楽に使えます。
ドキュメントに記載のある通り以下のイメージが公式より配布されているので、それを利用します。
modelscope-registry.cn-hangzhou.cr.aliyuncs.com/modelscope-repo/modelscope:ubuntu22.04-cuda12.4.0-py311-torch2.6.0-vllm0.8.3-modelscope1.25.0-swift3.3.0.post1
modelscope-registry.us-west-1.cr.aliyuncs.com/modelscope-repo/modelscope:ubuntu22.04-cuda12.4.0-py311-torch2.6.0-vllm0.8.3-modelscope1.25.0-swift3.3.0.post1
上記のイメージを通常通りdocker runすればそのまま利用可能です。
注意点として、イメージに入っているms-swiftのライブラリが古くQwen3に対応していないので、アップデートが必要です。コンテナ内のシェルでアップデートするか、新しくイメージをビルドするなどしてください。
docker run --gpus all -it \
modelscope-registry.cn-hangzhou.cr.aliyuncs.com/modelscope-repo/modelscope:ubuntu22.04-cuda12.4.0-py311-torch2.6.0-vllm0.8.3-modelscope1.25.0-swift3.3.0.post1 \
/bin/bash
# ms-swiftを更新
pip install --no-cache-dir -U \
"git+https://github.com/modelscope/ms-swift.git@main"
Qwen3の学習
データセット準備
学習に利用するデータセットを準備します。以下の公式ドキュメントを元に基本的な部分のみを簡単に解説します。
ms-swiftが標準で受け付けるデータセットの形式は以下の4つです。
- Messages形式
{"messages": [{"role": "system", "content": "<system>"}, {"role": "user", "content": "<query1>"}, {"role": "assistant", "content": "<response1>"}, {"role": "user", "content": "<query2>"}, {"role": "assistant", "content": "<response2>"}]} - ShareGPT形式
{"system": "<system>", "conversation": [{"human": "<query1>", "assistant": "<response1>"}, {"human": "<query2>", "assistant": "<response2>"}]} - Alpaca形式
{"system": "<system>", "instruction": "<query-inst>", "input": "<query-input>", "output": "<response>"} - Query-Response形式
{"system": "<system>", "query": "<query2>", "response": "<response2>", "history": [["<query1>", "<response1>"]]}
最も基本となる形式は1つ目のMessages形式のようで、基本的にはこれに合わせておけば問題ありません。
例えばSFTに利用するデータ形式の一例は以下の通りです。
{"messages": [{"role": "system", "content": "You are a useful and harmless assistant"}, {"role": "user", "content": "Tell me tomorrow's weather"}, {"role": "assistant", "content": "Tomorrow's weather will be sunny"}]}
{"messages": [{"role": "system", "content": "You are a useful and harmless math calculator"}, {"role": "user", "content": "What is 1 + 1?"}, {"role": "assistant", "content": "It equals 2"}, {"role": "user", "content": "What about adding 1?"}, {"role": "assistant", "content": "It equals 3"}]}
継続事前学習を行いたい場合は以下のようにassistantのroleのみを持ったデータセットを渡すことでそのまま利用可能となります。
{"messages": [{"role": "assistant", "content": "I love music"}]}
{"messages": [{"role": "assistant", "content": "Coach, I want to play basketball"}]}
{"messages": [{"role": "assistant", "content": "Which is more authoritative, tomato and egg rice or the third fresh stir-fry?"}]}
その他様々なオプションがあるので、気になる方はドキュメントをご確認ください。
モデル変換（HF→mcore）
Megatron-LMによる学習を行う前に、HF形式のモデルをmcore形式に変換する必要があります。
Megatron-SWIFTではswift exportコマンドで簡単に変換を行うことが出来ます。
下記コマンドはQwen/Qwen3-30B-A3B-Baseをmcore形式に変換する例です。
CUDA_VISIBLE_DEVICES=0 \
swift export \
--model Qwen/Qwen3-30B-A3B-Base \
--to_mcore true \
--torch_dtype bfloat16 \
--output_dir /path/to/Qwen3-30B-A3B-Base-mcore
また、--test_convert_precision trueを引数に追加すると、変換前後のモデルでの精度比較を行い変換に問題がないかをチェックしてくれます。
なお、ms-swiftでは基本的にデフォルトでHugging FaceではなくModelScopeからモデルやデータセットをダウンロードします。HFからダウンロードしたい場合は--use_hf trueの引数を追加してください。
学習の実行
megatron ptコマンドで継続事前学習を、megatron sftコマンドでSFTを実行できます。なお、megatron ptは--use_chat_template=falseのmegatron sftコマンドと同等です。
以下は私がH200x8の環境でQwen3-30B-A3Bの学習に利用したコマンドの例です。
-
継続事前学習の例
CUDA_DEVICE_MAX_CONNECTIONS=1 \ CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ NPROC_PER_NODE=8 \ megatron pt \ --load /path/to/Qwen3-30B-A3B-Base-mcore \ --dataset '/path/to/dataset' \ --tensor_model_parallel_size 2 \ --expert_model_parallel_size 4 \ --context_parallel_size 1 \ --sequence_parallel true \ --moe_grouped_gemm true \ --moe_shared_expert_overlap true \ --moe_aux_loss_coeff 0.01 \ --micro_batch_size 1 \ --global_batch_size 256 \ --recompute_granularity full \ --recompute_method uniform \ --recompute_num_layers 1 \ --train_iters 800 \ --finetune true \ --cross_entropy_loss_fusion true \ --lr 1e-5 \ --lr_warmup_iters 80 \ --min_lr 1e-6 \ --save /path/to/megatron_output/Qwen3-30B-A3B-CPT \ --save_interval 100 \ --max_length 32768 \ --truncation_strategy right \ --num_workers 1 \ --no_save_optim true \ --no_save_rng true \ --dataset_num_proc 64 \ --packing true \ --streaming true \ --use_flash_attn true \ --wandb_project qwen3 \ --wandb_exp_name 30b-a3b-cpt \ --log_interval 1 -
SFTの例
CUDA_DEVICE_MAX_CONNECTIONS=1 \ CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ NPROC_PER_NODE=8 \ megatron sft \ --load /path/to/megatron_output/Qwen3-30B-A3B-CPT/vx-xxx \ --dataset '/path/to/dataset' \ --tensor_model_parallel_size 2 \ --expert_model_parallel_size 4 \ --context_parallel_size 1 \ --sequence_parallel true \ --moe_grouped_gemm true \ --moe_shared_expert_overlap true \ --moe_aux_loss_coeff 0.01 \ --micro_batch_size 1 \ --global_batch_size 64 \ --recompute_granularity full \ --recompute_method uniform \ --recompute_num_layers 1 \ --train_iters 300 \ --finetune true \ --cross_entropy_loss_fusion true \ --lr 1e-5 \ --lr_warmup_iters 30 \ --min_lr 1e-6 \ --save /path/to/megatron_output/Qwen3-30B-A3B-SFT \ --save_interval 100 \ --max_length 32768 \ --num_workers 1 \ --no_save_optim true \ --no_save_rng true \ --dataset_num_proc 64 \ --packing true \ --streaming false \ --dataset_shuffle true \ --use_flash_attn true \ --wandb_project qwen3 \ --wandb_exp_name 30b-a3b-sft \ --log_interval 1 \ --use_chat_template true
上記の設定+flash attention 3で、私の環境では450~500 TFLOPsほどの学習速度が出ていました。
その他のコマンドライン引数の詳細については下記ドキュメントをご確認ください。
モデル変換（mcore→HF）
最後に学習後のモデルをmcore形式からHF形式に変換します。
HF→mcoreの場合と同様、swift exportコマンドで簡単に変換できます。
CUDA_VISIBLE_DEVICES=0 \
swift export \
--mcore_model /path/to/megatron_output/Qwen3-30B-A3B-SFT/vx-xxx \
--to_hf true \
--torch_dtype bfloat16 \
--output_dir /path/to/hf_model/Qwen3-30B-A3B-SFT-hf
HF→mcoreと同様、--test_convert_precision trueを引数に追加すると変換前後のモデルでの精度比較を行い変換に問題がないかをチェックしてくれます。
あとはこれをhuggingface-cli等でHFにアップロードすれば完了です。
まとめ
この記事では、ms-swiftというフレームワークを使ってMegatron-LMをベースにLLMをファインチューニングする方法について解説しました。
元々Megatron-LMによる学習がかなり効率が良いというのは知っていたのですが、環境構築等のハードルが高く敬遠している部分がありました。このMegatron-SWIFTではMegatronが上手く統合されており、かなり簡単に利用することが出来ました。
もしこの記事を読んで興味が湧いたという方がいればぜひ手元で試してみてください。
余談
今回H200x8の環境を用意するにあたり安さだけに釣られてHPC-AIという謎のクラウドGPUサービスを使用しました。
なんとH200x8が$16.72/hourという異常な安さで利用できます（その代わり8台単位でしか借りられない）。流石に安すぎて怪しんでいたのですが、今のところは問題なく使えています・
下記リンクから飛ぶと私の招待コード経由で登録したことになりますが、残念ながら招待コードを使っていただいても現状特に報酬等はないようです…（公式に問い合わせて確認しました）。
一応リンクだけ貼っておくので、ご興味のある方はお試しください。
Discussion