PCS でのマルチインスタンス GPU (MIG) AWS の設定

NVIDIA マルチインスタンス GPU (MIG) は、サポートされている GPU を分離された GPU インスタンスに分割し、それぞれに独自のメモリとコンピューティングリソースがあります。Slurm は、各 MIG プロファイルを独自の汎用リソース (GRES) タイプとしてスケジュールするため、複数のジョブで 1 つの物理 GPU をハードウェア分離と共有できます。

MIG デバイスを で静的に記述することはできませんgres.conf。デバイスファイルと UUIDsは、起動されたノードで MIG インスタンスが作成された後にのみ存在するため、事前に設定する Fileまたは Type値はありません。 AWS PCS では、 のみを有効にするgres.confレコードとAutoDetect、プロファイルを宣言するGres設定を使用して MIG を設定します。 slurmd は、ノードの起動時に MIG デバイスを検出します。 AWS PCS はAutoDetect、Slurm バージョン 25.11 以降で 、つまり MIG をサポートします。

MIG を設定すると、各ノードの GPUs がパーティション分割され、結果のプロファイルがコンピューティングノードグループの Slurm に宣言されます。どちらも単独では機能しないため、両方を 1 回の変更で適用します。Slurm がノードのパーティションを知らない場合、ノードにはスケジューリング可能な GRES がありません。宣言されたプロファイルを作成するノードがない場合、ノードは登録時にドレインされます。推奨されるノードライフサイクルアクションでは、1 つのコンピューティングノードグループの更新でパーティショニングスクリプトと GRES 設定の両方が保持されます。

前提条件

起動ごとに MIG ノードが持つ必要があるもの

MIG モードと MIG インスタンスはディスク上ではなく GPU 上に存在するため、カスタム AMI はそれらを伝送できず、インスタンスが停止、中断、または終了されるたびに失われます。起動するたびに、 slurmdを開始する前に、次のすべてが true である必要があります。

MIG モードが有効になっていても MIG インスタンスが存在しない場合、 は親 GPU slurmdを見つけますが、MIG デバイスを見つけません。 GPUs をログに記録しMIG mode is enabled, but no MIG devices were found、宣言されたプロファイルDiscarding the following config-only GPU due to lack of File specificationごとに を記録します。ノードは がGres宣言する数よりも少ない GPUs を報告するため、コントローラーは登録を拒否し、 という理由でノードをドレインしますgres/gpu count reported lower than configured。ノードで GPU をリクエストするジョブは、 で失敗しますRequired node not available。その後に MIG インスタンスを作成しても、ノードのドレインは解除されません。再開または置き換える必要があります。「」で説明されているノードライフサイクルアクションは、登録前にノードが失敗することでこれノードライフサイクルアクションを使用して設定を適用するを回避します。

インスタンスで MIG を設定する

インスタンスを設定すると、次の 2 つのオペレーションが実行されます。ノードライフサイクルアクションから実行します。これにより、再構築する必要がある AMI ではなく、コンピューティングノードグループにロジックが保持されます。「ノードライフサイクルアクションを使用して設定を適用する」を参照してください。その時点で別のプロセスがすでに GPUs を保持している場合は、cloud-initブートフックまたはカスタム AMI のsystemdユニットからより前に実行します。「MIG モードが保留中のままの場合」を参照してください。コマンドの詳細については、「NVIDIA マルチインスタンス GPU ユーザーガイド」を参照してください。

ノードライフサイクルアクションを使用して設定を適用する

ノードライフサイクルアクションを次のように設定します。コンピューティングノードグループに 1 つを追加する方法については、「」を参照してくださいPCS AWS でノードライフサイクルアクションを設定する。

起動ごとに実行されるスクリプトはべき等である必要があります。次の概要は、最初の引数として GPU ごとのレイアウトを取り、一度に 1 つの GPU を処理します。そのレイアウトが既にある GPU をスキップし、他のすべてのパーティションを再分割して、完了できない場合はゼロ以外を終了し、ノードが失敗します。ノードを再起動することはありません。引数がスクリプトに到達する方法については、「」を参照してくださいスクリプトに引数を渡す。

#!/usr/bin/env bash
# Partition every GPU as $1 specifies, then verify the result. Expected
# nvidia-smi failures are handled inline, so this doesn't use 'set -e'.
set -uo pipefail

# $1 - per-GPU MIG profile names, comma-separated, largest first, as in Gres.
PROFILES="${1:-}"
if ! echo "$PROFILES" | grep -Eq '^[0-9]+g\.[0-9]+gb(,[0-9]+g\.[0-9]+gb)*$'; then
    echo "ERROR: pass the per-GPU MIG profiles as the first argument," >&2
    echo "       for example 3g.20gb,2g.10gb,1g.5gb,1g.5gb" >&2
    exit 1
fi
EXPECTED=$(echo "$PROFILES" | tr ',' '\n' | sort)

# Fail, rather than exit 0, when there is no GPU to configure.
GPUS=$(nvidia-smi --query-gpu=index --format=csv,noheader) || {
    echo "ERROR: nvidia-smi failed. This node has no usable NVIDIA driver." >&2
    exit 1
}
if [ -z "$GPUS" ]; then
    echo "ERROR: no NVIDIA GPUs detected." >&2
    exit 1
fi

# The MIG profiles on GPU $1, in the same form as EXPECTED. nvidia-smi -L lists
# what NVML reports, which is what slurmd autodetects.
mig_devices() {
    nvidia-smi -L | awk -v gpu="$1" '
        /^GPU [0-9]+:/ { current = ($2 + 0); next }
        $1 == "MIG" && current == gpu { print $2 }
    ' | sort
}

for i in $GPUS; do
    # 1. Skip a GPU that already has the requested layout.
    if [ "$(mig_devices "$i")" = "$EXPECTED" ]; then
        echo "GPU $i: already configured as $PROFILES"
        continue
    fi

    # 2. Enable MIG mode. Ignore the exit status; step 3 checks the result.
    enable_output=$(nvidia-smi -i "$i" -mig 1 2>&1) || true

    # 3. Fail if the mode isn't applied. Don't reboot: that would interrupt
    #    the AWS PCS bootstrap.
    current=$(nvidia-smi -i "$i" --query-gpu=mig.mode.current --format=csv,noheader)
    pending=$(nvidia-smi -i "$i" --query-gpu=mig.mode.pending --format=csv,noheader)
    if [ "$current" != "Enabled" ]; then
        echo "ERROR: GPU $i MIG mode is '$current', pending '$pending'." >&2
        echo "$enable_output" >&2
        if [ "$pending" = "Enabled" ]; then
            echo "The enable is pending, so a process is holding the GPU. Stop it," >&2
            echo "or configure MIG mode in a boothook instead." >&2
        elif [ "$current" = "[N/A]" ] || [ "$current" = "N/A" ]; then
            echo "This GPU does not support MIG." >&2
        fi
        exit 1
    fi

    # 4. Recreate the layout from scratch, because instances can't be added to
    #    a partial one. Nothing uses them yet, because slurmd hasn't started.
    nvidia-smi mig -i "$i" -dci >/dev/null 2>&1 || true
    nvidia-smi mig -i "$i" -dgi >/dev/null 2>&1 || true
    if ! nvidia-smi mig -i "$i" -cgi "$PROFILES" -C; then
        echo "ERROR: MIG instance creation failed on GPU $i" >&2
        exit 1
    fi

    # 5. Check which profiles exist, not only how many, so a layout that
    #    doesn't match Gres fails here instead of at job time.
    actual=$(mig_devices "$i")
    if [ "$actual" != "$EXPECTED" ]; then
        echo "ERROR: GPU $i has the wrong MIG layout." >&2
        echo "  expected: $(echo "$EXPECTED" | tr '\n' ' ')" >&2
        echo "  found:    $(echo "$actual" | tr '\n' ' ')" >&2
        exit 1
    fi
    echo "GPU $i: configured as $PROFILES"
done

スクリプトと GRES 設定を 1 回の呼び出しで設定するには、 の例を参照してくださいMIG プロファイルを Slurm に宣言する。

スクリプトを適応させる場合は、両方のチェックを保持します。モードチェックがないと、MIG モードEnabledが MIG デバイスなしで登録されることはなく、原因はslurmdログにのみ表示されます。レイアウトチェックを行わないと、間違ったプロファイルを作成するノードは正常に登録され、宣言されたプロファイルをリクエストするジョブのみに失敗します。スクリプトの出力が書き込まれる場所については、「」を参照してくださいログ記録とデバッグ。

MIG モードが保留中のままの場合

GPU が開いているプロセスがある間、MIG モードの有効化は失敗GPUs。 は、 などで有効化を保留中としてnvidia-smi報告しIn use by another client、それらのプロセスが終了するかノードが再起動された後にのみ有効になります。 AWS PCS コンピューティングノードでは、プロセスは通常、NVIDIA DCGM や AMI に付属するエージェントなどの GPU モニタリングエージェントです。起動後に MIG が設定されている場合はそれslurmd自体slurmdです。ライフサイクルアクションから再起動できないため、次のいずれかの修正を使用します。

MIG プロファイルを Slurm に宣言する

コンピューティングノードグループで両方の GRES 設定を宣言します。

aws pcs update-compute-node-group \
    --cluster-identifier my-cluster \
    --compute-node-group-identifier my-cng-1 \
    --slurm-configuration '{
      "gresCustomSettings": [
        { "AutoDetect": "nvml" }
      ],
      "slurmCustomSettings": [
        { "parameterName": "Gres", "parameterValue": "gpu:3g.20gb:8,gpu:2g.10gb:8,gpu:1g.5gb:16" }
      ]
    }' \
    --node-lifecycle-actions '{
      "stages": {
        "nodeBootstrapped": [
          {
            "name": "configure-mig",
            "scriptSource": {
              "scriptLocation": "s3://my-bucket/configure-mig.sh"
            },
            "arguments": ["3g.20gb,2g.10gb,1g.5gb,1g.5gb"],
            "executionPolicy": "EVERY_BOOT",
            "onError": "TERMINATE"
          }
        ]
      }
    }'

2 つの設定間の制約については、「」を参照してくださいgres.conf と Gres 設定の間の制約。PCS AWS は起動時に検出されたデバイスをカウントできないため、 のみを有効にするレコードの Gres設定でプロファイル名とカウントを受け入れますAutoDetect。実際の MIG インスタンスとの整合性は、ノードの登録時にのみ検証されます。

ノードが起動して登録されると、検出されたプロファイルをソケットアフィニティ、たとえば でアドバタイズしますGres=gpu:3g.20gb:8(S:0),gpu:1g.5gb:28(S:1)。などのプロファイルをリクエストするジョブは --gres=gpu:3g.20gb:11 つの MIG インスタンスに配置され、 で UUID を取得するため CUDA_VISIBLE_DEVICES (例: MIG-e7aa6185-06d5-53e9-a4d9-00f7670f741e)、CUDA はジョブをその MIG インスタンスに制限します。

制限事項

MIG 設定ではデバイスの名前がないため、コントローラーはノードの起動と登録後にのみノードの実際の GRES を認識します。ノードが一度登録されるまで、コントローラーは Gres設定のプロファイル数のみを保持します。これにより、次の制限が発生します。これは、 のみを有効にする GPU レコードAutoDetect、つまりすべての MIG コンピューティングノードグループに適用されます。

ノードが最初に登録される前に割り当てられたジョブは GPU なしで開始できます

新しく作成されたコンピューティングノードグループの最初の起動時など、電源が切られて登録されていないノードに割り当てられたジョブには、プロファイルカウントが付与されますが、デバイスはまだわからないため、特定のデバイスはありません。ノードが起動してジョブが開始されると、ジョブには MIG デバイスが割り当てられません。 CUDA_VISIBLE_DEVICES は設定されず、PCS AWS は cgroup を介したデバイスアクセスを制限しないため、ジョブはノード上のすべての MIG デバイスを使用できます。Slurm はジョブをRUNNINGそのように報告し、ジョブがカウントされた MIG インスタンスは、そのプロファイルをリクエストする次のジョブに割り当てられるため、両方のジョブが同じ MIG インスタンスで終了する可能性があります。1 回登録した後にノードに割り当てられたジョブは、後の電源オフサイクルと電源アップサイクルの後を含め、想定どおりにデバイスを取得します。

この条件を検出するには、設定CUDA_VISIBLE_DEVICESされているジョブの開始時に を確認し、設定されていない場合はジョブを終了または再キューに入れます。

if [ -z "${CUDA_VISIBLE_DEVICES-}" ]; then
    scontrol requeue "$SLURM_JOB_ID"
    exit 1
fi

再キューに入れられたジョブは、デフォルトで 5 秒のスrequeue_delayケジューラパラメータの後に再度実行できるようになります。その後、ノードは登録されているため、ジョブは次の起動時にデバイスを取得します。