BioNeMo Framework とは
BioNeMo Framework は、創薬のための生成 AI モデルを開発するため、NVIDIA が提供しているフレームワークで、研究開発者が生成 AI、大規模言語モデル (LLM)、および基盤モデルの構築とデプロイを容易に行うことができます。本記事では BioNemo Framework を利用して、大規模タンパク質モデルの事前学習、ファインチューニングと推論方法について説明します。
この記事を執筆している時点では、BioNeMo Framework の最新バージョン v1.4 がリリースされ、コンテナーが NGC で一般公開されており、ダウンロードが可能です。コンテナーには以下のゲノム、タンパク質、低分子、シングル セルに関する 11 種類のモデルが含まれていますが、本記事はタンパク質モデルの ESM-2nv を例として説明します。
- DNABERT
- Enformer
- MegaMolBART
- MolMIM
- EquiDock
- DiffDock
- ESM-1nv
- ESM-2nv
- ProtT5
- OpenFold
- Geneformer
必要な計算環境
- x86 Linux システム
- Docker (GPU 対応、docker engine 19.03 以上)
- モデル トレーニングを行う場合は NVIDIA GPU が必要
- GPU: Pascal 世代以降の GPU (例えば、Pascal、Volta、Turing、Ampere、Hopper)
- bfloat16 精度は Ampere 世代の GPU 以上が必要
また、今回は以下の環境でチュートリアルを行っています。
- GPU: DGX A100 (8x NVIDIA A100 80 GB GPUs)
- Driver: 535.129.03
- OS: Ubuntu 22.04.3 LTS
- Docker: 23.0.4
- Container: nvcr.io/nvidia/clara/bionemo-framework:1.4
BioNemo Framework コンテナーを入手
docker login nvcr.io
Username: $oauthtoken
Password: <API キーを入力してください>
- BioNemo Framework の最新コンテナー v1.4 を Pull
docker pull nvcr.io/nvidia/clara/bionemo-framework:1.4[Option] Weights and Biases のセットアップ
モデルのトレーニング進行状況やチャートは、 Weights and Biases を通じて可視化できます。ログを有効にするために、 API キーを設定してください。
BioNemo Framework はローカル マシンやクラウドでも実行可能ですが、今回の記事はローカル マシンでの使い方について説明します。
Docker コンテナーの起動
docker run -it --rm --gpus all -v /home/bionemo/:/workspace/bionemo/mydata/ nvcr.io/nvidia/clara/bionemo-framework:1.4 bash-v オプションはローカルのディレクトリをコンテナー内にマウントします。もし独自のデータを利用して継続事前学習、ファインチューニング、および推論したい場合は、ローカルのディレクトリ /home/bionemo/ にデータを置いておけば、コンテナー内の /workspace/bionemo/mydata/ でそのデータが確認できます。
Jupyter Lab の起動
事前学習とファインチューニングのステップ、Jupyter Lab は必須ではありませんが、推論ステップは必要になります。Jupyter Lab を利用したい場合は、コンテナーの中で、下記のコマンドを入力して Jupyter Lab を起動してください。
jupyter lab --allow-root --ip=* --port=8888 --no-browser --NotebookApp.token='' --NotebookApp.allow_origin='*' --ContentsManager.allow_hidden=True --notebook-dir=/workspace/bionemo & sleep infinityそして、ブラウザーを開いて、http://hostname:8888 に接続して notebook のページが表示されれば大丈夫です。
Jupyter Lab を起動した場合は、Jupyter Lab のターミナルを開いて、これからのステップを実行してください。そうではない場合は、ローカル マシンのターミナルを開いて実行してください。
データの事前準備
ESM-2nv のトレーニング中には 2 つのデータセットを使用します。Uniref50 はトレーニング、検証、テスト用に分割されています。データのサイズと多様性を増加させるために、トレーニング中に Uniref50 のミニバッチがサンプリングされますが、このバッチの各シーケンスは、対応する Uniref90 クラスターのシーケンスに置き換えられます。詳細は「Language models of protein sequences at the scale of evolution enable accurate structure prediction (進化の規模でのタンパク質配列の言語モデルが正確な構造予測を実現)」を参照してください。Uniref90 はトレーニング時にのみ使用され、検証やテストには使用されません。
データの前処理を行う前に、以下のコマンドでサンプル データを解凍してください。
cd /workspace/bionemo/examples/tests/test_data/
unzip uniref202104_esm2_qc_test200_val200.zip
データの前処理
次に、下記の場所にあるコードを実行してデータの前処理を行いますが、その前に、データのパスを設定する必要があります。/workspace/bionemo/examples/protein/esm2nv/pretrain.py
データのパスを正しく設定するために、下記の YAML 設定ファイルを変更するか、データの前処理やモデル トレーニングを行う際にコマンドの一部としてそれらのパラメーターを更新することができます。/workspace/bionemo/examples/protein/esm2nv/conf/base_config.yaml
では、以下のコマンドでデータの前処理を行います。
cd /workspace/bionemo
python examples/protein/esm2nv/pretrain.py\
--config-path=conf\
--config-name=pretrain_esm2_650M\
++do_training=False\
++model.data.val_size=500\
++model.data.test_size=100\
++model.data.uf50_datapath=/workspace/bionemo/examples/tests/test_data/uniref202104_esm2_qc_test200_val200/uniref50_train_filt.fasta\
++model.data.uf90_datapath=/workspace/bionemo/examples/tests/test_data/uniref202104_esm2_qc_test200_val200/ur90_ur50_sampler.fasta\
++model.data.cluster_mapping_tsv=/workspace/bionemo/examples/tests/test_data/uniref202104_esm2_qc_test200_val200/mapping.tsv\
++model.data.dataset_path=/workspace/bionemo/examples/tests/test_data/uniref202104_esm2_qc_test200_val200/uf50\
++model.data.uf90.uniref90_path=/workspace/bionemo/examples/tests/test_data/uniref202104_esm2_qc_test200_val200/uf90\
++model.data.train.uf50_datapath=/workspace/bionemo/examples/tests/test_data/uniref202104_esm2_qc_test200_val200/uniref50_train_filt.fasta\
++model.data.train.uf90_datapath=/workspace/bionemo/examples/tests/test_data/uniref202104_esm2_qc_test200_val200/ur90_ur50_sampler.fasta\
++model.data.train.cluster_mapping_tsv=/workspace/bionemo/examples/tests/test_data/uniref202104_esm2_qc_test200_val200/mapping.tsv\
++model.data.val.uf50_datapath=/workspace/bionemo/examples/tests/test_data/uniref202104_esm2_qc_test200_val200/uniref50_train_filt.fasta\
++model.data.test.uf50_datapath=/workspace/bionemo/examples/tests/test_data/uniref202104_esm2_qc_test200_val200/uniref50_train_filt.fasta
-- で始まるパラメーターは、コマンドライン引数として pretrain.py に渡されます。例えば、config-path と config-name は、設定 YAML ファイルのフォルダーと yaml ファイル名を指定します。このパスは pretrain.py に対して相対的です。conf は examples/protein/esm2nv/conf を指し、pretrain_esm2_650M は examples/protein/esm2nv/conf/pretrain_esm2_650M.yaml を指します。
++ で始まるパラメーターは、YAML ファイルで設定可能です。 例えば、base_config.yaml から継承された pretrain_esm2_650M.yaml では、以下のパラメーターを見つけることができます:
do_training: データの前処理のみを行い、トレーニングは行わないように False に設定します。model.data.val_size と model.data.test_size: 検証およびテスト データセットのサイズを指定します。model.data.uf50_datapath: uniref50 fasta ファイルへのパスを指定します。model.data.uf90_datapath: uniref90 fasta ファイルへのパスを指定します。model.data.cluster_mapping_tsv: uniref50 クラスターを uniref90 シーケンスにマッピングするファイルへのパスを指定します。model.data.dataset_path: 前処理された uniref50 データの出力ディレクトリへのパスを指定します。このフォルダーはトレイン、検証、テストの分割を含むことになります。model.data.uf90.uniref90_path: 前処理された uniref90 データの出力ディレクトリへのパスを指定します。このフォルダーは u90_csvs というフォルダーのみを持ち、uniref90 はトレーニングでのみ使用されるため、トレイン/テスト/検証の分割は持ちません。model.data.train.uf50_datapath: uniref50 fasta ファイルへのパスを指定します。model.data.train.uf90_datapath: uniref90 fasta ファイルへのパスを指定します。model.data.train.cluster_mapping_tsv: uniref50 クラスターを uniref90 シーケンスにマッピングするファイルへのパスを指定します。-
model.data.val.uf50_datapath: uniref50 fasta ファイルへのパスを指定します。 model.data.test.uf50_datapath: uniref50 fasta ファイルへのパスを指定します。
また、上記のようにコマンドラインを通じて引数を上書きするのではなく、直接 YAML ファイルを変更することもできます。処理が完了すると、前処理されたデータは /workspace/bionemo/examples/tests/test_data/uniref202104_esm2_qc_test200_val200/uf50/uf50/ にあります。
もし独自のデータを利用して事前学習、ファインチューニングや推論をしたい場合は、パスを /workspace/bionemo/mydata/ に指定してください。ただし、データの構造やフォーマットをサンプル データに揃える必要があります。
モデルの事前学習
今回の例で UniRef50 と UniRef90 のサブセットの前処理が完了したので、ESM-2nv モデルの事前学習を開始することができます。BioNeMo Framework では、事前学習されたモデルのチェックポイントが提供されています。例えば、ESM-1nv、ESM-2nv (650m、3b、15b)、ProtT5nv、MegaMolBART などです。これらのモデルの重みは NVIDIA の NGC からダウンロードできます。すでにモデルのチェックポイントを持っている場合は、事前学習を再開することも可能です。
では、以下のコマンドでモデルを事前学習します。
cd /workspace/bionemo
python examples/protein/esm2nv/pretrain.py \
--config-path=conf \
--config-name=pretrain_esm2_650M \
++do_training=True \
++do_testing=False \
++model.data.dataset_path=examples/tests/test_data/uniref202104_esm2_qc_test200_val200/uf50 \
++model.data.uf90.uniref90_path=examples/tests/test_data/uniref202104_esm2_qc_test200_val200/uf50/uf90 \
++trainer.devices=2 \
++model.tensor_model_parallel_size=2 \
++model.micro_batch_size=8 \
++trainer.max_steps=50 \
++trainer.val_check_interval=5 \
++exp_manager.create_wandb_logger=True \
++exp_manager.checkpoint_callback_params.save_top_k=5
パラメーターの説明:
do_training: モデルをトレーニングするために True に設定します。これはデータが前処理されたことを前提としています。do_testing: テストをスキップするために False に設定します。model.data.dataset_path: トレーニング/検証/テスト分割が含まれる前処理済みの uniref50 データ フォルダーのパスを指定します。model.data.uf90.uniref90_path: 前処理済みの uniref90 データのパスを指定します。このフォルダーには u90_csvs という別のフォルダーが含まれている必要があります。u90_csvs フォルダー内には、x000.csv から x049.csv のファイルが存在する必要があります。trainer.devices: 使用する GPU の数を指定します。model.tensor_model_parallel_size: テンソル モデル並列サイズを設定します。model.micro_batch_size: バッチ サイズを設定します。メモリ エラーが発生しない限り、これをできるだけ増やします。trainer.max_steps: トレーニングの最大ステップ数を指定します。デモのために 50 に設定しました。1 ステップ = 1 バッチの処理です。まず、total_batches = サンプル総数 / バッチ サイズを計算します。N エポックでトレーニングしたい場合は、max_stepsを N *total_batchesに設定します。trainer.val_check_interval: 検証セットを実行する間隔を指定します。exp_manager.create_wandb_logger: wandb へのログ記録を無効にするために False に設定します。True の場合は、wandb API キーを提供する必要があります。exp_manager.checkpoint_callback_params.save_top_k: 保存する最適のチェックポイントの数を指定します。
トレーニングされた結果が /workspace/bionemo/results/nemo_experiments/ に保存されます。
モデルのファインチューニング
BioNemo Framework のサンプル コードは 3 つの下流タスクを提供しています。1 つ目はタンパク質の 10 の細胞内局在部位を予測すること、2 つ目はタンパク質の融解温度を予測すること、3 つ目はタンパク質の 3 状態構造を予測することです。本記事では 3 つ目のタスクを例として説明します。具体的には、シーケンスの各アミノ酸について、それがヘリックス、シート、またはコイルのどれにあるかを予測します。
サンプル データは下記のフォルダーにあります。/workspace/bionemo/examples/tests/test_data/protein/downstream/
次に、事前学習済みモデルをダウンロードします。モデルをダウンロードするため、ngc をインストールして、ngc config を設定する必要があります。
wget -q -O /tmp/ngccli_linux.zip --content-disposition https://api.ngc.nvidia.com/v2/resources/nvidia/ngc-apps/ngc_cli/versions/3.38.0/files/ngccli_linux.zip && unzip -o /tmp/ngccli_linux.zip -d /tmp && chmod u+x /tmp/ngc-cli/ngc && rm /tmp/ngccli_linux.zipそして、下記のように ngc config set をしてください。
/tmp/ngc-cli/ngc config set
<順番に API キー、CLI output format, org, team, ace を入力してください。API キーは NGC の API キーを入力してください。Org は ’no-org’以外のものを選択して、その他は「Enter」でデフォルト値を入力すれば大丈夫です>
最後に、下記のコマンドを入力すれば、すべてのモデルをダウンロードできます。
python download_models.py all --verbose --source ngc --download_dir ${BIONEMO_HOME}/models 特定のモデルのみをダウンロードしたい場合、例えば、今回のチュートリアルで使用する esm2nv_650m モデルの場合は、以下のコマンドでダウンロードできます。
python download_models.py --download_dir ${BIONEMO_HOME}/models esm2nv_650m ダウンロードされたモデルは .nemo というファイルで、 /workspace/bionemo/models に保存されます。
そして、以下の YAML ファイルを開きます。/workspace/bionemo/examples/protein/esm2nv/conf/downstream_flip_sec_str.yaml
以下のパラメーターは必要に応じて設定してください。
restore_from_path: 事前学習済みのモデルのチェックポイントの.nemoファイルのパスに設定します。trainer.devices, trainer.num_nodes: 使用する GPU とノードの数に設定します。trainer.max_epochs: トレーニングしたいエポック数に設定します。trainer.val_check_interval: 検証を実行するステップ数に設定します。model.micro_batch_size: トレーニングのマイクロバッチサイズに設定します。data.task_name: 任意の名前を設定します。data.task_type: 現在のオプションには、token-level-classification、classification (sequence level)、および regression (sequence level) があります。preprocessed_data_path:dataset_pathの親フォルダーのパスに設定します。dataset_path: train/val/test フォルダーが含まれるフォルダーのパスに設定します。例えば、上記のpath/to/dataというパスに設定します。dataset.train, dataset.val, dataset.test: CSV 名または範囲に設定します。sequence_column: シーケンスを含む列の名前に設定します。例: sequencetarget_column: ターゲットを含む列の名前に設定します。例: scl_labeltarget_size: 分類のための各ラベルのクラス数。num_classes:target_sizeに設定します。
そして、下記のコマンドでファインチューニングを実行します。BioNemo Framework v1.4 は新たな機能として、LoRa というファインチューニングの手法を追加しました。Lora は事前学習済み大規模言語モデルのすべての重みをファインチューニングするのではなく、大規模な重み行列を近似する 2 つの小規模な行列をファンチューニングする効率的なファインチューニングの手法です。
cd /workspace/bionemo
python examples/protein/downstream/downstream_flip.py\
--config-path="../esm2nv/conf"\
--config-name=downstream_sec_str_LORA\
++model.data.dataset_path=/workspace/bionemo/examples/tests/test_data/protein/downstream/
推論
下記の YAML ファイルを開いて、/workspace/bionemo/examples/protein/esm2nv/conf/infer.yaml
以下の情報を更新してください。提供する事前学習済みモデルを使用するか、自分でトレーニングしたモデルを使用することができます。
downstream_task:
restore_from_path: "${oc.env:BIONEMO_HOME}/models/protein/esm2nv/esm2nv_650M_converted.nemo" # 事前学習済みモデルのパス
Jupyter Lab を起動して、下記の Jupyter Notebook を開いて、/workspace/bionemo/examples/protein/esm2nv/nbs/inference_interactive.ipynb
セルを順次実行してください。
補足情報: 独自データの準備
ローカルマシンのディレクトリ /home/bionemo/ に、train、val、test というフォルダーを準備してください。それぞれのフォルダーの中に、x000.csv, x001.csv,… という名前のファイルを準備してください。
- 事前学習データの準備
事前学習データのフォーマットについては、下記のファイルを参考にしてください。/workspace/bionemo/examples/tests/test_data/uniref202104_esm2_qc_test200_val200/uf50/train/x000.csvrecode_idとsequence列が必須で、他の列は任意です。また、下記の YAML ファイルの data の部分についてもデータの数やパスなどを変更する必要があります。/workspace/bionemo/examples/protein/esm2nv/conf/base_config.yaml - ファインチューニング データの準備
ファインチューニング データのフォーマットについては、下記のファイルを参考にしてください。/workspace/bionemo/examples/tests/test_data/protein/downstream/train/x000.csvid、sequenceとtarget (3state か 8state)列が必須で、他の列は任意です。また、下記の YAML ファイルの data の部分についてもデータの数、パス、ターゲット情報などを変更する必要があります。/workspace/bionemo/examples/protein/esm2nv/conf/downstream_sec_str_LORA.yaml
まとめ
本記事では、BioNeMo Framework を利用して大規模なタンパク質モデルの事前学習、ファインチューニングおよび推論する方法について解説しました。BioNeMo Framework が、タンパク質に関する研究や開発の効率化の役に立てれば嬉しく思います。第 2 弾は大規模な低分子モデルについて解説する予定です。
