こんにちは、代表の藤倉です。
本記事では医療分野に特化したLLM(大規模言語モデル)GatorTronの導入について解説していきます。日本語資料は当然として英語資料も少ないのはキツかった。
執筆時の現在だとGatorTronのモデルが置いてあるだけで公式のチュートリアルは一切ありません。
この記事があることで多少は救われる人がいると思うので、落とし穴の回避策も含めて本記事にまとめていきたいと思います。
GatorTronは同じくNVIDIAが提供しているフレームワーク(NeMo)用の.nemo形式とPyTrochモデルのチェックアウト(.pt)形式の2つで公開されています。
まずはNeMo版についてですが、肝心のNeMo上で読み込みができずGithub上に報告があがってはいますが、1ヶ月以上もcloseしない状況です(https://github.com/NVIDIA/NeMo/issues/5654)。
認知度の高いOSSだと上記のような状況自体はなんら珍しいものではないですが、それに対してユーザ側でなんとかしようというトライアンドエラーに関する記事がほとんどなく、コミュニティがまだ小さい印象を受けました。というかNeMoに関するユーザの記事もめちゃくちゃ少ないのでトライアンドエラーの難易度が高い。
従って、本記事ではGatorTronのNeMo版の使用は避け、PyTorchで扱えるようにしています。
※ちなみにNeMo Megatronのエンタープライズ版なるものがあるようで、こちらであれば公式のNVIDIAからサポートを受けながら開発できるものだと思われます(https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/nlp/nemo_megatron/intro.html)。
また、GatorTronは厳密にはMegatronというLLMが前身にあり、それを医療分野にファインチューニングしたものとなります(https://huggingface.co/nvidia/megatron-bert-cased-345m)。
後述しますが、実際のGatorTronのファイル名もMegatronになっているため、ここで混乱しないように意識しておくことが大事です。
以下のNVIDIAのページからモデルをダウンロードします。
https://catalog.ngc.nvidia.com/orgs/nvidia/teams/clara/models/gatortron_og
右上のDownloadのボタンをクリックしWGETを選択します。
するとクリップボードにwgetを使ったスクリプトがコピーされますので、導入したいマシンのコンソール上に貼り付けてダウンロードを実行します。
データ容量が5GBとなかなか重たいです、早速ですがコーヒーでも淹れて待ちましょう。
待ってる間に50杯は飲めると思います。
ダウンロード前のコンテンツリストかダウンロード後に中身を見ればわかりますが、GatorTronというモデル名ではなくMegatronBERTというモデル名でチェックポイントが入っています。
その他の準備も合わせて並行して進めてしまいましょう。
次にMegatronを使用できるようにセッティングしていきます。
基本的にはこちらの案内(https://huggingface.co/nvidia/megatron-bert-cased-345m)に従ってMegatronを使用できるようにします。
まずは環境変数を設定していきます。
export MYDIR=$HOME
公式の例では$HOMEを指定していますが、$HOMEは基本的に汚したくないためtmpディレクトリなどを作成してそこに設定すると良いと思います。
次にmegatronをダウンロードしていきます。
mkdir -p $MYDIR/nvidia/megatron-bert-cased-345m wget --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/megatron_bert_345m/versions/v0.1_cased/zip -O $MYDIR/nvidia/megatron-bert-cased-345m/checkpoint.zip python3 $MYDIR/transformers/src/transformers/models/megatron_bert/convert_megatron_bert_checkpoint.py $MYDIR/nvidia/megatron-bert-cased-345m/checkpoint.zip
上記の2行目でダウンロードを行い、3行目で圧縮ファイルからファイルの作成を行っています。
ここまでのステップが完了すると、config.json
とpytorch_model.bin
が生成されています。
GatorTronにはPyTorchのチェックアウト(MegatronBERT.pt)のみがあり、Megatronにはconfig.jsonとmegatronのバイナリ(pytorch_model.bin)があります。
つまり、モデルアーキテクチャとしてはMegatronを読み込み、重みパラメータはMegatronBERT.ptを指定します。
これを表現しようとすると以下のコードになります。
import os import torch gatortron_path = '/workspace/model/gatortron/MegatronBERT.pt' state_dict = torch.load(gatortron_path)["model"] megatron_dir = '/workspace/gatortron/nvidia/megatron-bert-cased-345m' # docker上のpathにする model = MegatronBertModel.from_pretrained(megatron_dir, state_dict=state_dict) model
上記の図のようにモデルアーキテクチャが表示されたら読み込み成功です。
筆者の場合はDocker上にJupyterNotebookを展開しているため、gatortron_pathはDocker内でのpathになります。
Dockerを使わない方はダウンロードしてきたmegatron-bert-cased-345mのpathをそのまま指定すると良いでしょう。
Dockerを使う方場合はCOPYで移しておくかdocker run時にボリュームのマウンティングを行う必要があります。
導入については以上となります、お疲れさまでした
使い方についてはまた次の記事で紹介したいと思います。
最後にいつもので締めさせていただきます。
サイシキでは医療分野のような専門性の高い分野における強いAIの開発にも対応してございます。
自然言語処理について開発やサポートが必要になりましたら弊社までご相談いただければ幸いです。
資料請求も以下からお問い合わせいただけます。
お問い合わせ: https://www.sai-shiki.com/contact/
メールアドレス: office@sai-shiki.com