🤖

GCEのNVIDIA A100 GPUでマルチインスタンス GPU（MIG）を使用する

2024/04/06に公開

Google Cloud

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の仕組みについてはこちらの記事が日本語でわかりやすくまとめてくださってます。