ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix for multi gpu setup training with a single GPU.

Browse source
This commit is contained in:
Alex Noble 2026-03-12 06:59:40 +00:00 • committed by Daniel Han
commit 08abf2f8a8
1 changed files with 8 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

8
unsloth/tokenizer_utils.py
View file

@ -1040,10 +1040,18 @@ def patch_sft_trainer_tokenizer():
"\n"
"import subprocess, re, gc, numpy as np\n"
"a = np.array([0,])\n"
"index_for_cuda = -1\n"
"if 'CUDA_VISIBLE_DEVICES' in os.environ:\n"
" index_for_cuda = os.environ['CUDA_VISIBLE_DEVICES']\n"
" if ',' in index_for_cuda:\n"
" raise RuntimeError('Unsloth currently does not support multi GPU setups - but we are working on it!')\n"
" index_for_cuda = int(index_for_cuda)\n"
"try:\n"
" a = subprocess.check_output('nvidia-smi --query-gpu=memory.used --format=csv', shell = True)\n"
" a = re.findall(rb'([\\d]{1,})[\\s]{1,}M', a)\n"
" a = np.array([int(x.decode('utf-8'))/1024 for x in a])\n"
" if index_for_cuda != -1:\n"
" a = np.array([a[index_for_cuda],])\n"
"except:\n"
" if not torch.cuda.is_available():\n"
" raise RuntimeError('Unsloth: We do not support AMD / Intel machines yet - it is a work in progress!')\n"