PyTorchのインストールから始めてLINE社が公開した36億パラメータのjapanese-large-lmを動かしてみた

環境

• Windows11
• Intel Corei9-11900K
• メモリ32GB（実行には15GB程度の余剰が必要）
• NVIDIA GeForce RTX 3070
• Anaconda3がインストール済み

CUDAのインストール

• こちらを参考に、CUDA11.8をインストールした

PyTorchのインストール

• anacondaで「pytorch」という環境を作成した（pythonのバージョンは3.11.4）。
• こちらを参考に、以下の手順でPyTorchをインストールした。

$ conda activate pytorch
$ pip install --upgrade pip setuptools
$ pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

必要なPythonライブラリのインストール

• 以下の手順で必要なPythonライブラリをインストールした。下の２つはなくても動いたが、Warningが発生したのでインストールした。

$ pip install transformers
$ pip install sentencepiece
$ pip install xformers

japanese-large-lmによるGeneration

• 公式ブログを参考に、下記のPythonスクリプトを実行した。結果として、5つの文章が生成された。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, set_seed

print("Preparing model...")
model = AutoModelForCausalLM.from_pretrained("line-corporation/japanese-large-lm-3.6b", torch_dtype=torch.float16)

print("Preparing tokenizer...")
tokenizer = AutoTokenizer.from_pretrained("line-corporation/japanese-large-lm-3.6b", use_fast=False)

print("Preparing generator...")
generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)
set_seed(101)

print("Generating text...") 
text = generator(
    "おはようございます、今日の天気は",
    max_length=30,
    do_sample=True,
    pad_token_id=tokenizer.pad_token_id,
    num_return_sequences=5,
)

for t in text:
    print(t)