Fix for multi gpu setup training with a single GPU.

unsloth/tokenizer_utils.py

@@ -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"