Skip to content

feat(train): add AUTO_FREE_GPUS mode to train_multi_gpu.sh - #14

Merged
tkgstrator merged 1 commit into
developfrom
feat/auto-free-gpus-v2
Jul 13, 2026
Merged

feat(train): add AUTO_FREE_GPUS mode to train_multi_gpu.sh#14
tkgstrator merged 1 commit into
developfrom
feat/auto-free-gpus-v2

Conversation

@tkgstrator

Copy link
Copy Markdown
Contributor

Re-opens #12 with a commitlint-compliant body (100-char wrap).

What

scripts/train_multi_gpu.sh gains two new env vars:

  • AUTO_FREE_GPUS=1 — before launching torchrun, query nvidia-smi for per-GPU memory usage and set CUDA_VISIBLE_DEVICES to the GPUs whose used memory is under FREE_GPU_MEM_MB.
  • FREE_GPU_MEM_MB (default 500) — the "free" threshold in MiB.

NPROC_PER_NODE is derived from the count of surviving GPUs so torchrun uses only those.

Why

The 8x A100 box is shared with other running jobs. Kicking off a naive NPROC_PER_NODE=8 DDP run would OOM those GPUs or trample the neighbors. This mode picks up only the idle GPUs so multiple training jobs coexist safely on the same host.

Test plan

  • detect_free_gpus 500 on the current box returns 1,3,4,5,6,7 (skips busy GPU 0 and 2)
  • Launch DDP with AUTO_FREE_GPUS=1 and confirm CUDA_VISIBLE_DEVICES is set and torchrun uses only the free 6 GPUs
  • If zero GPUs qualify, script exits with a helpful error (nvidia-smi output printed)

🤖 Generated with Claude Code

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 with each other on shared
hardware.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@tkgstrator
tkgstrator merged commit 198dee3 into develop Jul 13, 2026
3 checks passed
@tkgstrator
tkgstrator deleted the feat/auto-free-gpus-v2 branch July 13, 2026 04:34
@tkgstrator tkgstrator mentioned this pull request Jul 13, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant