Configuring Multi-Instance GPU (MIG) in AWS PCS

NVIDIA Multi-Instance GPU (MIG) partitions a supported GPU into isolated GPU instances, each with its own memory and compute resources. Slurm schedules each MIG profile as its own Generic Resource (GRES) type, so multiple jobs can share one physical GPU with hardware isolation.

MIG devices can't be described statically in gres.conf: their device files and UUIDs only exist after the MIG instances are created on the booted node, so there is no File or Type value to configure ahead of time. On AWS PCS you configure MIG with a gres.conf record that only enables AutoDetect, together with a Gres setting that declares the profiles. slurmd then discovers the MIG devices when the node boots. AWS PCS supports AutoDetect, and therefore MIG, on Slurm version 25.11 and later.

Configuring MIG means partitioning the GPUs on each node and declaring the resulting profiles to Slurm on the compute node group. Neither works alone, so apply both in one change. If Slurm doesn't know about a node's partitions, the node has no schedulable GRES. If no node creates the declared profiles, the nodes are drained when they register. With the recommended node lifecycle action, one compute node group update carries both the partitioning script and the GRES settings.

Prerequisites

What a MIG node must have on every boot

MIG mode and MIG instances live on the GPU, not on disk, so a custom AMI can't carry them, and they're lost whenever the instance is stopped, suspended, or terminated. On every boot, all of the following must be true before slurmd starts:

If MIG mode is enabled but no MIG instances exist, slurmd finds the parent GPUs but no MIG device. It logs MIG mode is enabled, but no MIG devices were found, then Discarding the following config-only GPU due to lack of File specification for each declared profile. Because the node reports fewer GPUs than Gres declares, the controller rejects the registration and drains the node with the reason gres/gpu count reported lower than configured. Jobs that request a GPU on the node fail with Required node not available. Creating the MIG instances afterward doesn't undrain the node; you have to resume or replace it. The node lifecycle action described in Apply the configuration with a node lifecycle action avoids this by failing the node before it registers.

Configure MIG on the instances

Configuring the instances takes the following two operations. Run them from a node lifecycle action, which keeps the logic on the compute node group instead of in an AMI you have to rebuild. See Apply the configuration with a node lifecycle action. If another process already holds the GPUs at that point, run them earlier from a cloud-init boothook or a systemd unit in a custom AMI. See If MIG mode stays pending. For more information about the commands, see NVIDIA Multi-Instance GPU User Guide.

Apply the configuration with a node lifecycle action

Configure the node lifecycle action as follows. For how to add one to a compute node group, see Configure node lifecycle actions in AWS PCS.

A script that runs on every boot must be idempotent. The following outline takes the per-GPU layout as its first argument and handles one GPU at a time. It skips a GPU that already has that layout, repartitions any other, and exits non-zero if it can't finish, which fails the node. It never reboots the node. For how arguments reach a script, see Passing arguments to scripts.

#!/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

To set the script and the GRES settings in one call, see the example in Declare the MIG profiles to Slurm.

If you adapt the script, keep both checks. Without the mode check, a node whose MIG mode never became Enabled goes on to register without its MIG devices, and the cause shows only in the slurmd log. Without the layout check, a node that creates the wrong profiles registers normally and fails only the jobs that request the declared profiles. For where the script's output is written, see Logging and debugging.

If MIG mode stays pending

Enabling MIG mode fails while any process has the GPUs open. nvidia-smi reports the enable as pending, for example with In use by another client, and it takes effect only after those processes exit or the node reboots. On a AWS PCS compute node the process is usually a GPU monitoring agent, such as NVIDIA DCGM or an agent that ships in the AMI, or slurmd itself if MIG is configured after slurmd starts. You can't reboot from the lifecycle action, so use one of the following fixes:

Declare the MIG profiles to Slurm

Declare both GRES settings on the compute node group:

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"
          }
        ]
      }
    }'

For the constraints between the two settings, see Constraints between gres.conf and the Gres setting. Because AWS PCS can't count devices that are discovered at boot, it accepts any profile names and counts in the Gres setting for a record that only enables AutoDetect. Consistency with the actual MIG instances is only verified when the node registers.

When a node boots and registers, it advertises the discovered profiles with their socket affinity, for example Gres=gpu:3g.20gb:8(S:0),gpu:1g.5gb:28(S:1). A job that requests a profile, such as --gres=gpu:3g.20gb:1, is placed on a single MIG instance and gets its UUID in CUDA_VISIBLE_DEVICES (for example, MIG-e7aa6185-06d5-53e9-a4d9-00f7670f741e), so CUDA restricts the job to that MIG instance.

Limitations

Because a MIG configuration names no devices, the controller only knows a node's real GRES after the node boots and registers. Until a node has registered once, the controller holds only the profile counts from the Gres setting. This causes the following limitation. It applies to any GPU record that only enables AutoDetect, and therefore to every MIG compute node group.

Jobs allocated before a node first registers can start without a GPU

A job allocated to a node that is powered down and has never registered, for example on the first boot of a newly created compute node group, is granted a profile count but no specific device, because the devices aren't known yet. When the node then boots and the job starts, the job has no MIG device assigned: CUDA_VISIBLE_DEVICES is not set, and because AWS PCS doesn't constrain device access through cgroups, the job can use every MIG device on the node. Slurm reports the job as RUNNING though, and the MIG instance the job was counted against is assigned to the next job that requests that profile, so both jobs can end up on the same MIG instance. Jobs allocated to the node after it has registered once, including after later power-down and power-up cycles, get their device as expected.

To detect this condition, verify at the start of the job that CUDA_VISIBLE_DEVICES is set, and exit or requeue the job when it isn't:

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

A requeued job becomes eligible to run again after the requeue_delay scheduler parameter, 5 seconds by default. By then the node has registered, so the job gets a device on its next start.