メルマガ登録
このたびブレインパッドは、LLM/Generative AIに関する研究プロジェクトを立ち上げ、この「Platinum Data Blog」を通じてLLM/Generative AIに関するさまざまな情報を発信をしています。この記事では、2023年3月に発表されたPyTorch 2.0のcompileを試し、速度改善がどの程度になるのか調査した結果をご紹介します。
こんにちは。アナリティクスサービス部の加茂です。
昨今のLLMの潮流に乗るように、機械学習の速度を向上させる取り組みも活発に行われており、JAXを利用したWhisperの高速化やmojo言語の発表など、日々賑わいを見せています。
今回は少しリリースから時間が経っていますが、2023年3月に発表されたPyTorch 2.0のcompileを試し、速度改善がどの程度になるのかを簡単に調査してみたいと思います。
推論よりも学習において速度の向上の効果が期待できると公式のページに記載があったため、タスクとしては言語モデルのfine-tuningを選択しました。
方法は主にこちらのurl を参考にさせていただきました。
transformersのgithubリポジトリのexampleにある、examples/pytorch/language-modeling/run_clm.py
を利用して、
同様の学習データで少し軽量なモデル(rinna社のjapanese-gpt2-xsmall)のfine-tuningを行いました。
run_clm.py
は実行時に引数を渡すだけでcompileの有無などを指定できます。
下表は関連する3つのオプションとその簡単な説明です。これらのcompileに関連するオプションは試験的に導入しており、将来的に変更される可能性があると書かれているため、もし試される方がいらっしゃったらバージョンにご注意ください。
オプション | 説明 |
---|---|
torch_compile (bool, *optional*, defaults to False) | PyTorch 2.0によるモデルのコンパイルを行うか否かを指定する |
torch_compile_backend (str, *optional*) | torch.compile で使用するbackendを指定する |
torch_compile_mode (str, *optional*) | torch.compile で使用するmodeを指定する |
実際にオプションが使われているコードの箇所を確認すると、transfomers.Trainer
の_wrap_model
メソッド内で .compile
が呼び出されていることが確認できます。
# torch.compile() needs to be called after wrapping the model with FSDP or DDP
# to ensure that it accounts for the graph breaks required by those wrappers
if self.args.torch_compile:
model = torch.compile(model, backend=self.args.torch_compile_backend, mode=self.args.torch_compile_mode)
今回はtorch_compile_mode
の各パターンも試してみたいと思います。
run_clm.py
を実行します。
--torch_compile=True
を付与して実行--torch_compile=True --torch_compile_mode=reduce-overhead
を付与して実行--torch_compile=Ture --torch_compile_mode=max-autotune
を付与して実行python transformers/examples/pytorch/language-modeling/run_clm.py
--model_name_or_path=rinna/japanese-gpt2-xsmall
--train_file=train_data/dolly-oasst1-ja.txt
--output_dir=output
--do_train
--optim=adafactor
--num_train_epochs=5
--save_steps=721
--logging_steps=72
--learning_rate=1e-03
--per_device_train_batch_size=4
--save_safetensors=True
--logging_dir=logs
今回の環境では問題なく全てのオプションが動作し、各実行条件とタスクの学習時間は次の通りになりました。
実行条件 | 計算時間 (sec.) |
---|---|
compileなし (PyTorch 1.13) | 1016.0 |
compileなし (PyTorch 2.0) | 1059.0 |
compileあり (default) | 497.9 |
compileあり (reduce-overhead) | 497.6 |
compileあり (max-autotune) | 769.1 |
compileなしと比較するとおよそ計算時間が半分程度になっていることが分かります。
公式ページではNVIDIA A100 GPUで43%速くなったと報告されており、今回はそれを少し上回る改善が確認できました。
速度改善が10%未満に留まっている検証記事も散見される(*1, *2)ため、利用するGPUによってかなり改善幅が変わるようです。
2点ほど細かい点に触れておくと、まずcompileなしの場合、PyTorch 2.0 と PyTorch 1.13では若干2.0の方が遅くなっています。2.0を利用していたとしても、compileなしでは速度改善しないことが確認できます。
また、modeがdefault
と reduce-overhead
ではほぼ同程度となっておりますが、max-autotune
は他の2つと比べると計算時間が長くなっています。公式によると、max-autotune
は最も高速なモデルにcompileを行う代わりに、compileに非常に時間がかかると説明されています。今回の環境とタスクでは、compileに要する時間の影響の方が大きかったと考えられます。
本記事の結果を簡単にまとめますと、PyTorch 2.0を利用して高速化したい場合、一旦はcompileありでdefaultもしくはreduce-overheadのmodeを利用すると良いと結論付けられそうです。
あなたにオススメの記事
2023.12.01
生成AI(ジェネレーティブAI)とは?ChatGPTとの違いや仕組み・種類・活用事例
2023.09.21
DX(デジタルトランスフォーメーション)とは?今さら聞けない意味・定義を分かりやすく解説【2024年最新】
2023.11.24
【現役社員が解説】データサイエンティストとは?仕事内容やAI・DX時代に必要なスキル
2023.09.08
DX事例26選:6つの業界別に紹介~有名企業はどんなDXをやっている?~【2024年最新版】
2023.08.23
LLM(大規模言語モデル)とは?生成AIとの違いや活用事例・課題
2024.03.22
生成AIの評価指標・ベンチマークとそれらに関連する問題点や限界を解説