こんにちは、代表の藤倉です。
本記事では医療分野に特化したLLM（大規模言語モデル）GatorTronの導入について解説していきます。
日本語資料は当然として英語資料も少ないのはキツかった。
執筆時の現在だとGatorTronのモデルが置いてあるだけで公式のチュートリアルは一切ありません。
この記事があることで多少は救われる人がいると思うので、落とし穴の回避策も含めて本記事にまとめていきたいと思います。

事前準備

方針

GatorTronは同じくNVIDIAが提供しているフレームワーク（NeMo）用の.nemo形式とPyTrochモデルのチェックアウト(.pt)形式の２つで公開されています。

まずはNeMo版についてですが、肝心のNeMo上で読み込みができずGithub上に報告があがってはいますが、１ヶ月以上も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になっているため、ここで混乱しないように意識しておくことが大事です。

GatorTronのダウンロード

以下のNVIDIAのページからモデルをダウンロードします。

https://catalog.ngc.nvidia.com/orgs/nvidia/teams/clara/models/gatortron_og

右上のDownloadのボタンをクリックしWGETを選択します。
するとクリップボードにwgetを使ったスクリプトがコピーされますので、導入したいマシンのコンソール上に貼り付けてダウンロードを実行します。

データ容量が5GBとなかなか重たいです、早速ですがコーヒーでも淹れて待ちましょう。
待ってる間に50杯は飲めると思います。

ダウンロード前のコンテンツリストかダウンロード後に中身を見ればわかりますが、GatorTronというモデル名ではなくMegatronBERTというモデル名でチェックポイントが入っています。

その他の準備も合わせて並行して進めてしまいましょう。

Megatronを使えるようにする

次に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を構築する

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