423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-12-15 19:30:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

support llama pro #2338 , add rslora

Browse Source
This commit is contained in:
hiyouga
2024-02-15 02:27:36 +08:00
parent 8a1b389086
commit 7924ffc55d
24 changed files with 438 additions and 203 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/llmtuner/extras/misc.py
View File

@@ -10,6 +10,7 @@ from transformers.utils import (
WEIGHTS_NAME,
is_torch_bf16_gpu_available,
is_torch_cuda_available,
is_torch_mps_available,
is_torch_npu_available,
is_torch_xpu_available,
)
@@ -133,6 +134,8 @@ def get_current_device() -> torch.device:
device = "xpu:{}".format(os.environ.get("LOCAL_RANK", "0"))
elif is_torch_npu_available():
device = "npu:{}".format(os.environ.get("LOCAL_RANK", "0"))
elif is_torch_mps_available():
device = "mps:{}".format(os.environ.get("LOCAL_RANK", "0"))
elif is_torch_cuda_available():
device = "cuda:{}".format(os.environ.get("LOCAL_RANK", "0"))
else: