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

Update __init__.py

Browse source
This commit is contained in:
Daniel Han-Chen 2023-11-30 23:17:41 +11:00
commit 03766c243d
1 changed files with 9 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

9
unsloth/__init__.py
View file

@ -43,9 +43,18 @@ elif n_gpus > 1:
"but we require 'CUDA_VISIBLE_DEVICES=0'\n"\
"We shall set it ourselves."
)
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
elif "CUDA_DEVICE_ORDER" not in os.environ:
warnings.warn(
f"Unsloth: 'CUDA_DEVICE_ORDER' is not set "\
"but we require 'CUDA_DEVICE_ORDER=PCI_BUS_ID'\n"\
"We shall set it ourselves."
)
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
else:
warnings.warn("Unsloth: 'CUDA_VISIBLE_DEVICES' is not set. We shall set it ourselves.")
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
pass