From a5064231eeb294f291d9cb5e74cb5afe0e7f2064 Mon Sep 17 00:00:00 2001
From: tkgstrator <nasawake.am@gmail.com>
Date: Mon, 13 Jul 2026 04:07:38 +0000
Subject: [PATCH] feat(train): add AUTO_FREE_GPUS for GPU auto-selection by
 free memory

Automatically restrict CUDA_VISIBLE_DEVICES to GPUs with used memory below FREE_GPU_MEM_MB (default 500 MiB) when AUTO_FREE_GPUS=1. Prevents multiple training jobs from interfering on shared hardware.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 scripts/train_multi_gpu.sh | 35 +++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/scripts/train_multi_gpu.sh b/scripts/train_multi_gpu.sh
index 04360b7..c21c654 100755
--- a/scripts/train_multi_gpu.sh
+++ b/scripts/train_multi_gpu.sh
@@ -7,11 +7,18 @@
 #   PER_GPU_WORKERS   DataLoader workers per rank (defaults to 4; total across
 #                     ranks is NPROC_PER_NODE * PER_GPU_WORKERS)
 #   BACKBONE          Model backbone (defaults to mobilenet_v3_small)
+#   AUTO_FREE_GPUS    When 1, restrict training to GPUs whose used memory is
+#                     below FREE_GPU_MEM_MB (default 500). Sets
+#                     CUDA_VISIBLE_DEVICES and NPROC_PER_NODE to the survivors
+#                     so shared boxes don't step on running jobs.
+#   FREE_GPU_MEM_MB   "Free" threshold in MiB (default 500). Only used when
+#                     AUTO_FREE_GPUS=1.
 #
 # Examples:
 #   ./scripts/train_multi_gpu.sh                           # all GPUs, defaults
 #   NPROC_PER_NODE=8 BACKBONE=convnext_tiny ./scripts/train_multi_gpu.sh
 #   NPROC_PER_NODE=4 EPOCHS=100 ./scripts/train_multi_gpu.sh
+#   AUTO_FREE_GPUS=1 BACKBONE=mobilenet_v3_small ./scripts/train_multi_gpu.sh
 
 set -euo pipefail
 
@@ -20,6 +27,34 @@ detect_gpu_count() {
     uv run python -c "import torch; print(torch.cuda.device_count())" 2>/dev/null || echo 1
 }
 
+# Select GPU indices whose used memory is under FREE_GPU_MEM_MB. Prints a
+# comma-separated list (or an empty string if none qualify).
+detect_free_gpus() {
+    local threshold="${1:-500}"
+    nvidia-smi --query-gpu=index,memory.used --format=csv,noheader,nounits 2>/dev/null \
+        | awk -F, -v t="${threshold}" '{
+              gsub(/ /, "", $1); gsub(/ /, "", $2);
+              if ($2+0 < t+0) picks[n++] = $1
+          }
+          END {
+              for (i=0; i<n; i++) printf "%s%s", (i?",":""), picks[i]
+          }'
+}
+
+if [[ "${AUTO_FREE_GPUS:-0}" == "1" ]]; then
+    FREE_GPU_MEM_MB=${FREE_GPU_MEM_MB:-500}
+    free_list=$(detect_free_gpus "${FREE_GPU_MEM_MB}")
+    if [[ -z "${free_list}" ]]; then
+        echo "[train_multi_gpu] AUTO_FREE_GPUS=1 but no GPU has memory < ${FREE_GPU_MEM_MB}MiB" >&2
+        nvidia-smi --query-gpu=index,memory.used --format=csv >&2
+        exit 1
+    fi
+    export CUDA_VISIBLE_DEVICES="${free_list}"
+    # torchrun sees the remapped devices, so nproc is just the count.
+    NPROC_PER_NODE=$(awk -F, '{print NF}' <<<"${free_list}")
+    echo "[train_multi_gpu] AUTO_FREE_GPUS: CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} (threshold=${FREE_GPU_MEM_MB}MiB)"
+fi
+
 NPROC_PER_NODE=${NPROC_PER_NODE:-$(detect_gpu_count)}
 MASTER_PORT=${MASTER_PORT:-29500}
 PER_GPU_WORKERS=${PER_GPU_WORKERS:-4}