2023年3月上旬にPyTorch 2.0がリリースされる予定です。
このPyTorch 2.0の目玉機能にtorch.compile()というものがあります。これを使ってTransformersの日本語文章生成を試してみました。
環境
- Google Colaboratory
- GPU有効（Nvidia Tesla T4）
- PyTorch 2.0.0.dev20230213+cu117
環境構築
Google Colaboratoryを利用します。
2023年2月16日現在、PyTorch 2.0 Nightlyビルドは日々開発が行われているせいかインストールするたびに挙動が不安定になる事があります。なので、バージョンを指定してインストールします。
!pip install numpy --pre https://download.pytorch.org/whl/nightly/cu117/torch-2.0.0.dev20230212%2Bcu117-cp38-cp38-linux_x86_64.whl --force-reinstall
またTransformersもインストールしておきます
!pip install transformers sentencepiece
文章生成
rinnaの公開しているGPTモデルを利用します。好みのモデル名を設定します。
MODEL_NAME = 'rinna/japanese-gpt2-xsmall'
# 好みで選択
# MODEL_NAME = 'rinna/japanese-gpt2-medium'
# MODEL_NAME = 'rinna/japanese-gpt-1b'
指定したモデルを読み込みます。
import torch
from transformers import T5Tokenizer, AutoModelForCausalLM
tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
if torch.cuda.is_available():
model = model.to("cuda")
次がPyTorch 2.0で異なるところです。torch.compile()の1行を追加します。
model_optimized = torch.compile(model)
文章を生成する実装をします。
input_ids = tokenizer.encode(
"昔々あるところに",
add_special_tokens=False,
return_tensors="pt"
)
def generate_by(model):
return model.generate(
input_ids.to(model.device),
pad_token_id=tokenizer.pad_token_id,
do_sample=True,
)[0]
print(tokenizer.decode(generate_by(model_optimized)))
出力※毎回異なります。
昔々あるところに、お世継ぎを残すために産まれたばかりの王子様と魔女と4人の魔女がいました。
（以下略）
以上です。torch.compile()以外は今までと同じであり、この1行を通っていないモデルでも普通に動きます。PyTorch 2.0と1.x系は互換性が高いようです。
PyTorch 2.0での変化
公式ブログによると、PyTorch 2.0でtorch.compile()を噛ますことでTransformersの多くのモデルにおいて処理が高速化したようです。しかし今回利用したGPUとモデルでは、色々試してみたものの速度の変化が全く見られませんでした。
# 生成速度は変わらず
generate_by(model)
generate_by(model_optimized)
torch.compile()では様々なモードを設定することができるようです。しかし、これらによる変化も今回のコードでは分かりませんでした。
model_optimized = torch.compile(model)
# フレームワークのオーバーヘッドを減らすため最適化するが、メモリが多く消費される。小さなモデルで効果的
model_optimized = torch.compile(model, mode="reduce-overhead")
# 最速のモデルを生成するようになるが、コンパイルに時間がかかる
model_optimized = torch.compile(model, mode="max-autotune")