Add multi-GPU support with --gpus flag

Adds --gpus flag to obliterate, run, and tourney commands for controlling
which GPUs to use (sets CUDA_VISIBLE_DEVICES). Works both locally and with
--remote. Models are automatically split across selected GPUs via
accelerate's device_map="auto". Also adds gpus field to remote YAML config.
This commit is contained in:
Stella Biderman
2026-03-12 12:35:29 -04:00
parent 34032b7821
commit cbdb772eb9
4 changed files with 74 additions and 5 deletions
+6
View File
@@ -9,6 +9,11 @@
#
# You can also use --remote on any command instead of a YAML section:
# obliteratus obliterate meta-llama/Llama-3.1-8B-Instruct --remote root@gpu-node --ssh-key ~/.ssh/id_rsa
#
# Multi-GPU: Models are automatically split across all available GPUs via
# accelerate's device_map="auto". Use --gpus or the gpus: field to select
# specific GPUs:
# obliteratus obliterate model --remote root@gpu-node --gpus 0,1,2,3
model:
name: meta-llama/Llama-3.1-8B-Instruct
@@ -39,3 +44,4 @@ remote:
remote_dir: /tmp/obliteratus_run
python: python3
sync_results: true
# gpus: "0,1,2,3" # uncomment to select specific GPUs (default: all)