if gpuid is specified, force specific order

This commit is contained in:
Concedo 2024-07-30 22:58:25 +08:00
parent 7e95b80211
commit 2f04f848e1

View file

@ -425,15 +425,19 @@ def set_backend_props(inputs):
if not args.tensor_split: if not args.tensor_split:
if (args.usecublas and "0" in args.usecublas): if (args.usecublas and "0" in args.usecublas):
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0" os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["HIP_VISIBLE_DEVICES"] = "0" os.environ["HIP_VISIBLE_DEVICES"] = "0"
elif (args.usecublas and "1" in args.usecublas): elif (args.usecublas and "1" in args.usecublas):
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1" os.environ["CUDA_VISIBLE_DEVICES"] = "1"
os.environ["HIP_VISIBLE_DEVICES"] = "1" os.environ["HIP_VISIBLE_DEVICES"] = "1"
elif (args.usecublas and "2" in args.usecublas): elif (args.usecublas and "2" in args.usecublas):
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2" os.environ["CUDA_VISIBLE_DEVICES"] = "2"
os.environ["HIP_VISIBLE_DEVICES"] = "2" os.environ["HIP_VISIBLE_DEVICES"] = "2"
elif (args.usecublas and "3" in args.usecublas): elif (args.usecublas and "3" in args.usecublas):
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "3" os.environ["CUDA_VISIBLE_DEVICES"] = "3"
os.environ["HIP_VISIBLE_DEVICES"] = "3" os.environ["HIP_VISIBLE_DEVICES"] = "3"
else: else: