GCPのNVIDIA A100 GPUでマルチインスタンス GPU（MIG）を使用する
NVIDIA H100やA100ではマルチインスタンス GPU(以後、MIG)機能を使うことで1枚のGPUカードを複数のインスタンスに分割することができるようになります。分割することでA100のポテンシャルを使い切ることが容易にできます（使用用途にもよりますが）。
Google CloudがGCEで提供するA100を推論に使うことになったのですが、そのままではA100のスペックを使い切ることが難しく料金が高いということもあり使い切りたかったのでMIGを使うことになりましたが、セットアップでハマりどころがありました。
GCEならではの事情
MIGを使うには下記3つのステップが必要です。
- MIGの有効化
- GPUのリセット、もしくはGPU搭載のVMの再起動
- A100をインスタンスに分割する
GCEにあるA100ではセキュリティの観点からGPUのリセットが許可されておらず一度VMを再起動する必要があります。またインスタンスの分割設定はVMの停止時にリセットされ永続化されないためVMを起動したり、停止する使い方をする場合、その都度設定をし直す必要があります。
これらを人間が毎回やるわけにはいかなかったためこのMIGの有効化を自動化する仕組みをつくりました。
自動化
- VMが起動
- MIGを有効化する
- VM再起動
- 起動時にMIGが有効化されていれば分割設定を行う
という4ステップが必要です。
今回はこれらの処理を自動化するスクリプトを実行するプロセスをsystemdのserviceとして登録しています。
serviceを使う場合、上の4ステップはこのように実行されます。
- serviceの起動コマンドとしてスクリプトを呼び出す
- スクリプトはMIGが有効になっていなければMIGを有効化し、VMを再起動する
- serviceの起動コマンドとして同じスクリプトが呼び出される
- スクリプトはMIGが有効になっていればインスタンスの分割設定を行う
具体的なコードはこちらです
service
/lib/systemd/system/enable-nvidia-gpu-mig.service などに書き出しています。
[Unit]
Description=Enable NVIDIA GPU MIG
After=nvidia-persistenced.service
[Service]
Type=oneshot
ExecStart=/usr/local/bin/enable-nvidia-gpu-mig
ExecStop=/usr/bin/echo enable-nvidia-gpu-mig prepared
RemainAfterExit=true
[Install]
WantedBy=multi-user.target
serviceが実行するスクリプト
/usr/local/bin/enable-nvidia-gpu-mig に書き出しています。
#!/bin/bash
echo "Start enable-nvidia-gpu-mig service..."
if nvidia-smi -i 0 --query-gpu=pci.bus_id,mig.mode.current --format=csv | grep -q "Disabled"; then
echo "Ordering enabling NVIDIA GPU MIG"
sudo nvidia-smi -mig 1
echo "Rebooting to enable NVIDIA GPU MIG"
sudo reboot
else
echo "Creating Virtual GPU Instances"
sudo nvidia-smi mig -cgi 9,3g.20gb -C
fi
serviceの登録コマンド
systemctlを使うことでこのserviceを登録することができます。
# 実行権限付与
sudo chmod 755 /usr/local/bin/enable-nvidia-gpu-mig
# 登録
sudo systemctl enable enable-nvidia-gpu-mig
参考情報
公式のガイド
MIGの仕組みについてはこちらの記事が日本語でわかりやすくまとめてくださってます。
Discussion